Skip to content

Commit b8120b4

Browse files
authored
Update video_mode.py (#368)
Added support for video files in MTS format
1 parent f60f666 commit b8120b4

File tree

1 file changed

+18
-3
lines changed

1 file changed

+18
-3
lines changed

src/video_mode.py

+18-3
Original file line numberDiff line numberDiff line change
@@ -13,14 +13,29 @@
1313
def open_path_as_images(path, maybe_depthvideo=False):
1414
"""Takes the filepath, returns (fps, frames). Every frame is a Pillow Image object"""
1515
suffix = pathlib.Path(path).suffix
16-
if suffix == '.gif':
16+
if suffix.lower() == '.gif':
1717
frames = []
1818
img = Image.open(path)
1919
for i in range(img.n_frames):
2020
img.seek(i)
2121
frames.append(img.convert('RGB'))
2222
return 1000 / img.info['duration'], frames
23-
if suffix in ['.avi'] and maybe_depthvideo:
23+
if suffix.lower() == '.mts':
24+
import imageio_ffmpeg
25+
import av
26+
container = av.open(path)
27+
frames = []
28+
for packet in container.demux(video=0):
29+
for frame in packet.decode():
30+
# Convert the frame to a NumPy array
31+
numpy_frame = frame.to_ndarray(format='rgb24')
32+
# Convert the NumPy array to a Pillow Image
33+
image = Image.fromarray(numpy_frame)
34+
frames.append(image)
35+
fps = float(container.streams.video[0].average_rate)
36+
container.close()
37+
return fps, frames
38+
if suffix.lower() in ['.avi'] and maybe_depthvideo:
2439
try:
2540
import imageio_ffmpeg
2641
# Suppose there are in fact 16 bits per pixel
@@ -40,7 +55,7 @@ def open_path_as_images(path, maybe_depthvideo=False):
4055
finally:
4156
if 'gen' in locals():
4257
gen.close()
43-
if suffix in ['.webm', '.mp4', '.avi']:
58+
if suffix.lower() in ['.webm', '.mp4', '.avi']:
4459
from moviepy.video.io.VideoFileClip import VideoFileClip
4560
clip = VideoFileClip(path)
4661
frames = [Image.fromarray(x) for x in list(clip.iter_frames())]

0 commit comments

Comments
 (0)