def vid_writer(output_vid_name, fps, frame_shape, frame_count=None):
""" manage vid cam """
fourcc = get_fourcc(output_vid_name)
# vid writer if shape isknows after rot
out_shape_cv = np_shape_to_cv(frame_shape[:2])
vid_writer = cv2.VideoWriter(output_vid_name, fourcc, fps, out_shape_cv)
# vid_writer.write(montage_image)
yield vid_writer
if frame_count is not None:
vid_writer.set(cv2.CAP_PROP_FRAME_COUNT, frame_count)
vid_writer.set(cv2.CAP_PROP_FPS, fps)
log.info("output vid: {}".format(output_vid_name))
vid_writer.release()
def main():
args = get_args()
input_imgs_file = os.path.expanduser(args.input_imgs)
n_views = args.num_views
image_size = (args.img_height, args.img_width)
assert os.path.isdir(input_imgs_file), "input not a dir"
input_imgs_file = get_view_pair_vid_files(n_views, input_imgs_file)
assert len(input_imgs_file), "no vids found"
view_pair_idx = args.view_idx
input_imgs_file = [
view_piar_vid[view_pair_idx] for view_piar_vid in input_imgs_file
]
input_imgs_file = sklearn.utils.shuffle(input_imgs_file)
fourcc = get_fourcc(args.output_vid_name)
log.info("output vid: {}".format(args.output_vid_name))
fps = args.fps
vid_writer = None
for frames in tqdm(
get_frames(input_imgs_file, args.mun_col * args.mun_row,
args.num_frames, image_size),
desc="frame",
total=args.num_frames,
):
imgs = [[frames[y] for y in range(x, x + args.mun_row)]
for x in range(0, len(frames), args.mun_row)]
margin = 2
montage_image = montage(
imgs,
margin_color_bgr=[0, 0, 0],
margin_top=margin,
margin_bottom=margin,
margin_left=margin,
margin_right=margin,
margin_separate_vertical=margin,
margin_separate_horizontal=margin,
)
montage_image = convert_to_uint8(montage_image)
if vid_writer is None:
# vid writer if shape isknows after rot
out_shape_cv = np_shape_to_cv(montage_image.shape[:2])
vid_writer = cv2.VideoWriter(args.output_vid_name, fourcc, fps,
out_shape_cv)
vid_writer.write(montage_image)
vid_writer.set(cv2.CAP_PROP_FRAME_COUNT, args.num_frames)
vid_writer.set(cv2.CAP_PROP_FPS, fps)
vid_writer.release()
def get_frame(self, frame_index):
"""Gets a frame at a specified index in a video."""
cap = cv2.VideoCapture(self._vid_path)
if self.approximate_frameCount is None:
self._get_vid_info(cap)
if frame_index < 0 or frame_index >= self.approximate_frameCount:
msg = "frame {} to high for video {} with appr. {} frames".format(
frame_index, self._vid_path, self.approximate_frameCount)
raise IndexError(msg)
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
ok, bgr = cap.read()
if not ok or bgr is None:
# approximate_frameCount was wrong
# reduce the trame index to get last frame
trys = 10
try_count = 1
while not ok and try_count <= trys:
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index - try_count)
ok, bgr = cap.read()
try_count += 1
msg = "read frame {} faild for video {} with appr. {} frames reduced to {}".format(
frame_index, self._vid_path, self.approximate_frameCount,
frame_index - try_count)
log.warning(msg)
cap.release()
assert bgr is not None, "faild reading frame for video frame {}, header frames {}".format(
frame_index, self.count_frames())
if (self.resize_shape is not None
and self.frameWidth != self.resize_shape[1]
and self.frameHeight != self.resize_shape[0]):
# cv2 mat size is fliped -> ::-1
bgr = cv2.resize(bgr, np_shape_to_cv(self.resize_shape),
cv2.INTER_NEAREST)
if self.to_rgb:
bgr = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
if self.dtype in [np.float32, np.float64]:
return np.asarray(bgr / 255.0, dtype=self.dtype)
else:
bgr = np.asarray(bgr, dtype=self.dtype)
if self.torch_transformer is not None:
bgr = self.torch_transformer(bgr)
return bgr