Update app.py

app.py

@@ -19,7 +19,7 @@ import spaces
 from registry import get_model
 from core.describe_scene import describe_scene
 from core.process import process_image
-from core.input_handler import resolve_input, validate_video, validate_image, show_preview_from_upload, show_preview_from_url
+from core.input_handler import resolve_input, validate_video, validate_image
 from utils.helpers import format_error, generate_session_id
 from huggingface_hub import hf_hub_download
 
@@ -84,32 +84,80 @@ def handle(mode, media_upload, url, run_det, det_model, det_confidence, run_seg,
     if not media:
         return None, format_error("No valid input provided. Please check your upload or URL."), None
 
-    results = []
-    for single_media in media:
-        if isinstance(single_media, str):  # Video file
-            valid, err = validate_video(single_media)
-            if not valid:
-                return None, format_error(err), None
-            cap = cv2.VideoCapture(single_media)
-            ret, frame = cap.read()
-            cap.release()
-            if not ret:
-                return None, format_error("Failed to read video frame."), None
-            single_media = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-
-        if isinstance(single_media, Image.Image):
-            valid, err = validate_image(single_media)
-            if not valid:
-                return None, format_error(err), None
-            try:
-                return process_image(single_media, run_det, det_model, det_confidence, run_seg, seg_model, run_depth, depth_model, blend)
-            except timeout_decorator.timeout_decorator.TimeoutError:
-                logger.error("Image processing timed out.")
-                return None, format_error("Processing timed out. Try a smaller image or simpler model."), None
+    first_input = media[0]
+    
+    # --- VIDEO PATH ---
+    if isinstance(first_input, str) and first_input.lower().endswith((".mp4", ".mov", ".avi")):
+        valid, err = validate_video(first_input)
+        if not valid:
+            return None, format_error(err), None
+        try:
+            _, msg, output_video_path = process_video(
+                video_path=first_input,
+                run_det=run_det,
+                det_model=det_model,
+                det_confidence=det_confidence,
+                run_seg=run_seg,
+                seg_model=seg_model,
+                run_depth=run_depth,
+                depth_model=depth_model,
+                blend=blend
+            )
+            return gr.update(visible=False), msg, output_video_path  # Hide img_out, show vid_out via click handler
+        except Exception as e:
+            logger.error(f"Video processing failed: {e}")
+            return None, format_error(str(e)), None
+
+    # --- IMAGE INPUT ---
+    elif isinstance(first_input, Image.Image):
+        valid, err = validate_image(first_input)
+        if not valid:
+            return None, format_error(err), None
+        try:
+            result_img, msg, output_zip = process_image(
+                image=first_input,
+                run_det=run_det,
+                det_model=det_model,
+                det_confidence=det_confidence,
+                run_seg=run_seg,
+                seg_model=seg_model,
+                run_depth=run_depth,
+                depth_model=depth_model,
+                blend=blend
+            )
+            return result_img, msg, output_zip  # Show img_out only
+        except timeout_decorator.timeout_decorator.TimeoutError:
+            logger.error("Image processing timed out.")
+            return None, format_error("Processing timed out. Try a smaller image or simpler model."), None
+        except Exception as e:
+            logger.error(f"Image processing failed: {e}")
+            return None, format_error(str(e)), None
 
     logger.warning("Unsupported media type resolved.")
     log_runtime(start_time)
-    return None, format_error("Invalid input. Please check your upload or URL."), None
+    return None, format_error("Unsupported input type."), None
+
+def show_preview_from_upload(files):
+    if not files:
+        return gr.update(visible=False), gr.update(visible=False)
+    file = files[0]
+    path = file.name.lower()
+    if path.endswith((".png", ".jpg", ".jpeg", ".webp")):
+        img = Image.open(file).convert("RGB")
+        return gr.update(value=img, visible=True), gr.update(visible=False)
+    elif path.endswith((".mp4", ".mov", ".avi")):
+        return gr.update(visible=False), gr.update(value=path, visible=True)
+    return gr.update(visible=False), gr.update(visible=False)
+
+def show_preview_from_url(url_input):
+    if not url_input:
+        return gr.update(visible=False), gr.update(visible=False)
+    path = url_input.strip().lower()
+    if path.endswith((".png", ".jpg", ".jpeg", ".webp")):
+        return gr.update(value=url_input, visible=True), gr.update(visible=False)
+    elif path.endswith((".mp4", ".mov", ".avi")):
+        return gr.update(visible=False), gr.update(value=url_input, visible=True)
+    return gr.update(visible=False), gr.update(visible=False)
 
 # Gradio Interface
 with gr.Blocks() as demo:
@@ -131,10 +179,6 @@ with gr.Blocks() as demo:
             # URL input
             url = gr.Textbox(label="URL (Image/Video)", visible=False)
             
-            # Preview uploaded media
-            preview_image = gr.Image(label="Preview Image", visible=False)
-            preview_video = gr.Video(label="Preview Video", visible=False)
-
             # Toggle visibility
             def toggle_inputs(selected_mode):
                 return [
@@ -146,8 +190,8 @@ with gr.Blocks() as demo:
                 ]
             
             mode.change(toggle_inputs, inputs=mode, outputs=[media_upload, url])
-            media_upload.change(show_preview_from_upload, inputs=media_upload, outputs=[preview_image, preview_video])
-            url.submit(show_preview_from_url, inputs=url, outputs=[preview_image, preview_video])
+            media_upload.change(show_preview_from_upload, inputs=media_upload, outputs=[img_out, vid_out])
+            url.submit(show_preview_from_url, inputs=url, outputs=[img_out, vid_out])
 
             # Visibility logic function
             def toggle_visibility(checked):
@@ -203,9 +247,12 @@ with gr.Blocks() as demo:
             # single_img_preview = gr.Image(label="Preview (Image)", visible=False)
             # gallery_preview = gr.Gallery(label="Preview (Gallery)", columns=3, height="auto", visible=False)
             # video_preview = gr.Video(label="Preview (Video)", visible=False)        
-            img_out = gr.Image(label="Scene Blueprint")
+            # Only one is shown at a time — image or video
+            img_out = gr.Image(label="Preview / Processed Output", visible=False)
+            vid_out = gr.Video(label="Preview / Processed Output", visible=False)
             json_out = gr.JSON(label="Scene JSON")
             zip_out = gr.File(label="Download Results")
+
                 
     
     
@@ -219,10 +266,28 @@ with gr.Blocks() as demo:
 
 
     # Button Click Event
+    run.click(
+    handle,
+    inputs=[
+        mode, media_upload, url,
+        run_det, det_model, det_confidence,
+        run_seg, seg_model,
+        run_depth, depth_model,
+        blend
+    ],
+    outputs=[img_out, json_out, zip_out]
+)
+
     run.click(
         handle,
-        inputs=[mode, media_upload, url, run_det, det_model, det_confidence, run_seg, seg_model, run_depth, depth_model, blend],
-        outputs=[img_out, json_out, zip_out]
+        inputs=[
+            mode, media_upload, url,
+            run_det, det_model, det_confidence,
+            run_seg, seg_model,
+            run_depth, depth_model,
+            blend
+        ],
+        outputs=[vid_out, json_out, zip_out]
     )
 
     # Footer Section