feat: make the frame_format configurable.

intelygenz · Oct 24, 2024 · 9cb7c22 · 9cb7c22
1 parent 677074a
commit 9cb7c22
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/phalp/configs/base.py b/phalp/configs/base.py
@@ -122,6 +122,7 @@ class ModelConfig:
 @dataclass
 class ExtraConfig:
     FOCAL_LENGTH: int = 5000
+    FRAME_FORMAT: str = "jpg"
 
 @dataclass
 class FullConfig:

diff --git a/phalp/utils/io.py b/phalp/utils/io.py
@@ -53,8 +53,8 @@ def get_frames_from_source(self):
 
                 fe = FrameExtractor(source_path)
                 log.info('Number of frames: ' + str(fe.n_frames))
-                fe.extract_frames(every_x_frame=1, img_name='', dest_path= self.cfg.video.output_dir + "/_DEMO/" + video_name + "/img/", start_frame=self.cfg.video.start_frame, end_frame=self.cfg.video.end_frame)
-                list_of_frames = sorted(glob.glob(self.cfg.video.output_dir + "/_DEMO/" + video_name + "/img/*.jpg"))
+                fe.extract_frames(every_x_frame=1, img_name='', dest_path= self.cfg.video.output_dir + "/_DEMO/" + video_name + "/img/", img_ext="." + self.cfg.EXTRA.FRAME_FORMAT, start_frame=self.cfg.video.start_frame, end_frame=self.cfg.video.end_frame)
+                list_of_frames = sorted(glob.glob(self.cfg.video.output_dir + "/_DEMO/" + video_name + "/img/*." + self.cfg.EXTRA.FRAME_FORMAT))
             else:
                 start_time, end_time = int(self.cfg.video.start_time[:-1]), int(self.cfg.video.end_time[:-1])
                 try:
@@ -85,7 +85,7 @@ def get_frames_from_source(self):
         # read from image folder
         elif(os.path.isdir(source_path)):
             video_name = source_path.split('/')[-1]
-            list_of_frames = sorted(glob.glob(source_path + "/*.jpg"))
+            list_of_frames = sorted(glob.glob(source_path + "/*." + self.cfg.EXTRA.FRAME_FORMAT))
 
         # pkl files are used to track ground truth videos with given bounding box
         # these gt_id, gt_bbox will be stored in additional_data, ground truth bbox should be in the format of [x1, y1, w, h]