Merge pull request #13 from bladeszasza/bugfix/generate_per_object_videos

Bugfix/generate per object videos

Author: bladeszasza

sowlv2/video_utils.py

@@ -4,10 +4,12 @@
 """
 import os
 from glob import glob
+import re
 import cv2  # pylint: disable=import-error
 from PIL import Image
 import numpy as np
 
+
 # Disable no-member for cv2 (OpenCV) for the whole file
 # pylint: disable=no-member
 
@@ -65,31 +67,88 @@ def images_to_video(image_files, video_path, fps=30):
     video_writer.release()
     print(f"Saved video {video_path}")
 
+def _parse_mask_filename(fname):
+    """
+    Parse a mask filename to extract sam_id_token and core_prompt_slug.
+    Returns (sam_id_token, core_prompt_slug) or (None, None) if not matched.
+    """
+    # Example: 000001_obj1_dog_mask.png
+    match = re.match(r"^\d+_(obj\d+)_([a-zA-Z0-9_]+)_mask\.png$", fname)
+    if match:
+        return match.group(1), match.group(2)
+    # Fallback: 000001_obj1_mask.png (no prompt)
+    match_simple = re.match(r"^\d+_(obj\d+)_mask\.png$", fname)
+    if match_simple:
+        return match_simple.group(1), None
+    return None, None
+
+def _collect_unique_tracked_objects(mask_files):
+    """
+    Collect unique (sam_id_token, core_prompt_slug) pairs from mask filenames.
+    Returns a dict with keys as (sam_id_token, core_prompt_slug).
+    """
+    unique_tracked_objects = {}
+    for f_path in mask_files:
+        fname = os.path.basename(f_path)
+        sam_id_token, core_prompt_slug = _parse_mask_filename(fname)
+        if sam_id_token is not None:
+            key = (sam_id_token, core_prompt_slug)
+            if key not in unique_tracked_objects:
+                unique_tracked_objects[key] = {
+                    "sam_id_token": sam_id_token,
+                    "core_prompt_slug": core_prompt_slug
+                }
+        else:
+            print(f"Warning: Filename {fname} did not match expected pattern.")
+    return unique_tracked_objects
+
+def _get_obj_files(output_dir, sam_id_token, core_prompt_slug):
+    """
+    Get sorted mask and overlay files for a given object.
+    """
+    if core_prompt_slug:
+        mask_pattern = os.path.join(
+            output_dir, f"*_{sam_id_token}_{core_prompt_slug}_mask.png")
+        overlay_pattern = os.path.join(
+            output_dir, f"*_{sam_id_token}_{core_prompt_slug}_overlay.png")
+        video_prefix = f"{sam_id_token}_{core_prompt_slug}"
+    else:
+        mask_pattern = os.path.join(output_dir, f"*_{sam_id_token}_mask.png")
+        overlay_pattern = os.path.join(output_dir, f"*_{sam_id_token}_overlay.png")
+        video_prefix = sam_id_token
+    mask_files = sorted(glob(mask_pattern))
+    overlay_files = sorted(glob(overlay_pattern))
+    return mask_files, overlay_files, video_prefix
+
 def generate_per_object_videos(output_dir, fps=30):
     """
     Generate per-object videos from mask and overlay images.
-    Each object will have its own video for masks and overlays.
+    Each object (identified by sam_id and core_prompt) will have its own
+    video for masks and overlays.
     """
-    mask_pattern = os.path.join(output_dir, "*_obj*_mask.png")
-    mask_files = sorted(glob(mask_pattern))
+    all_mask_files_pattern = os.path.join(output_dir, "*_mask.png")
+    all_mask_files = sorted(glob(all_mask_files_pattern))
 
-    objects = set()
-    for f in mask_files:
-        try:
-            # Assuming filename format like '000001_obj1_mask.png'
-            obj_id = os.path.basename(f).split('_')[1] # Extracts 'obj1'
-            objects.add(obj_id)
-        except IndexError:
-            print(f"Warning: Could not parse object ID from filename {f}. Skipping.")
-            continue
+    if not all_mask_files:
+        print(f"No mask files found in {output_dir} matching pattern.")
+        return
+
+    unique_tracked_objects = _collect_unique_tracked_objects(all_mask_files)
+    if not unique_tracked_objects:
+        print(f"No objects successfully parsed from filenames in {output_dir}.")
+        return
 
+    for key in sorted(unique_tracked_objects.keys()):
+        obj_info = unique_tracked_objects[key]
+        sam_id_token = obj_info["sam_id_token"]
+        core_prompt_slug = obj_info["core_prompt_slug"]
 
-    for obj in sorted(list(objects)): # Convert set to sorted list for deterministic order
-        obj_mask_files = sorted(glob(os.path.join(output_dir, f"*_{obj}_mask.png")))
-        obj_overlay_files = sorted(glob(os.path.join(output_dir, f"*_{obj}_overlay.png")))
+        mask_files, overlay_files, video_file_prefix = _get_obj_files(
+            output_dir, sam_id_token, core_prompt_slug
+        )
 
-        mask_video_path = os.path.join(output_dir, f"{obj}_mask_video.mp4")
-        overlay_video_path = os.path.join(output_dir, f"{obj}_overlay_video.mp4")
+        mask_video_path = os.path.join(output_dir, f"{video_file_prefix}_mask_video.mp4")
+        overlay_video_path = os.path.join(output_dir, f"{video_file_prefix}_overlay_video.mp4")
 
-        images_to_video(obj_mask_files, mask_video_path, fps)
-        images_to_video(obj_overlay_files, overlay_video_path, fps)
+        images_to_video(mask_files, mask_video_path, fps)
+        images_to_video(overlay_files, overlay_video_path, fps)