mirror of
https://github.com/harry7557558/spirula-studio.git
synced 2026-10-04 03:48:23 +08:00
369 lines
13 KiB
Python
369 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
|
|
# extract frames from video or rosbag file
|
|
# pick good frames that minimize motion blur
|
|
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import os
|
|
|
|
import subprocess
|
|
|
|
quality = 95
|
|
rotate = 0
|
|
scale = 1.0
|
|
|
|
ENHANCE = False # unsupported for now
|
|
|
|
|
|
def write_image(image, filename, filename_enh=None, _enhance_model=[]):
|
|
if rotate != 0:
|
|
rotate_param = {
|
|
90: cv2.ROTATE_90_CLOCKWISE,
|
|
180: cv2.ROTATE_180,
|
|
-90: cv2.ROTATE_90_COUNTERCLOCKWISE,
|
|
270: cv2.ROTATE_90_COUNTERCLOCKWISE
|
|
}[rotate]
|
|
image = cv2.rotate(image, rotate_param)
|
|
if image is None:
|
|
return
|
|
|
|
# resize image
|
|
tile = 8 if ENHANCE else 1
|
|
if scale < 1.0:
|
|
h, w = image.shape[:2]
|
|
image = cv2.resize(image, (int(scale*w/tile+0.5)*tile, int(scale*h/tile+0.5)*tile),
|
|
interpolation=cv2.INTER_AREA)
|
|
elif scale == int(scale) and scale != 1.0:
|
|
iscale = int(scale)
|
|
h, w, c = image.shape
|
|
h, w = (h//(iscale*tile))*tile, (w//(iscale*tile))*tile
|
|
image = image[:h*iscale, :w*iscale]
|
|
image = np.mean(np.mean(
|
|
image.reshape(h, iscale, w, iscale, c),
|
|
axis=1), axis=2).astype(np.uint8)
|
|
|
|
if image is None:
|
|
return
|
|
|
|
def write_im(image, filename):
|
|
if 0 <= quality <= 100:
|
|
filename += '.jpg'
|
|
encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]
|
|
cv2.imwrite(filename, image, encode_param)
|
|
else:
|
|
filename += '.png'
|
|
cv2.imwrite(filename, image)
|
|
|
|
write_im(image, filename)
|
|
|
|
if filename_enh is not None:
|
|
raise NotImplementedError()
|
|
if len(_enhance_model) == 0:
|
|
import colmap.image_enhance as enhance
|
|
image_enhance_model = enhance.load_image_enhance_model_02()
|
|
image_enhance_model = image_enhance_model.float()
|
|
_enhance_model.append(enhance.torch)
|
|
_enhance_model.append(image_enhance_model)
|
|
torch = enhance.torch
|
|
else:
|
|
torch, image_enhance_model = _enhance_model
|
|
|
|
image = torch.from_numpy(image).cuda().permute(2,0,1).unsqueeze(0).float()
|
|
image = image_enhance_model(image)
|
|
image = torch.clip(image, 0.0, 255.0)[0].byte().permute(1,2,0).cpu().numpy()
|
|
write_im(image, filename_enh)
|
|
|
|
return filename
|
|
|
|
|
|
class FrameSelector:
|
|
def __init__(self, image_dir, enhance_dir, max_frames, skip, keep):
|
|
os.makedirs(image_dir, exist_ok=True)
|
|
if enhance_dir is not None:
|
|
os.makedirs(enhance_dir, exist_ok=True)
|
|
|
|
self.image_dir = image_dir
|
|
self.enhance_dir = enhance_dir
|
|
self.max_frames = max_frames
|
|
self.skip = skip
|
|
self.keep = keep
|
|
|
|
self.frame_count = 0
|
|
self.written_frame_count = 0
|
|
self.iqms = []
|
|
self.frames_buf = []
|
|
|
|
def add_frame(self, frame, comment=""):
|
|
if self.written_frame_count >= self.max_frames:
|
|
print("Maximum number of frames reached.")
|
|
return False
|
|
|
|
# check motion blur
|
|
s = 256
|
|
gray = cv2.cvtColor(cv2.resize(frame, (2*s, 2*s)), cv2.COLOR_BGR2GRAY)
|
|
mu, sigma = np.mean(gray), np.std(gray)
|
|
gray = (gray-mu)#/sigma
|
|
|
|
iqm = cv2.Laplacian(gray, cv2.CV_64F).var()
|
|
|
|
self.iqms.append((mu, sigma, iqm))
|
|
|
|
self.frames_buf.append((self.frame_count, iqm, frame, comment))
|
|
if len(self.frames_buf) > max(self.keep, 1):
|
|
del self.frames_buf[0]
|
|
|
|
if self.keep != 0:
|
|
self.frame_count += 1
|
|
if self.frame_count % self.skip == 0:
|
|
fid, iqm, frame, comment = sorted(self.frames_buf, key=lambda x: x[1], reverse=True)[0]
|
|
filename = os.path.join(self.image_dir, f"{fid:05d}")
|
|
filename_enh = None if self.enhance_dir is None else os.path.join(self.enhance_dir, f"{fid:05d}")
|
|
if False: # gamma adjustment
|
|
gamma = 0.6
|
|
lut = (255*np.linspace(0, 1, 256)**gamma).astype(np.uint8)
|
|
frame = lut[frame]
|
|
filename = write_image(frame, filename, filename_enh)
|
|
print(filename, comment)
|
|
self.written_frame_count += 1
|
|
if self.keep == 0:
|
|
self.frame_count += 1
|
|
|
|
return True
|
|
|
|
def conclude(self):
|
|
print(f"Extracted {self.written_frame_count} frames.")
|
|
|
|
|
|
|
|
def extract_rosbag_frames(bag_path, image_dir, enhance_dir, max_frames, skip, keep, required_topic=None):
|
|
|
|
def _imgmsg_to_cv2(img_msg):
|
|
|
|
dtype = np.dtype("uint8")
|
|
dtype = dtype.newbyteorder('>' if img_msg.is_bigendian else '<')
|
|
image_opencv = np.ndarray(shape=(img_msg.height, img_msg.width, 3),
|
|
dtype=dtype, buffer=img_msg.data)
|
|
|
|
if img_msg.encoding == "rgb8":
|
|
image_opencv = cv2.cvtColor(image_opencv, cv2.COLOR_BGR2RGB)
|
|
elif img_msg.encoding != "bgr8":
|
|
raise ValueError(f"Unsupported encoding: {img_msg.encoding}")
|
|
|
|
return image_opencv
|
|
|
|
def _compressed_imgmsg_to_cv2(compressed_img_msg):
|
|
np_arr = np.frombuffer(compressed_img_msg.data, np.uint8)
|
|
return cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
|
|
|
|
def get_image(msg):
|
|
if msg._type == 'sensor_msgs/Image':
|
|
cv_image = _imgmsg_to_cv2(msg)
|
|
else: # CompressedImage
|
|
cv_image = _compressed_imgmsg_to_cv2(msg)
|
|
return cv_image
|
|
|
|
frame_selector = FrameSelector(image_dir, enhance_dir, max_frames, skip, keep)
|
|
|
|
all_image_topics = set()
|
|
frame_count = 0
|
|
with __import__('rosbag').Bag(bag_path, 'r') as bag:
|
|
for topic, msg, t in bag.read_messages():
|
|
if required_topic is not None and topic != required_topic:
|
|
continue
|
|
if msg._type in ['sensor_msgs/Image', 'sensor_msgs/CompressedImage']:
|
|
all_image_topics.add(topic)
|
|
if frame_count % skip != 0:
|
|
frame_count += 1
|
|
continue
|
|
image = get_image(msg)
|
|
|
|
if not frame_selector.add_frame(image, 'topic='+topic):
|
|
break
|
|
|
|
frame_selector.conclude()
|
|
|
|
if required_topic is None:
|
|
print("Detected image topics:", all_image_topics)
|
|
|
|
|
|
def _get_video_codec(video_path):
|
|
cmd = [
|
|
"ffprobe", "-v", "error", "-select_streams", "v:0",
|
|
"-show_entries", "stream=codec_name", "-of", "csv=p=0", video_path
|
|
]
|
|
try:
|
|
result = subprocess.run(cmd, capture_output=True, text=True, timeout=10)
|
|
return result.stdout.strip().lower()
|
|
except (FileNotFoundError, subprocess.TimeoutExpired):
|
|
return None
|
|
|
|
|
|
def _iter_frames_cv2(video_path):
|
|
video = cv2.VideoCapture(video_path)
|
|
if not video.isOpened():
|
|
print(f"Error: could not open {video_path}")
|
|
return
|
|
try:
|
|
while True:
|
|
success, frame = video.read()
|
|
if not success:
|
|
break
|
|
yield frame
|
|
finally:
|
|
video.release()
|
|
|
|
|
|
def _iter_frames_pyav(video_path):
|
|
try:
|
|
import av
|
|
except ImportError:
|
|
print("Error: PyAV is required for AV1 video decoding.")
|
|
print("Install it with: pip install av")
|
|
exit(1)
|
|
container = av.open(video_path)
|
|
try:
|
|
for frame in container.decode(video=0):
|
|
yield frame.to_ndarray(format='bgr24')
|
|
finally:
|
|
container.close()
|
|
|
|
|
|
def iter_video_frames(video_path):
|
|
codec = _get_video_codec(video_path)
|
|
if codec == 'av1':
|
|
print(f"AV1 codec detected, using PyAV decoder.")
|
|
return _iter_frames_pyav(video_path)
|
|
return _iter_frames_cv2(video_path)
|
|
|
|
|
|
def extract_video_frames(video_path, image_dir, enhance_dir, max_frames, skip, keep):
|
|
|
|
def get_video_streams(video_path):
|
|
cmd = [
|
|
"ffprobe", "-v", "error", "-select_streams", "v",
|
|
"-show_entries", "stream=index", "-of", "csv=p=0", video_path
|
|
]
|
|
try:
|
|
result = subprocess.run(cmd, capture_output=True, text=True)
|
|
except FileNotFoundError:
|
|
print("ffprobe not found. Please install ffmpeg.")
|
|
exit(0)
|
|
if result.returncode != 0:
|
|
raise RuntimeError(f"ffprobe failed: {result.stderr}")
|
|
streams = [int(x.strip()) for x in result.stdout.splitlines() if x.strip()]
|
|
return streams
|
|
|
|
streams = get_video_streams(video_path) if video_path.endswith(".insv") else [0]
|
|
|
|
if len(streams) == 0:
|
|
print(f"Error: no video streams found in {video_path}")
|
|
return
|
|
|
|
for idx, stream_idx in enumerate(streams):
|
|
# Handle multi-track video by splitting track via ffmpeg if needed
|
|
if len(streams) > 1:
|
|
print(f"Extracting track {stream_idx} to cam{idx}")
|
|
temp_path = os.path.join(image_dir, f"temp_cam{idx}.mp4")
|
|
subprocess.run([
|
|
"ffmpeg", "-y", "-i", video_path,
|
|
"-map", f"0:v:{idx}", "-c", "copy", temp_path
|
|
], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
track_video_path = temp_path
|
|
out_image_dir = os.path.join(image_dir, f"cam{idx}")
|
|
out_enhance_dir = os.path.join(enhance_dir, f"cam{idx}") if enhance_dir is not None else None
|
|
else:
|
|
track_video_path = video_path
|
|
out_image_dir = image_dir
|
|
out_enhance_dir = enhance_dir
|
|
|
|
frame_selector = FrameSelector(out_image_dir, out_enhance_dir, max_frames, skip, keep)
|
|
|
|
for frame in iter_video_frames(track_video_path):
|
|
if not frame_selector.add_frame(frame):
|
|
break
|
|
|
|
frame_selector.conclude()
|
|
|
|
if len(streams) > 1:
|
|
os.remove(track_video_path)
|
|
|
|
print("Frame extraction completed.")
|
|
|
|
|
|
def extract_frames(filename, dirname, max_frames=100000, skip=1, keep=-1, ros_topic=""):
|
|
if ENHANCE:
|
|
image_dir = os.path.join(dirname, 'images_raw')
|
|
enhance_dir = os.path.join(dirname, 'images')
|
|
else:
|
|
image_dir = os.path.join(dirname, 'images')
|
|
enhance_dir = None
|
|
|
|
if not os.path.exists(image_dir):
|
|
os.makedirs(image_dir)
|
|
if enhance_dir is not None and not os.path.exists(enhance_dir):
|
|
os.makedirs(enhance_dir)
|
|
|
|
if keep == -1:
|
|
keep = int(0.5*skip+0.5)
|
|
|
|
if filename.split('.')[-1].lower() == 'bag':
|
|
if ros_topic == "":
|
|
ros_topic = None
|
|
extract_rosbag_frames(filename, image_dir, enhance_dir, max_frames, skip, keep, ros_topic)
|
|
|
|
else: # video
|
|
extract_video_frames(filename, image_dir, enhance_dir, max_frames, skip, keep)
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
parser = argparse.ArgumentParser(
|
|
description="Extract frames from a rosbag or a video.")
|
|
parser.add_argument("input_file", nargs=1, help="The input rosbag or video file.")
|
|
parser.add_argument("--max_frames", "-n", default="100000", help="Maximum number of frames to output.")
|
|
parser.add_argument("--skip", "-s", default="1", help="Take one frame every this number of frames.")
|
|
parser.add_argument("--keep", "-k", default="-1", help="A fraction of --skip, nonzero for selecting images with minimum motion blur")
|
|
parser.add_argument("--topic", default="", help="rosbag topic. If not set, it will exact all detected frames.")
|
|
parser.add_argument("--quality", "-q", default="95", help="Quality to save JPEG images. Save lossless PNG if this is not between 0 and 100.")
|
|
parser.add_argument("--rotate", "-r", default="0", help="Rotate clockwise this degrees, must be one of {0, 90, 180, 270}.")
|
|
parser.add_argument("--scale", default="1", help="Relative image size scale, less than 1.0 for downscale.")
|
|
parser.add_argument("--mask", default="0", help="Whether to use export scripts to generate masks.")
|
|
args = parser.parse_args()
|
|
|
|
quality = int(args.quality)
|
|
rotate = int(args.rotate)
|
|
scale = float(args.scale)
|
|
filename = args.input_file[0]
|
|
dirname = filename[:filename.rfind('.')]
|
|
extract_frames(
|
|
filename, dirname,
|
|
int(args.max_frames), int(args.skip), int(args.keep), args.topic
|
|
)
|
|
|
|
cur_dir = os.path.dirname(__file__)
|
|
|
|
# generate masks
|
|
if int(args.mask) != 0:
|
|
# run_colmap.bash
|
|
open(os.path.join(dirname, 'run_colmap.bash'), 'w').write(f"""
|
|
# {' '.join(__import__('sys').argv)}
|
|
|
|
{open(os.path.join(cur_dir, "run_colmap.bash")).read()}""".lstrip())
|
|
# run_sam2.bash
|
|
open(os.path.join(dirname, 'run_sam2.bash'), 'w').write(f"""
|
|
# Run this from SAM2 install directory, after downloading checkpoints
|
|
python3 {os.path.join(cur_dir, "SAM2-GUI", "mask_app.py")} --root_dir {os.path.abspath(dirname)}
|
|
""".lstrip())
|
|
# run_lang_sam.bash
|
|
|
|
# no mask
|
|
else:
|
|
# run_colmap.bash
|
|
open(os.path.join(dirname, 'run_colmap.bash'), 'w').write(f"""
|
|
# {' '.join(__import__('sys').argv)}
|
|
|
|
{open(os.path.join(cur_dir, "run_colmap.bash")).read()}""".lstrip())
|