nermadie/2.5D_Depth_Studio

1from fastapi import FastAPI, File, UploadFile2from fastapi.middleware.cors import CORSMiddleware3from fastapi.responses import JSONResponse4import torch5import cv26import numpy as np7from PIL import Image8import io9import base6410from transformers import DPTImageProcessor, DPTForDepthEstimation11from scipy.ndimage import binary_dilation, binary_erosion, gaussian_filter12 13app = FastAPI()14 15app.add_middleware(16    CORSMiddleware,17    allow_origins=["*"],18    allow_credentials=True,19    allow_methods=["*"],20    allow_headers=["*"],21)22 23device = "cuda" if torch.cuda.is_available() else "cpu"24print(f"Using device: {device}")25 26# Load depth model (prefer the larger model when available)27try:28    processor = DPTImageProcessor.from_pretrained("Intel/dpt-large")29    model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")30except:31    processor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas")32    model = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas")33 34model.to(device)35model.eval()36 37 38def predict_depth(image: Image.Image):39    """Generate high-quality depth map with enhanced background detail"""40    inputs = processor(images=image, return_tensors="pt").to(device)41 42    with torch.no_grad():43        outputs = model(**inputs)44        depth = outputs.predicted_depth45 46    depth = depth.squeeze().cpu().numpy()47 48    # Resize to original size49    depth = cv2.resize(50        depth, (image.width, image.height), interpolation=cv2.INTER_CUBIC51    )52 53    # Normalize while preserving details54    depth = (depth - depth.min()) / (depth.max() - depth.min())55 56    # Apply CLAHE with a higher clipLimit to increase contrast/detail in the background57    depth = (depth * 255).astype(np.uint8)58    clahe = cv2.createCLAHE(clipLimit=5.0, tileGridSize=(8, 8))59    depth = clahe.apply(depth)60 61    # Apply bilateral filter to preserve edges while smoothing noise62    depth = cv2.bilateralFilter(depth, 9, 75, 75)63 64    # Light normalization to increase variation in the background65    depth = cv2.normalize(depth, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)66 67    return depth68 69 70def edge_aware_inpaint(image, mask, edge_mask=None):71    """72    Smarter inpainting that better respects edges.73    """74    # Dilate mask to cover the full inpaint region75    kernel = np.ones((11, 11), np.uint8)76    mask_dilated = cv2.dilate(mask, kernel, iterations=3)77 78    # Use TELEA instead of NS for more natural results79    inpainted = cv2.inpaint(80        image, mask_dilated, inpaintRadius=5, flags=cv2.INPAINT_TELEA81    )82 83    # Blend with surrounding pixels84    mask_blur = cv2.GaussianBlur(mask_dilated.astype(float) / 255.0, (21, 21), 0)85    mask_blur = np.stack([mask_blur] * 3, axis=-1)86 87    result = inpainted * mask_blur + image * (1 - mask_blur)88 89    return result.astype(np.uint8)90 91 92def create_soft_mask(hard_mask, blur_amount=15):93    """Create a soft mask with more natural edges."""94    # Morphological operations to smooth95    kernel = np.ones((5, 5), np.uint8)96    mask = cv2.morphologyEx(hard_mask, cv2.MORPH_CLOSE, kernel, iterations=2)97    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1)98 99    # Gaussian blur for soft edges100    mask_float = mask.astype(np.float32) / 255.0101    mask_soft = cv2.GaussianBlur(mask_float, (0, 0), sigmaX=blur_amount)102 103    # Adjust mask contrast104    mask_soft = np.power(mask_soft, 0.8)  # Softer falloff105    mask_soft = np.clip(mask_soft, 0, 1)106 107    return mask_soft108 109 110def advanced_layer_separation(image_np, depth):111    """112    Automatically split the image into multiple layers based on the depth map.113    Facebook-style with a configurable number of layers.114    """115    NUM_LAYERS = 128  # Increase to 128 layers for richer background detail116 117    h, w = image_np.shape[:2]118 119    # Ensure depth matches image120    if depth.shape[:2] != (h, w):121        depth = cv2.resize(depth, (w, h), interpolation=cv2.INTER_CUBIC)122 123    # Create percentile-based thresholds, focusing more density in mid-range124    depth_flat = depth.flatten()125    # Use non-uniform percentiles: 70% background, 25% midground, 5% foreground126    bg_layers = int(NUM_LAYERS * 0.7)  # 70% for background127    mg_layers = int(NUM_LAYERS * 0.25)  # 25% for midground128    fg_layers = NUM_LAYERS - bg_layers - mg_layers  # 5% for foreground129 130    percentiles = np.concatenate(131        [132            np.linspace(0, 35, bg_layers),  # Background: 70% layers, high detail133            np.linspace(35, 80, mg_layers),  # Midground: 25% layers134            np.linspace(80, 100, fg_layers + 1),  # Foreground: 5% layers, fewer layers135        ]136    )137    thresholds = [np.percentile(depth_flat, p) for p in percentiles]138 139    def create_soft_mask_adaptive(mask, blur_amount, is_foreground=False):140        """Create soft alpha edges (foreground needs more blur)."""141        kernel_size = 9 if is_foreground else 7142        kernel = np.ones((kernel_size, kernel_size), np.uint8)143        mask_clean = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)144        mask_clean = cv2.dilate(mask_clean, kernel, iterations=1)145        mask_float = mask_clean.astype(np.float32) / 255.0146        # Foreground uses more blur for smoother edges147        blur_sigma = (blur_amount / 10.0) * (2.0 if is_foreground else 1.0)148        mask_blur = cv2.GaussianBlur(mask_float, (0, 0), blur_sigma)149        return np.clip(mask_blur, 0.0, 1.0)150 151    layers = []152 153    # Layer 0: Backplate (full opaque background)154    all_fg_mask = (depth >= thresholds[1]).astype(np.uint8) * 255155    inpaint_mask = cv2.dilate(all_fg_mask, np.ones((7, 7), np.uint8), iterations=2)156    backplate = edge_aware_inpaint(image_np, inpaint_mask)157    backplate = cv2.GaussianBlur(backplate, (7, 7), 1.5)158    bp_alpha = np.ones((h, w), dtype=np.uint8) * 255159    layers.append(np.dstack([backplate, bp_alpha]))160 161    # Create layers from far to near162    for i in range(NUM_LAYERS):163        # Create mask for this layer164        if i == NUM_LAYERS - 1:165            mask = (depth >= thresholds[i]).astype(np.uint8) * 255166        else:167            mask = ((depth >= thresholds[i]) & (depth < thresholds[i + 1])).astype(168                np.uint8169            ) * 255170 171        # Determine if this layer is foreground (top 5%)172        is_foreground = i >= (NUM_LAYERS - fg_layers)173 174        # Smooth mask (foreground uses more blur)175        blur_amount = max(8, 16 - i // 4)176        mask_soft = create_soft_mask_adaptive(mask, blur_amount, is_foreground)177 178        # Only inpaint the first 8 layers (background) to optimize speed179        if i < 8:180            # Compute combined mask of all layers in front181            combined_fg = np.zeros_like(depth, dtype=np.uint8)182            for j in range(i + 1, NUM_LAYERS):183                if j == NUM_LAYERS - 1:184                    fg = (depth >= thresholds[j]).astype(np.uint8)185                else:186                    fg = (187                        (depth >= thresholds[j]) & (depth < thresholds[j + 1])188                    ).astype(np.uint8)189                combined_fg = np.clip(combined_fg + fg, 0, 1)190 191            # Inpaint background192            if combined_fg.sum() > 100:193                inpaint_mask_layer = cv2.dilate(194                    combined_fg * 255, np.ones((5, 5), np.uint8), iterations=1195                )196                layer_img = edge_aware_inpaint(image_np, inpaint_mask_layer)197                # Reduce blur as depth increases198                blur_radius = max(1, 5 - i // 2)199                if blur_radius > 1:200                    layer_img = cv2.GaussianBlur(201                        layer_img,202                        (blur_radius * 2 + 1, blur_radius * 2 + 1),203                        blur_radius / 2.0,204                    )205            else:206                layer_img = image_np.copy()207        else:208            # Near layers don't need inpainting209            layer_img = image_np.copy()210 211        # Build alpha channel212        alpha = (mask_soft * 255).astype(np.uint8)213        layer = np.dstack([layer_img, alpha])214        layers.append(layer)215 216    return layers, depth217 218 219def image_to_base64(img_array, has_alpha=False):220    """Convert numpy array to base64"""221    if has_alpha:222        img = Image.fromarray(img_array, mode="RGBA")223    else:224        img = Image.fromarray(img_array)225 226    buffered = io.BytesIO()227    img.save(buffered, format="PNG", optimize=True)228    return base64.b64encode(buffered.getvalue()).decode()229 230 231@app.post("/api/process")232async def process_image(file: UploadFile = File(...)):233    """234    Process image with Facebook-quality 3D effect235    """236    try:237        # Read image238        contents = await file.read()239        image = Image.open(io.BytesIO(contents)).convert("RGB")240 241        # Resize to optimize performance while keeping quality242        max_size = 1024  # Increased from 800 to 1024243        if max(image.size) > max_size:244            ratio = max_size / max(image.size)245            new_size = tuple(int(dim * ratio) for dim in image.size)246            image = image.resize(new_size, Image.Resampling.LANCZOS)247 248        image_np = np.array(image)249 250        # Generate depth map251        print("Generating depth map...")252        depth = predict_depth(image)253 254        # Create advanced layers255        print("Creating layers...")256        layers, depth_processed = advanced_layer_separation(image_np, depth)257 258        # Convert to base64 and generate metadata for N layers259        layers_b64 = []260        num_layers = len(layers)261 262        for i, layer in enumerate(layers):263            layer_b64 = image_to_base64(layer, has_alpha=True)264 265            # Automatically compute depth and name for each layer266            if i == 0:267                layer_name = "backplate"268                layer_depth = 0.0269            else:270                # Linear depth from 0 to 1271                layer_depth = (i - 1) / (num_layers - 2) if num_layers > 2 else 0.5272                layer_name = f"layer_{i-1}"273 274            layers_b64.append(275                {276                    "index": i,277                    "name": layer_name,278                    "depth": layer_depth,279                    "data": f"data:image/png;base64,{layer_b64}",280                }281            )282 283        # Original image284        image_b64 = image_to_base64(image_np)285 286        # Depth visualization287        depth_colored = cv2.applyColorMap(depth_processed, cv2.COLORMAP_MAGMA)288        depth_colored = cv2.cvtColor(depth_colored, cv2.COLOR_BGR2RGB)289        depth_b64 = image_to_base64(depth_colored)290 291        # Depth data for mesh rendering (normalized 0..1)292        depth_normalized = (depth_processed.astype(np.float32) / 255.0).tolist()293 294        print("Processing complete!")295 296        return JSONResponse(297            {298                "success": True,299                "image": f"data:image/png;base64,{image_b64}",300                "depth": f"data:image/png;base64,{depth_b64}",301                "depth_data": depth_normalized,302                "layers": layers_b64,303                "width": image.width,304                "height": image.height,305                "use_mesh": True,  # Flag for the frontend to use mesh rendering306            }307        )308 309    except Exception as e:310        import traceback311 312        traceback.print_exc()313        return JSONResponse({"success": False, "error": str(e)}, status_code=500)314 315 316@app.get("/")317def root():318    return {319        "message": "Facebook 3D Photo Backend - Improved Version",320        "status": "running",321        "device": device,322    }323 324 325if __name__ == "__main__":326    import uvicorn327 328    uvicorn.run(app, host="0.0.0.0", port=8000)329