diff --git a/transcoder/frame_grabber.py b/transcoder/frame_grabber.py index 8ef99b6..6e56cd4 100644 --- a/transcoder/frame_grabber.py +++ b/transcoder/frame_grabber.py @@ -44,9 +44,11 @@ class FileFrameGrabber(FrameGrabber): yield Image.fromarray(frame_array) @staticmethod - def _output_dir(filename) -> str: - # TODO: should include palette - return ".".join(filename.split(".")[:-1]) + def _output_dir(filename, video_mode, palette) -> str: + return "%s/%s/%s" % ( + ".".join(filename.split(".")[:-1]), + video_mode.name, + palette.name) def _palette_arg(self) -> str: return "P%d" % self.palette.value @@ -57,11 +59,9 @@ class FileFrameGrabber(FrameGrabber): We do the encoding in a background thread to parallelize. """ - frame_dir = self._output_dir(self.filename) - try: - os.mkdir(frame_dir) - except FileExistsError: - pass + frame_dir = self._output_dir( + self.filename, self.video_mode, self.palette) + os.makedirs(frame_dir, exist_ok=True) q = queue.Queue(maxsize=10) diff --git a/transcoder/frame_grabber_test.py b/transcoder/frame_grabber_test.py index e424ba3..aee23ec 100644 --- a/transcoder/frame_grabber_test.py +++ b/transcoder/frame_grabber_test.py @@ -1,23 +1,35 @@ import unittest import frame_grabber +import palette +import video_mode class TestFileFrameGrabber(unittest.TestCase): def test_output_dir(self): self.assertEqual( - "/foo/bar", - frame_grabber.FileFrameGrabber._output_dir("/foo/bar.mp4") + "/foo/bar/DHGR/NTSC", + frame_grabber.FileFrameGrabber._output_dir( + "/foo/bar.mp4", video_mode.VideoMode.DHGR, palette.Palette.NTSC + ) ) self.assertEqual( - "/foo/bar.blee", - frame_grabber.FileFrameGrabber._output_dir("/foo/bar.blee.mp4") + "/foo/bar.blee/HGR/IIGS", + frame_grabber.FileFrameGrabber._output_dir( + "/foo/bar.blee.mp4", + video_mode.VideoMode.HGR, + palette.Palette.IIGS + ) ) self.assertEqual( - "/foo/bar blee", - frame_grabber.FileFrameGrabber._output_dir("/foo/bar blee.mp4") + "/foo/bar blee/DHGR/IIGS", + frame_grabber.FileFrameGrabber._output_dir( + "/foo/bar blee.mp4", + video_mode.VideoMode.DHGR, + palette.Palette.IIGS + ) )