def downscale_frames(self,
subdir,
max_size,
ext,
align=16,
full_subdir="color_full"):
full_dir = pjoin(self.path, full_subdir)
down_dir = pjoin(self.path, subdir)
mkdir_ifnotexists(down_dir)
if self.check_frames(down_dir, ext):
# Frames are already extracted and checked OK.
return
for i in range(self.frame_count):
full_file = "%s/frame_%06d.png" % (full_dir, i)
down_file = ("%s/frame_%06d." + ext) % (down_dir, i)
suppress_messages = (i > 0)
image = image_io.load_image(full_file,
max_size=max_size,
align=align,
suppress_messages=suppress_messages)
image = image[..., ::-1] # Channel swizzle
if ext == "raw":
image_io.save_raw_float32_image(down_file, image)
else:
cv2.imwrite(down_file, image * 255)
self.check_frames(down_dir, ext)
def extract_frames(self):
frame_dir = "%s/color_full" % self.path
mkdir_ifnotexists(frame_dir)
if self.check_frames(frame_dir, "png"):
# Frames are already extracted and checked OK.
return
if not os.path.exists(self.video_file):
sys.exit("ERROR: input video file '%s' not found.",
self.video_file)
cmd = "%s -i %s -start_number 0 -vsync 0 %s/frame_%%06d.png" % (
ffmpeg,
self.video_file,
frame_dir,
)
print(cmd)
os.popen(cmd).read()
count = len(os.listdir(frame_dir))
if count != self.frame_count:
sys.exit("ERROR: %d frames extracted, but %d PTS entries." %
(count, self.frame_count))
self.check_frames(frame_dir, "png")
def extract_frames(self):
frame_dir = "%s/color_full" % self.path
mkdir_ifnotexists(frame_dir) # 创建目录 视频文件目录/color_full
if self.check_frames(
frame_dir,
"png"): # 用于检查视频文件目录/color_full否已经有frames文件，如果有则不需要再提取
# Frames are already extracted and checked OK.
return
if not os.path.exists(self.video_file):
sys.exit("ERROR: input video file '%s' not found.",
self.video_file)
cmd = "%s -i %s -start_number 0 -vsync 0 %s/frame_%%06d.png" % (
ffmpeg,
self.video_file,
frame_dir,
) # ffmpeg -i 视频文件 -start_number 0 -vsync 0 视频文件目录/color_full/frame_%%06d.png
print(cmd)
os.popen(cmd).read()
count = len(os.listdir(frame_dir))
if count != self.frame_count:
sys.exit(
"ERROR: %d frames extracted, but %d PTS entries." # 提取到的帧数与记录的pts数目不一致
% (count, self.frame_count))
self.check_frames(frame_dir, "png")
def compute_flow(self, index_pairs, checkpoint):
"""Run Flownet2 with specific <checkpoint> (FlowNet2 or finetuned on KITTI)
Note that we don't fit homography first for FlowNet2-KITTI model.
"""
model_name = checkpoint.lower()
if model_name == "flownet2-kitti":
model_file = get_model_from_url(
"https://www.dropbox.com/s/mme80czrpbqal7k/flownet2-kitti.pth.tar?dl=1",
model_name + ".pth",
)
else:
model_file = f"checkpoints/{model_name}.pth"
mkdir_ifnotexists("%s/flow" % self.path)
if self.check_flow_files(index_pairs):
return
frame_dir = "%s/color_flow" % self.path
frame1_fns = [
"%s/frame_%06d.png" % (frame_dir, pair[0]) for pair in index_pairs
]
frame2_fns = [
"%s/frame_%06d.png" % (frame_dir, pair[1]) for pair in index_pairs
]
out_fns = [
"%s/flow/flow_%06d_%06d.raw" % (self.path, i, j)
for (i, j) in index_pairs
]
tmp = image_io.load_raw_float32_image(
pjoin(self.path, "color_down", "frame_{:06d}.raw".format(0)))
size = tmp.shape[:2][::-1]
print("Resizing flow to", size)
args = dotdict()
args.pretrained_model_flownet2 = model_file
args.im1 = list(frame1_fns)
args.im2 = list(frame2_fns)
args.out = list(out_fns)
args.size = size
args.fp16 = False
args.homography = 'KITTI' not in checkpoint
args.rgb_max = 255.0
args.visualize = False
optical_flow_flownet2_homography.process(args)
self.check_flow_files(index_pairs)
