Update main.py

main.py

@@ -144,6 +144,43 @@ def generate_shadow(car_rgba, bg, paste_pos, actual_tyre_y, v_type="low"):
     shadow_rgba = np.zeros((bg_h, bg_w, 4), dtype=np.uint8)
     shadow_rgba[..., 3] = (np.clip(shadow_canvas, 0, 1) * 220).astype(np.uint8)
     return Image.fromarray(shadow_rgba, "RGBA")
+# ──────────────────────────────────────────────────────────────────
+# NEW: SPYNE-GRADE SHADOW & REFLECTION LOGIC
+# ──────────────────────────────────────────────────────────────────
+
+def extract_original_shadow(car_pil, mask_np):
+    """Original photo se asli shadow nikalne ka logic."""
+    orig_np = np.array(car_pil.convert("L")) # Grayscale
+    # Mask ko thora phailao (dilate) taake tires ke niche ka area mil jaye
+    kernel = np.ones((20,20), np.uint8)
+    dilated_mask = cv2.dilate(mask_np, kernel, iterations=2)
+    shadow_area_mask = dilated_mask - mask_np
+    
+    # Shadow map: Light areas ko white aur dark areas (shadow) ko black/gray karo
+    # Original ground agar whiteish hai toh ye bohot acha kaam karta hai
+    shadow_map = orig_np.astype(float) / 255.0
+    shadow_map = np.clip(shadow_map * 1.2, 0, 1) # Contrast barhao
+    return shadow_map, shadow_area_mask
+
+def generate_reflection(car_rgba, bg_size, pos):
+    """Car ki halki si perchaayi (reflection) floor par."""
+    bg_w, bg_h = bg_size
+    cw, ch = car_rgba.size
+    px, py = pos
+    
+    # Car ko flip karo aur blur kardo
+    reflection = car_rgba.transpose(Image.FLIP_TOP_BOTTOM)
+    reflection = reflection.resize((cw, int(ch * 0.3)), Image.LANCZOS)
+    
+    # Reflection Canvas
+    refl_canvas = Image.new("RGBA", (bg_w, bg_h), (0,0,0,0))
+    # Car ke bilkul niche paste karo (tyre line par)
+    refl_canvas.paste(reflection, (px, py + ch - 5))
+    
+    # Bohat halki opacity (Spyne level is 5-10%)
+    refl_np = np.array(refl_canvas).astype(np.float32)
+    refl_np[..., 3] *= 0.08 
+    return Image.fromarray(refl_np.astype(np.uint8), "RGBA")
 
 # ──────────────────────────────────────────────────────────────────
 # CORE PIPELINE & POSITIONING
@@ -203,51 +240,48 @@ def auto_position_car(car_rgba: Image.Image, bg: Image.Image):
     return car_res, (px, py), tyre_y
 
 def run_pipeline(car_pil, bg_pil):
-    # Standardize sizes
+    # 1. Standardize
     car_pil = car_pil.resize((1024, int(1024 * car_pil.height/car_pil.width)), Image.LANCZOS)
     bg_pil = bg_pil.resize((1600, int(1600 * bg_pil.height/bg_pil.width)), Image.LANCZOS)
 
-    # 1. Classification (New)
     v_type = classify_vehicle(car_pil)
-    logger.info(f"Detected Vehicle Type: {v_type}")
-
+    
     # 2. Segmentation
-    import torch
     bundle = get_model("birefnet")
-    orig_w, orig_h = car_pil.size
     inp = bundle["transform"](car_pil.convert("RGB")).unsqueeze(0).to(bundle["device"]).float()
     with torch.no_grad():
         bundle["model"].float()
-      #  inp = inp.float()
         preds = bundle["model"](inp)
     mask = torch.sigmoid(preds[-1]).squeeze().cpu().numpy()
-    mask = cv2.resize(mask, (orig_w, orig_h))
+    mask = cv2.resize(mask, car_pil.size)
     
-    car_rgb = np.array(car_pil.convert("RGB"))
-    car_rgba = Image.fromarray(np.dstack([car_rgb, (mask*255).astype(np.uint8)]), "RGBA")
-    
-    # 3. Positioning
-    car_positioned, pos, actual_tyre_y = auto_position_car(car_rgba, bg_pil)
+    # 3. Extract Original Shadow (Secret Sauce)
+    orig_shadow_map, shadow_mask = extract_original_shadow(car_pil, mask)
+
+    # 4. Positioning
+    car_rgba_temp = Image.fromarray(np.dstack([np.array(car_pil), (mask*255).astype(np.uint8)]), "RGBA")
+    car_positioned, pos, actual_tyre_y = auto_position_car(car_rgba_temp, bg_pil)
     
-    # 4. Color Matching
+    # 5. Color Matching
     c_arr = np.array(car_positioned)
     rgb_fixed = harmonize_color(c_arr[..., :3], c_arr[..., 3]/255.0, np.array(bg_pil))
     car_final = Image.fromarray(np.dstack([rgb_fixed, c_arr[..., 3]]), "RGBA")
+
+    # 6. COMPOSITING LAYERS (Spyne Style)
+    bg_final = bg_pil.convert("RGBA")
     
-    # 5. Dynamic Shadow Generation
-    shadow_layer = generate_shadow(car_final, bg_pil, pos, actual_tyre_y, v_type=v_type)
+    # Layer A: Ambient Shadow (Jo humne pehle banayi thi)
+    ambient_layer = generate_shadow(car_final, bg_pil, pos, actual_tyre_y, v_type=v_type)
+    bg_final = Image.alpha_composite(bg_final, ambient_layer)
     
-    # 6. Final Compositing
-    bg_np = np.array(bg_pil.convert("RGB")).astype(np.float32)
-    sh_np = np.array(shadow_layer).astype(np.float32)
-    sh_alpha = (sh_np[..., 3] / 255.0)[..., None]
+    # Layer B: Reflection (New)
+    reflection_layer = generate_reflection(car_final, bg_pil.size, pos)
+    bg_final = Image.alpha_composite(bg_final, reflection_layer)
     
-    # Multiply effect for shadow
-    darkened_bg = bg_np * (1.0 - sh_alpha * 0.92)
-    output_img = Image.fromarray(darkened_bg.astype(np.uint8))
-    output_img.paste(car_final, pos, car_final)
+    # Layer C: The Car
+    bg_final.paste(car_final, pos, car_final)
     
-    return output_img.convert("RGB")
+    return bg_final.convert("RGB")
 
 @app.post("/composite")
 async def composite(car_image: UploadFile = File(...), background_image: UploadFile = File(...)):