support zip

app.py

@@ -14,7 +14,7 @@ from datetime import datetime
 import glob
 import gc
 import time
-
+import zipfile
 from stream3r.models.stream3r import STream3R
 from stream3r.stream_session import StreamSession
 from stream3r.models.components.utils.load_fn import load_and_preprocess_images
@@ -31,6 +31,49 @@ device = "cuda"
 
 model = STream3R.from_pretrained("yslan/STream3R")
 
+def handle_zip(zip_file):
+    outdir = "uploads"
+    os.makedirs(outdir, exist_ok=True)
+    with zipfile.ZipFile(zip_file.name, "r") as zf:
+        zf.extractall(outdir)
+    return f"Extracted {len(os.listdir(outdir))} files"
+
+# --- add near your imports/helpers ---
+ALLOWED_IMG_EXT = {".png", ".jpg", ".jpeg", ".bmp", ".webp"}
+
+def _is_within_dir(base_dir: str, path: str) -> bool:
+    # Prevent zip-slip: ensure extracted path stays inside base_dir
+    base_dir = os.path.abspath(base_dir)
+    path = os.path.abspath(path)
+    return os.path.commonpath([base_dir]) == os.path.commonpath([base_dir, path])
+
+def extract_images_from_zip(zip_path: str, outdir: str) -> list[str]:
+    """
+    Extracts only image files from a zip into outdir.
+    Returns list of extracted file paths.
+    """
+    os.makedirs(outdir, exist_ok=True)
+    extracted = []
+    with zipfile.ZipFile(zip_path, "r") as zf:
+        for member in zf.infolist():
+            # Skip directories and non-image files
+            name = member.filename
+            if name.endswith("/"):
+                continue
+            ext = os.path.splitext(name)[1].lower()
+            if ext not in ALLOWED_IMG_EXT:
+                continue
+            # Construct final path safely
+            dest_path = os.path.join(outdir, os.path.basename(name))
+            # Zip-slip guard (in case filename has ../ etc.)
+            if not _is_within_dir(outdir, dest_path):
+                continue
+            with zf.open(member) as src, open(dest_path, "wb") as dst:
+                shutil.copyfileobj(src, dst)
+            extracted.append(dest_path)
+    return extracted
+
+
 # -------------------------------------------------------------------------
 # 1) Core model inference
 # -------------------------------------------------------------------------
@@ -116,10 +159,13 @@ def run_model(target_dir: str, model: STream3R, mode: str="causal", streaming: b
 # -------------------------------------------------------------------------
 # 2) Handle uploaded video/images --> produce target_dir + images
 # -------------------------------------------------------------------------
-def handle_uploads(input_video, input_images):
+def handle_uploads(input_video, input_images, input_zip=None):
     """
-    Create a new 'target_dir' + 'images' subfolder, and place user-uploaded
-    images or extracted frames from video into it. Return (target_dir, image_paths).
+    Create a new 'target_dir' + 'images' subfolder.
+    - Copies uploaded images
+    - Optionally extracts images from a ZIP
+    - Optionally extracts frames from a video (1 fps)
+    Returns (target_dir, image_paths).
     """
     start_time = time.time()
     gc.collect()
@@ -130,36 +176,32 @@ def handle_uploads(input_video, input_images):
     target_dir = os.path.join("demo_cache", f"input_images_{timestamp}")
     target_dir_images = os.path.join(target_dir, "images")
 
-    # Clean up if somehow that folder already exists
     if os.path.exists(target_dir):
         shutil.rmtree(target_dir)
-    os.makedirs(target_dir)
-    os.makedirs(target_dir_images)
+    os.makedirs(target_dir_images, exist_ok=True)
 
-    image_paths = []
+    image_paths: list[str] = []
 
-    # --- Handle images ---
-    if input_images is not None:
+    # --- Handle images (list) ---
+    if input_images:
         for file_data in input_images:
-            if isinstance(file_data, dict) and "name" in file_data:
-                file_path = file_data["name"]
-            else:
-                file_path = file_data
+            file_path = file_data["name"] if isinstance(file_data, dict) and "name" in file_data else file_data
             dst_path = os.path.join(target_dir_images, os.path.basename(file_path))
             shutil.copy(file_path, dst_path)
             image_paths.append(dst_path)
 
-    # --- Handle video ---
-    if input_video is not None:
-        if isinstance(input_video, dict) and "name" in input_video:
-            video_path = input_video["name"]
-        else:
-            video_path = input_video
+    # --- Handle ZIP (extract images) ---
+    if input_zip:
+        zip_path = input_zip["name"] if isinstance(input_zip, dict) and "name" in input_zip else input_zip
+        extracted = extract_images_from_zip(zip_path, target_dir_images)
+        image_paths.extend(extracted)
 
+    # --- Handle video (extract frames at 1 fps) ---
+    if input_video:
+        video_path = input_video["name"] if isinstance(input_video, dict) and "name" in input_video else input_video
         vs = cv2.VideoCapture(video_path)
-        fps = vs.get(cv2.CAP_PROP_FPS)
-        frame_interval = int(fps * 1)  # 1 frame/sec
-
+        fps = vs.get(cv2.CAP_PROP_FPS) or 30.0
+        frame_interval = max(1, int(fps * 1))  # 1 frame/sec
         count = 0
         video_frame_num = 0
         while True:
@@ -172,30 +214,30 @@ def handle_uploads(input_video, input_images):
                 cv2.imwrite(image_path, frame)
                 image_paths.append(image_path)
                 video_frame_num += 1
+        vs.release()
 
-    # Sort final images for gallery
-    image_paths = sorted(image_paths)
+    image_paths = sorted(set(image_paths))  # de-dupe + sort
 
     end_time = time.time()
-    print(f"Files copied to {target_dir_images}; took {end_time - start_time:.3f} seconds")
+    print(f"Prepared {len(image_paths)} files in {target_dir_images}; took {end_time - start_time:.3f}s")
     return target_dir, image_paths
 
 
+
 # -------------------------------------------------------------------------
 # 3) Update gallery on upload
 # -------------------------------------------------------------------------
-def update_gallery_on_upload(input_video, input_images):
+def update_gallery_on_upload(input_video, input_images, input_zip):
     """
-    Whenever user uploads or changes files, immediately handle them
-    and show in the gallery. Return (target_dir, image_paths).
-    If nothing is uploaded, returns "None" and empty list.
+    Handle any new uploads (video, images, or zip) and render preview.
     """
-    if not input_video and not input_images:
+    if not input_video and not input_images and not input_zip:
         return None, None, None, None
-    target_dir, image_paths = handle_uploads(input_video, input_images)
+    target_dir, image_paths = handle_uploads(input_video, input_images, input_zip)
     return None, target_dir, image_paths, "Upload complete. Click 'Reconstruct' to begin 3D processing."
 
 
+
 # -------------------------------------------------------------------------
 # 4) Reconstruction: uses the target_dir plus any viz parameters
 # -------------------------------------------------------------------------
@@ -460,7 +502,8 @@ with gr.Blocks(
     with gr.Row():
         with gr.Column(scale=2):
             input_video = gr.Video(label="Upload Video", interactive=True)
-            input_images = gr.File(file_count="multiple", label="Upload Images", interactive=True, max_files=100)
+            input_images = gr.File(file_count="multiple", label="Upload Images", interactive=True)
+            input_zip = gr.File(file_types=[".zip"], label="Upload ZIP of Images", interactive=True)
 
             image_gallery = gr.Gallery(
                 label="Preview",
@@ -757,13 +800,18 @@ with gr.Blocks(
     # -------------------------------------------------------------------------
     input_video.change(
         fn=update_gallery_on_upload,
-        inputs=[input_video, input_images],
+        inputs=[input_video, input_images, input_zip],
         outputs=[reconstruction_output, target_dir_output, image_gallery, log_output],
     )
     input_images.change(
         fn=update_gallery_on_upload,
-        inputs=[input_video, input_images],
+        inputs=[input_video, input_images, input_zip],
+        outputs=[reconstruction_output, target_dir_output, image_gallery, log_output],
+    )
+    input_zip.change(
+        fn=update_gallery_on_upload,
+        inputs=[input_video, input_images, input_zip],
         outputs=[reconstruction_output, target_dir_output, image_gallery, log_output],
     )
 
-    demo.queue(max_size=20).launch(show_error=True, share=True)
+    demo.queue(max_size=20).launch(show_error=True, share=False)