diff --git a/src/video_mode.py b/src/video_mode.py index 69d593d..b9bc519 100644 --- a/src/video_mode.py +++ b/src/video_mode.py @@ -13,14 +13,29 @@ def open_path_as_images(path, maybe_depthvideo=False): """Takes the filepath, returns (fps, frames). Every frame is a Pillow Image object""" suffix = pathlib.Path(path).suffix - if suffix == '.gif': + if suffix.lower() == '.gif': frames = [] img = Image.open(path) for i in range(img.n_frames): img.seek(i) frames.append(img.convert('RGB')) return 1000 / img.info['duration'], frames - if suffix in ['.avi'] and maybe_depthvideo: + if suffix.lower() == '.mts': + import imageio_ffmpeg + import av + container = av.open(path) + frames = [] + for packet in container.demux(video=0): + for frame in packet.decode(): + # Convert the frame to a NumPy array + numpy_frame = frame.to_ndarray(format='rgb24') + # Convert the NumPy array to a Pillow Image + image = Image.fromarray(numpy_frame) + frames.append(image) + fps = float(container.streams.video[0].average_rate) + container.close() + return fps, frames + if suffix.lower() in ['.avi'] and maybe_depthvideo: try: import imageio_ffmpeg # Suppose there are in fact 16 bits per pixel @@ -40,7 +55,7 @@ def open_path_as_images(path, maybe_depthvideo=False): finally: if 'gen' in locals(): gen.close() - if suffix in ['.webm', '.mp4', '.avi']: + if suffix.lower() in ['.webm', '.mp4', '.avi']: from moviepy.video.io.VideoFileClip import VideoFileClip clip = VideoFileClip(path) frames = [Image.fromarray(x) for x in list(clip.iter_frames())]