Update main.py

main.py

@@ -1,10 +1,10 @@
 """
-Marksman AI - Professional Spyne-Grade Engine
-Logic: Original Shadow Extraction + Multi-Layer Compositing + Floor Reflection
+Marksman AI - SPYNE ENTERPRISE REPLICATION
+Logic: Distance-Field Shadows, Multiply Blending, Grounded Reflection
 """
 
 import io, logging, traceback, numpy as np, cv2, torch
-from PIL import Image, ImageOps, ImageFilter
+from PIL import Image, ImageFilter
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
@@ -12,11 +12,11 @@ from fastapi.middleware.cors import CORSMiddleware
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-app = FastAPI(title="Marksman AI Pro API", version="6.0.0")
+app = FastAPI(title="Spyne-Grade Engine", version="7.0.0")
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
 # ──────────────────────────────────────────────────────────────────
-# 1. AI ENSEMBLE (Models)
+# 1. MODELS (Clean Logic)
 # ──────────────────────────────────────────────────────────────────
 _models = {}
 
@@ -24,94 +24,91 @@ def get_model(name: str):
     if name not in _models:
         device = "cuda" if torch.cuda.is_available() else "cpu"
         if name == "birefnet":
-            from transformers import AutoModelForImageSegmentation
+            from transformers import AutoModelForImageSegmentation, AutoConfig
             from torchvision import transforms
             model = AutoModelForImageSegmentation.from_pretrained("ZhengPeng7/BiRefNet_dynamic", trust_remote_code=True)
             model.to(device).eval().float()
-            _models[name] = {
-                "model": model, 
-                "device": device,
-                "transform": transforms.Compose([
-                    transforms.Resize((1024, 1024)),
-                    transforms.ToTensor(),
-                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-                ])
-            }
+            _models[name] = {"model": model, "device": device, "transform": transforms.Compose([
+                transforms.Resize((1024, 1024)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+            ])}
         elif name == "yolo":
             from ultralytics import YOLO
             _models[name] = {"model": YOLO("yolov8n-cls.pt")}
     return _models[name]
 
 # ──────────────────────────────────────────────────────────────────
-# 2. THE SECRET SAUCE: COMPOSITING LOGIC
+# 2. THE SPYNE "SECRET": MULTIPLY BLENDING & SHADOW FIELD
 # ──────────────────────────────────────────────────────────────────
 
-def apply_spyne_lighting(car_np, bg_np, mask):
-    """Gaari ke colors ko background se match karne ki logic (Harmonization)"""
-    bg_avg = bg_np.mean(axis=(0,1))
-    car_avg = car_np[mask > 0.5].mean(axis=0)
-    
-    # 5% tint from background for realism
-    res = car_np.astype(np.float32)
-    for i in range(3):
-        res[..., i] = (res[..., i] * 0.95) + (bg_avg[i] * 0.05)
-        
-    # Luminance matching
-    diff = (bg_avg.mean() * 0.9 - car_avg.mean()) * 0.2
-    return np.clip(res + diff, 0, 255).astype(np.uint8)
+def multiply_blend(background, shadow_mask, opacity=0.7):
+    """
+    Asli Spyne logic: Shadow ko 'Paste' nahi karte, background ko 'Darken' karte hain.
+    Is se background ki tiles/road shadow ke niche nazar aati hain.
+    """
+    bg_float = background.astype(np.float32) / 255.0
+    # Shadow mask ko intensity mein convert karen
+    mask = shadow_mask.astype(np.float32) / 255.0
+    mask = mask[..., None] # Add channel dim
+    
+    # Multiply: New_Pixel = Old_Pixel * (1 - Mask * Opacity)
+    darkened = bg_float * (1.0 - (mask * opacity))
+    return (darkened * 255).astype(np.uint8)
 
-def create_pro_shadow(car_mask, bg_size, pos, v_type="low"):
-    """Multi-layer shadow: Silhouette + Contact + Ambient"""
+def generate_spyne_shadow(car_mask, bg_size, pos, v_type="low"):
+    """
+    Distance Transform based shadow: Jo car ke kareeb dark aur door jate blur ho jaye.
+    """
     cw, ch = car_mask.size
-    px, py = pos
-    mask_np = np.array(car_mask).astype(np.float32) / 255.0
+    mask_np = (np.array(car_mask) > 128).astype(np.uint8) * 255
+    
+    # 1. Create a flattened silhouette
+    s_h = int(ch * (0.22 if v_type == "high" else 0.14))
+    silhouette = cv2.resize(mask_np, (cw, s_h))
     
-    # 1. SILHOUETTE SQUASH
-    s_h = int(ch * (0.18 if v_type == "high" else 0.12))
-    s_base = cv2.resize(mask_np, (cw, s_h))
+    # 2. Distance Field (Edges ko soft karne ke liye)
+    dist_map = cv2.distanceTransform(silhouette, cv2.DIST_L2, 5)
+    dist_map = cv2.normalize(dist_map, None, 0, 1.0, cv2.NORM_MINMAX)
     
-    # 2. DUAL BLUR (Contact vs Ambient)
-    contact = cv2.GaussianBlur(s_base, (7, 3), 0) # Sharp grounding
-    ambient = cv2.GaussianBlur(s_base, (int(cw*0.08)|1, int(s_h*0.7)|1), 0) # Soft spread
+    # 3. Layer 1: Contact (Super sharp thin line)
+    contact = cv2.GaussianBlur(silhouette, (5, 3), 0)
     
-    combined = (contact * 0.75) + (ambient * 0.25)
+    # Layer 2: Ambient (The soft spread)
+    ambient_blur = int(cw * 0.1) | 1
+    ambient = cv2.GaussianBlur(silhouette, (ambient_blur, int(s_h * 0.8) | 1), 0)
     
-    canvas = Image.new("RGBA", bg_size, (0,0,0,0))
-    shadow_img = Image.fromarray((combined * 160).astype(np.uint8), "L") # 160 = Opacity
+    # Mix layers
+    final_mask = (contact.astype(float) * 0.6) + (ambient.astype(float) * 0.4)
     
-    # Position: Anchor to tire line
-    canvas.paste((0,0,0,255), (px, py + ch - int(s_h * 0.4)), mask=shadow_img)
+    # Full Canvas placement
+    canvas = np.zeros((bg_size[1], bg_size[0]), dtype=np.uint8)
+    # Anchor: Spyne hamesha shadow ko car ke 40% height offset par rakhta hai
+    target_y = pos[1] + ch - int(s_h * 0.42)
+    
+    y1, y2 = max(target_y, 0), min(target_y + s_h, bg_size[1])
+    x1, x2 = pos[0], min(pos[0] + cw, bg_size[0])
+    
+    if (y2 - y1) > 0:
+        canvas[y1:y2, x1:x2] = final_mask[:y2-y1, :x2-x1].astype(np.uint8)
+        
     return canvas
 
-def create_pro_reflection(car_rgba, bg_size, pos):
-    """Subtle floor reflection logic"""
-    cw, ch = car_rgba.size
-    refl = car_rgba.transpose(Image.FLIP_TOP_BOTTOM).resize((cw, int(ch * 0.25)), Image.LANCZOS)
-    refl = refl.filter(ImageFilter.GaussianBlur(radius=8))
-    
-    canvas = Image.new("RGBA", bg_size, (0,0,0,0))
-    canvas.paste(refl, (pos[0], pos[1] + ch - 4))
-    
-    # Fade out reflection
-    r_np = np.array(canvas).astype(np.float32)
-    r_np[..., 3] *= 0.06 # Spyne standard reflection opacity (6%)
-    return Image.fromarray(r_np.astype(np.uint8), "RGBA")
-
 # ──────────────────────────────────────────────────────────────────
-# 3. THE ENSEMBLE PIPELINE
+# 3. THE PIPELINE
 # ──────────────────────────────────────────────────────────────────
 
-def run_spyne_pipeline(car_pil, bg_pil):
-    # Standardize
+def run_pro_pipeline(car_pil, bg_pil):
+    # Resize
     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)
 
-    # A. Classification (SUV vs Sedan)
+    # A. Classification
     yolo = get_model("yolo")["model"]
     res_cls = yolo(car_pil, verbose=False)
     v_type = "high" if any(x in res_cls[0].names[res_cls[0].probs.top1].lower() for x in ['suv','truck','van','jeep']) else "low"
 
-    # B. AI Masking
+    # B. Segmentation
     biref = get_model("birefnet")
     inp = biref["transform"](car_pil.convert("RGB")).unsqueeze(0).to(biref["device"]).float()
     with torch.no_grad():
@@ -119,41 +116,50 @@ def run_spyne_pipeline(car_pil, bg_pil):
         preds = biref["model"](inp)
     mask = cv2.resize(torch.sigmoid(preds[-1]).squeeze().cpu().numpy(), car_pil.size)
     
-    # C. Auto-Positioning (Size 95%)
+    # C. Sizing (95% Spyne Standard)
     target_w = int(bg_pil.size[0] * 0.95)
     scale = target_w / car_pil.size[0]
     car_res = car_pil.resize((target_w, int(car_pil.size[1]*scale)), Image.LANCZOS)
-    mask_res = cv2.resize(mask, car_res.size)
+    mask_res = (cv2.resize(mask, car_res.size) * 255).astype(np.uint8)
     
-    # Ground detection
+    # D. Ground Detection
     bg_gray = np.array(bg_pil.convert("L"))
     edges = cv2.Canny(bg_gray[int(bg_pil.size[1]*0.7):, :], 50, 150)
     ground_y = int(bg_pil.size[1]*0.7 + np.argmax(np.sum(edges, axis=1))) if edges.max() > 0 else int(bg_pil.size[1]*0.88)
     
     px = (bg_pil.size[0] - target_w) // 2
-    py = ground_y - int(np.where(mask_res > 0.5)[0].max())
+    # Find exact tire contact point
+    bottom_y = np.where(mask_res > 128)[0].max()
+    py = ground_y - bottom_y
     pos = (px, max(5, min(py, bg_pil.size[1] - car_res.size[1] - 5)))
 
-    # D. Compositing (Layering)
-    car_harmonized = apply_spyne_lighting(np.array(car_res), np.array(bg_pil), mask_res)
-    car_final = Image.fromarray(np.dstack([car_harmonized, (mask_res*255).astype(np.uint8)]), "RGBA")
-    
-    result = bg_pil.convert("RGBA")
-    # 1. Reflection Layer
-    result = Image.alpha_composite(result, create_pro_reflection(car_final, bg_pil.size, pos))
-    # 2. Multi-Layer Shadow Layer
-    result = Image.alpha_composite(result, create_pro_shadow(car_final.split()[-1], bg_pil.size, pos, v_type))
-    # 3. Top Car Layer
-    result.paste(car_final, pos, car_final)
-    
-    return result.convert("RGB")
+    # E. The Compositing (NumPy Multiply)
+    bg_np = np.array(bg_pil.convert("RGB"))
+    car_np = np.array(car_res.convert("RGB"))
+    
+    # 1. Shadow Step (Multiply Blend)
+    shadow_mask = generate_spyne_shadow(Image.fromarray(mask_res), bg_pil.size, pos, v_type)
+    # Background ko shadow se darken karo (Opacity 0.75 for realism)
+    blended_bg = multiply_blend(bg_np, shadow_mask, opacity=0.75)
+    
+    # 2. Reflection Step (Subtle)
+    final_img = Image.fromarray(blended_bg).convert("RGBA")
+    refl = car_res.transpose(Image.FLIP_TOP_BOTTOM).resize((target_w, int(car_res.size[1]*0.25)))
+    refl_mask = Image.fromarray((np.array(refl.convert("L")) * 0.08).astype(np.uint8)) # 8% opacity
+    final_img.paste(refl, (pos[0], pos[1] + car_res.size[1] - 2), mask=refl_mask)
+    
+    # 3. Car Step
+    car_rgba = Image.fromarray(np.dstack([car_np, mask_res]))
+    final_img.paste(car_rgba, pos, car_rgba)
+    
+    return final_img.convert("RGB")
 
 @app.post("/composite")
 async def composite(car_image: UploadFile = File(...), background_image: UploadFile = File(...)):
     try:
         c_pil = Image.open(io.BytesIO(await car_image.read())).convert("RGB")
         b_pil = Image.open(io.BytesIO(await background_image.read())).convert("RGB")
-        res = run_spyne_pipeline(c_pil, b_pil)
+        res = run_pro_pipeline(c_pil, b_pil)
         buf = io.BytesIO()
         res.save(buf, format="JPEG", quality=95)
         buf.seek(0)