Update app.py

app.py

@@ -1,71 +1,77 @@
 import gradio as gr
 import numpy as np
 from PIL import Image
-import cv2  # OpenCV for video processing (if used)
+import cv2
+import os
+import uuid
+
 from extract_frames import video_to_keyframes
 from apply_mask import apply_mask_and_crop
 from run_gmm import run_gmm_inference
 from compose_video import compose_final_video
-# import the processing functions from original app
-# from heatmap_module import video_to_keyframes, apply_mask_and_crop, run_gmm_inference, compose_final_video
 
-# Helper to extract first frame for mask drawing
+
+# Ensure folders exist
+for path in [
+    "video_outputs/extracted_frames",
+    "video_outputs/masked_frames",
+    "video_outputs/output_heatmap",
+    "video_inputs",
+    "assets"
+]:
+    os.makedirs(path, exist_ok=True)
+
+
+# Get first frame for preview
 def get_first_frame(video_path):
     cap = cv2.VideoCapture(video_path)
     success, frame = cap.read()
     cap.release()
     if success:
-        # Convert BGR to RGB color for PIL/Gradio
         frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
         return Image.fromarray(frame)
-    else:
-        return None
-
-# Helper to get mask from drawn data
-def extract_mask_from_drawn(composite_image, background_image):
-    # Convert to numpy arrays for comparison
-    comp = np.array(composite_image)
-    bg = np.array(background_image)
-    if comp.shape != bg.shape:
-        # If background not same shape as composite, just threshold comp
-        gray = comp if comp.ndim == 2 else comp[..., :3].mean(axis=-1)
-        mask = (gray > 10).astype(np.uint8)  # simple threshold
-    else:
-        # Compute difference where composite != background (assuming draw color != background)
-        diff = np.any(comp != bg, axis=-1)
-        mask = diff.astype(np.uint8)
-    return mask * 255  # return as binary mask image (255 inside mask)
-
-def process_video(video_file, mask_image, drawn_editor, progress=gr.Progress()):
-    import os
-    import uuid
-
-    # Prepare all output folders
+    return None
+
+
+def process_video(video_file, progress=gr.Progress()):
     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")
+
+    # Clear old frames
+    for folder in [extracted_dir, masked_dir, heatmap_dir]:
+        for f in os.listdir(folder):
+            os.remove(os.path.join(folder, f))
+
+    # Load default mask
+    mask_path = "assets/default_mask.png"
+    if not os.path.exists(mask_path):
+        raise gr.Error("❌ Default mask not found at 'assets/default_mask.png'")
     
-    os.makedirs("video_outputs/extracted_frames", exist_ok=True)
-    os.makedirs("video_outputs/masked_frames", exist_ok=True)
-    os.makedirs("video_outputs/output_heatmap", exist_ok=True)
-    os.makedirs("video_inputs", exist_ok=True)
-    os.makedirs("assets", exist_ok=True)
-
-    # Choose mask: from upload or drawing
-    if mask_image is not None:
-        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 and "composite" in drawn_editor and "background" in drawn_editor:
-        mask = extract_mask_from_drawn(drawn_editor)
-    else:
-        raise gr.Error("Please provide a valid mask (uploaded or drawn).")
+    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
+    if mask is None:
+        raise gr.Error("❌ Failed to load default mask.")
 
     progress(0, desc="Extracting keyframes...")
     video_to_keyframes(video_file, extracted_dir)
 
+    # Load first frame to align mask size
+    first_frame_name = sorted(os.listdir(extracted_dir))[0]
+    first_frame = cv2.imread(os.path.join(extracted_dir, first_frame_name))
+
+    if first_frame is None:
+        raise gr.Error("❌ Failed to read first extracted keyframe.")
+
+    if mask.shape != first_frame.shape[:2]:
+        mask = cv2.resize(mask, (first_frame.shape[1], first_frame.shape[0]))
+
+    # Optional: get bounding box (coords) of table region
+    coords = cv2.findNonZero(mask)
+    if coords is None:
+        raise gr.Error("❌ No table region detected in default mask.")
+    x, y, w, h = cv2.boundingRect(coords)
+
     progress(0.3, desc="Applying mask and cropping...")
     apply_mask_and_crop(extracted_dir, mask, masked_dir)
 
@@ -77,53 +83,58 @@ def process_video(video_file, mask_image, drawn_editor, progress=gr.Progress()):
     result_path = os.path.join(base_dir, video_name)
     compose_final_video(mask, heatmap_dir, extracted_dir, result_path)
 
-    progress(1.0, desc="Done!")
+    progress(1.0, desc="Done ✅")
+    return "✅ Heatmap video generated successfully!", result_path, result_path
 
-    return "✅ Heatmap video generated!", result_path, result_path
 
-
-# Define the Gradio app layout
+# Layout
 custom_css = """
-.gradio-container {background: url('/gradio_api/file=background.jpg') center/cover no-repeat !important;
-                   background-color: #000 !important;}
-.panel {max-width: 800px; margin: 2rem auto; padding: 2rem; background: rgba(30,30,30, 0.8); border-radius: 8px;}
+.gradio-container {
+    background: url('/gradio_api/file=background.jpg') center/cover no-repeat !important;
+    background-color: #000 !important;
+}
+.panel {
+    max-width: 800px;
+    margin: 2rem auto;
+    padding: 2rem;
+    background: rgba(30,30,30, 0.8);
+    border-radius: 8px;
+}
 """
-with gr.Blocks(theme=gr.themes.Monochrome(), css=custom_css, title="Heatmap Generator") as demo:
-    gr.Markdown("## 🎥 Heatmap Generator", elem_classes="panel")
+
+with gr.Blocks(theme=gr.themes.Monochrome(), css=custom_css, title="UV Scan - Table Heatmap") as demo:
+    gr.Markdown("## 🎥 UV Scan – Table Heatmap Generator", elem_classes="panel")
+
     with gr.Row(elem_classes="panel"):
         video_input = gr.Video(label="Upload Video", format="mp4")
-    with gr.Tabs(elem_classes="panel"):
-        with gr.Tab("Upload Mask"):
-            mask_upload = gr.Image(label="Upload Mask Image", type="numpy")
-        with gr.Tab("Draw Mask"):
-            draw_info = gr.Markdown("*Draw mask on the frame:* Use brush to highlight the region of interest.")
-            mask_draw = gr.ImageEditor(label="Draw Mask", type="pil")
-    # Buttons
+
     with gr.Row(elem_classes="panel"):
         generate_btn = gr.Button("🔥 Generate Heatmap", variant="primary")
         reset_btn = gr.Button("Reset")
-        download_btn = gr.File(label="Download Video")
-    # Status and output
+        download_btn = gr.File(label="⬇️ Download Video")
+
     with gr.Row(elem_classes="panel"):
-        status_text = gr.Markdown("")  # to show status or final message
+        status_text = gr.Markdown("")
+
     with gr.Row(elem_classes="panel"):
         output_video = gr.Video(label="Output Video")
-    # Event handlers
-    # When video is uploaded, extract a frame and set it in the draw component
-    def prep_frame_for_drawing(video_file):
-        if video_file is None:
-            return None
-        frame = get_first_frame(video_file)
-        return {'background': frame, 'composite': frame}  # initial EditorValue
-    video_input.change(fn=prep_frame_for_drawing, inputs=video_input, outputs=mask_draw)
-    # Generate button triggers processing
-    generate_btn.click(fn=process_video, inputs=[video_input, mask_upload, mask_draw], outputs=[status_text, output_video])
-    # After video is generated, enable download (bind the file path from output)
-    # (Gradio may automatically handle download if output_video has a file source)
-    generate_btn.click(fn=lambda vid: vid, inputs=output_video, outputs=download_btn) 
-    # Reset button clears all
-    reset_btn.click(fn=lambda: (None, None, None, "", None), inputs=[], 
-                    outputs=[video_input, mask_upload, mask_draw, status_text, output_video])
-# Launch (if running locally; on HF Spaces this is handled automatically)
+
+    def on_video_upload(video_file):
+        return get_first_frame(video_file)
+
+    video_input.change(fn=on_video_upload, inputs=video_input, outputs=None)
+
+    generate_btn.click(
+        fn=process_video,
+        inputs=[video_input],
+        outputs=[status_text, output_video, download_btn]
+    )
+
+    reset_btn.click(
+        fn=lambda: (None, "", None, None),
+        inputs=[],
+        outputs=[video_input, status_text, output_video, download_btn]
+    )
+
 if __name__ == "__main__":
     demo.launch()