Update app.py

app.py

@@ -37,36 +37,47 @@ def extract_mask_from_drawn(composite_image, background_image):
     return mask * 255  # return as binary mask image (255 inside mask)
 
 def process_video(video_file, mask_image, drawn_editor, progress=gr.Progress()):
-    # video_file: path to uploaded video
-    # mask_image: numpy array (HxW or HxWx3) if uploaded, or None
-    # drawn_editor: dict with 'background', 'composite' from ImageEditor, or None
-    
-    # Decide mask source
-    mask = None
+    import os
+    import uuid
+
+    # Prepare all output folders
+    base_dir = "video_outputs"
+    extracted_dir = os.path.join(base_dir, "extracted_frames")
+    masked_dir = os.path.join(base_dir, "masked_frames")
+    heatmap_dir = os.path.join(base_dir, "output_heatmap")
+    os.makedirs(extracted_dir, exist_ok=True)
+    os.makedirs(masked_dir, exist_ok=True)
+    os.makedirs(heatmap_dir, exist_ok=True)
+
+    # Choose mask: from upload or drawing
     if mask_image is not None:
-        # Ensure mask is binary (if user uploaded a colored mask, convert to gray)
         mask = mask_image
         if mask.ndim == 3:
             mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
         _, mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
-    elif drawn_editor is not None:
-        comp = drawn_editor["composite"]
-        bg = drawn_editor["background"]
-        mask = extract_mask_from_drawn(comp, bg)
+    elif drawn_editor and "composite" in drawn_editor and "background" in drawn_editor:
+        mask = extract_mask_from_drawn(drawn_editor)
     else:
-        raise gr.Error("Please provide a mask (upload or draw).")
-    
+        raise gr.Error("Please provide a valid mask (uploaded or drawn).")
+
     progress(0, desc="Extracting keyframes...")
-    frames = video_to_keyframes(video_file)
+    video_to_keyframes(video_file, extracted_dir)
+
     progress(0.3, desc="Applying mask and cropping...")
-    cropped_frames = apply_mask_and_crop(frames, mask)
+    apply_mask_and_crop(extracted_dir, mask, masked_dir)
+
     progress(0.6, desc="Running inference on frames...")
-    output_frames = run_gmm_inference(cropped_frames)
+    run_gmm_inference(masked_dir, heatmap_dir)
+
     progress(0.85, desc="Composing final video...")
-    result_path = compose_final_video(output_frames, "heatmap_output.mp4")
-    progress(1.0, desc="Done")
-    
-    return "✅ Heatmap video generated!", result_path
+    video_name = f"heatmap_output_{uuid.uuid4().hex[:6]}.mp4"
+    result_path = os.path.join(base_dir, video_name)
+    compose_final_video(mask, heatmap_dir, extracted_dir, result_path)
+
+    progress(1.0, desc="Done!")
+
+    return "✅ Heatmap video generated!", result_path, result_path
+
 
 # Define the Gradio app layout
 custom_css = """