Upload space application files

app.py

@@ -0,0 +1,180 @@
+import gradio as gr
+import torch
+import numpy as np
+import cv2
+from PIL import Image
+from transformers import DPTForDepthEstimation, DPTImageProcessor
+from gradio_client import Client, handle_file
+import tempfile
+import os
+
+# === DEVICE ===
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# === DEPTH MODEL ===
+def load_depth_model():
+    # DPTImageProcessor is the modern replacement for FeatureExtractor
+    model = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to(device)
+    processor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas")
+    return model, processor
+
+@torch.no_grad()
+def estimate_depth(image_pil, model, processor):
+    # Keep original size for restoration later
+    original_size = image_pil.size  # (width, height)
+    
+    # Preprocess (processor handles resizing internally for the model)
+    inputs = processor(images=image_pil, return_tensors="pt").to(device)
+    
+    depth = model(**inputs).predicted_depth
+    
+    # Interpolate depth back to ORIGINAL image size
+    depth = torch.nn.functional.interpolate(
+        depth.unsqueeze(1),
+        size=(original_size[1], original_size[0]), # torch expects (H, W)
+        mode="bicubic",
+        align_corners=False,
+    ).squeeze().detach().cpu().numpy()
+    
+    # Normalize
+    depth_min, depth_max = depth.min(), depth.max()
+    if depth_max - depth_min > 0:
+        return (depth - depth_min) / (depth_max - depth_min)
+    return depth
+
+def depth_to_disparity(depth, max_disp=30):
+    # Invert depth: close objects (bright) shift more
+    return depth * max_disp
+
+def generate_right_and_mask(image, disparity):
+    """
+    Vectorized shift operation. 100x faster than for-loops.
+    """
+    height, width = image.shape[:2]
+    
+    # Create a grid of coordinates
+    x_coords, y_coords = np.meshgrid(np.arange(width), np.arange(height))
+    
+    # Calculate target coordinates (shift pixels to the left for right eye)
+    # Note: Disparity logic depends on convergence plane. 
+    # Usually: Right Eye View = Original - Disparity
+    shift = disparity.astype(int)
+    target_x = x_coords - shift
+    
+    # Initialize output and mask
+    right = np.zeros_like(image)
+    mask = np.ones((height, width), dtype=np.uint8) * 255 # 255 = hole/inpainting area
+    
+    # Valid indices mask
+    valid_mask = (target_x >= 0) & (target_x < width)
+    
+    # Flatten arrays for advanced indexing
+    flat_y = y_coords[valid_mask]
+    flat_x_target = target_x[valid_mask]
+    flat_x_source = x_coords[valid_mask]
+    
+    # Assign pixels
+    # Note: In case of collision (two pixels mapping to same spot), 
+    # this simple method overwrites. For better results, Z-buffering is needed,
+    # but this is sufficient for basic stereo.
+    right[flat_y, flat_x_target] = image[flat_y, flat_x_source]
+    
+    # Update Mask: Areas that were written to are NOT holes (0)
+    mask[flat_y, flat_x_target] = 0
+    
+    return right, mask
+
+# === LAMA INPAINTING (Via Gradio Client) ===
+# Note: You need a valid Space that accepts image + mask. 
+# Using a popular LaMa space as reference.
+try:
+    # Attempt to connect to a public LaMa space
+    # You can change this string to "frxngb23/lama-inpainting-api" if that space 
+    # supports the API client, otherwise use "any-other-lama-space"
+    lama_client = Client("asif-k/LaMa-Inpainting") 
+except Exception as e:
+    print(f"Could not connect to external LaMa client: {e}")
+    lama_client = None
+
+def run_lama_inpainting(image_bgr, mask):
+    if lama_client is None:
+        print("LaMa client unavailable, returning unfilled image.")
+        return image_bgr
+
+    # Prepare files for Gradio Client
+    # Convert BGR (OpenCV) to RGB for PIL
+    img_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
+    
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f_img, \
+         tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f_mask:
+        
+        Image.fromarray(img_rgb).save(f_img.name)
+        Image.fromarray(mask).save(f_mask.name)
+        
+        try:
+            # Predict using the external space
+            # Note: The api_name="/predict" or parameters might vary per Space.
+            # You must check the "View API" button at the bottom of the target Space.
+            result_path = lama_client.predict(
+                image=handle_file(f_img.name),
+                mask=handle_file(f_mask.name),
+                api_name="/predict" 
+            )
+            
+            # Result is a filepath
+            res_img = Image.open(result_path)
+            return cv2.cvtColor(np.array(res_img), cv2.COLOR_RGB2BGR)
+            
+        except Exception as e:
+            print(f"Inpainting failed: {e}")
+            return image_bgr # Return original with holes if fail
+        finally:
+            # Cleanup
+            os.remove(f_img.name)
+            os.remove(f_mask.name)
+
+# === APP LOGIC ===
+depth_model, depth_processor = load_depth_model()
+
+def stereo_pipeline(image_pil):
+    if image_pil is None:
+        return None, None
+        
+    image_cv = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
+
+    # 1. Estimate Depth
+    depth = estimate_depth(image_pil, depth_model, depth_processor)
+    
+    # 2. Calculate Disparity
+    disparity = depth_to_disparity(depth)
+    
+    # 3. Shift Pixels
+    right_img, mask = generate_right_and_mask(image_cv, disparity)
+    
+    # 4. Inpaint Holes
+    # Pass the mask where 255 indicates holes to be filled
+    right_filled = run_lama_inpainting(right_img, mask)
+
+    left = image_pil
+    right = Image.fromarray(cv2.cvtColor(right_filled, cv2.COLOR_BGR2RGB))
+    
+    return left, right
+
+# === GRADIO UI ===
+with gr.Blocks(title="2D to 3D Stereo") as demo:
+    gr.Markdown("## 2D to 3D Stereo Generator")
+    gr.Markdown("Generates a stereo pair using Depth Estimation and LaMa Inpainting.")
+    
+    with gr.Row():
+        input_img = gr.Image(type="pil", label="Input Image")
+    
+    with gr.Row():
+        out_left = gr.Image(label="Left Eye")
+        out_right = gr.Image(label="Right Eye (Generated)")
+        
+    btn = gr.Button("Generate 3D")
+    
+    btn.click(fn=stereo_pipeline, inputs=input_img, outputs=[out_left, out_right])
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,8 @@
+gradio
+gradio_client
+torch
+numpy
+opencv-python
+pillow
+transformers
+scipy