Update main.py

main.py

@@ -1,6 +1,6 @@
 """
-Marksman AI Pro - Multi-Model Ensemble Engine (Spyne Architecture)
-Models: YOLOv8-Cls (Perception) + BiRefNet (Segmentation) + OpenCV (Grounding)
+Marksman AI - Professional Spyne-Grade Engine
+Logic: Original Shadow Extraction + Multi-Layer Compositing + Floor Reflection
 """
 
 import io, logging, traceback, numpy as np, cv2, torch
@@ -12,11 +12,11 @@ from fastapi.middleware.cors import CORSMiddleware
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-app = FastAPI(title="Marksman AI Spyne-Grade API", version="5.0.0")
+app = FastAPI(title="Marksman AI Pro API", version="6.0.0")
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
 # ──────────────────────────────────────────────────────────────────
-# 1. ENSEMBLE MODEL LOADING
+# 1. AI ENSEMBLE (Models)
 # ──────────────────────────────────────────────────────────────────
 _models = {}
 
@@ -39,154 +39,123 @@ def get_model(name: str):
             }
         elif name == "yolo":
             from ultralytics import YOLO
-            model = YOLO("yolov8n-cls.pt")
-            model.to(device)
-            _models[name] = {"model": model}
+            _models[name] = {"model": YOLO("yolov8n-cls.pt")}
     return _models[name]
 
 # ──────────────────────────────────────────────────────────────────
-# 2. PHYSICS LOGIC: GROUNDING & SHADOWS
+# 2. THE SECRET SAUCE: COMPOSITING LOGIC
 # ──────────────────────────────────────────────────────────────────
 
-def generate_multi_layer_shadow(car_mask, bg_size, pos, v_type="low"):
-    """Spyne's Ultra-Soft Shadow Blend"""
-    bg_w, bg_h = bg_size
+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 create_pro_shadow(car_mask, bg_size, pos, v_type="low"):
+    """Multi-layer shadow: Silhouette + Contact + Ambient"""
     cw, ch = car_mask.size
     px, py = pos
     mask_np = np.array(car_mask).astype(np.float32) / 255.0
-
-    # Shadow Settings for White Floors (Subtle is key)
-    h_ratio = 0.16 if v_type == "high" else 0.10
-    # White background par opacity 110-130 best hai
-    opacity = 125 
-    
-    shadow_h = int(ch * h_ratio)
-    shadow_base = cv2.resize(mask_np, (cw, shadow_h))
     
-    # Layer 1: Contact (Tires ko chipkanay ke liye - Very Tight)
-    contact = cv2.GaussianBlur(shadow_base, (5, 3), 0)
-    # Layer 2: Ambient (Soft spread)
-    ambient = cv2.GaussianBlur(shadow_base, (int(cw * 0.1)|1, int(shadow_h * 0.9)|1), 0)
+    # 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))
     
-    # Blending Layers: 80% Contact, 20% Ambient
-    combined = (contact * 0.8) + (ambient * 0.2)
+    # 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
     
-    canvas = np.zeros((bg_h, bg_w), dtype=np.float32)
-    # POSITION FIX: Shadow ko tires ke thora sa aur upar "anchor" karna (30% offset)
-    target_y = py + ch - int(shadow_h * 0.45) 
-    
-    y1, y2 = max(target_y, 0), min(target_y + shadow_h, bg_h)
-    x1, x2 = px, min(px + cw, bg_w)
-    
-    if (y2 - y1) > 0:
-        canvas[y1:y2, x1:x2] = combined[:y2-y1, :x2-x1]
-        
-    shadow_rgba = np.zeros((bg_h, bg_w, 4), dtype=np.uint8)
-    shadow_rgba[..., 3] = (np.clip(canvas, 0, 1) * opacity).astype(np.uint8)
-    return Image.fromarray(shadow_rgba, "RGBA")
-
-def harmonize_lighting(car_rgb, car_mask, bg_rgb):
-    """Studio Brightness Adjustment"""
-    # Agar background white hai, toh car ko thora 'bright' hona chahiye
-    res = car_rgb.astype(np.float32)
+    combined = (contact * 0.75) + (ambient * 0.25)
     
-    # Global brightness boost for white studios (5% boost)
-    res = np.clip(res * 1.05, 0, 255)
+    canvas = Image.new("RGBA", bg_size, (0,0,0,0))
+    shadow_img = Image.fromarray((combined * 160).astype(np.uint8), "L") # 160 = Opacity
     
-    # Environment Matching
-    avg_env = bg_rgb.mean(axis=(0,1))
-    for c in range(3):
-        res[..., c] = (res[..., c] * 0.94) + (avg_env[c] * 0.06)
-        
-    return res.astype(np.uint8)
+    # Position: Anchor to tire line
+    canvas.paste((0,0,0,255), (px, py + ch - int(s_h * 0.4)), mask=shadow_img)
+    return canvas
 
-def generate_subtle_reflection(car_rgba, bg_size, pos):
-    """Studio-grade floor reflection"""
+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.22)), Image.LANCZOS)
-    refl = refl.filter(ImageFilter.GaussianBlur(radius=6))
+    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))
-    # Paste with 7% opacity (Spyne Standard)
     canvas.paste(refl, (pos[0], pos[1] + ch - 4))
+    
+    # Fade out reflection
     r_np = np.array(canvas).astype(np.float32)
-    r_np[..., 3] *= 0.07 
+    r_np[..., 3] *= 0.06 # Spyne standard reflection opacity (6%)
     return Image.fromarray(r_np.astype(np.uint8), "RGBA")
 
 # ──────────────────────────────────────────────────────────────────
-# 3. ENSEMBLE PIPELINE
+# 3. THE ENSEMBLE PIPELINE
 # ──────────────────────────────────────────────────────────────────
 
-def run_ensemble_pipeline(car_pil, bg_pil):
-    # Standardize input for AI models
+def run_spyne_pipeline(car_pil, bg_pil):
+    # 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)
 
-    # A. Perception Phase (Classification)
-    yolo_bundle = get_model("yolo")
-    cls_res = yolo_bundle["model"](car_pil, verbose=False)
-    v_class = cls_res[0].names[cls_res[0].probs.top1].lower()
-    v_type = "high" if any(x in v_class for x in ['suv','truck','van','jeep','bus']) else "low"
-    logger.info(f"Spyne Engine Classified: {v_type.upper()}")
+    # A. Classification (SUV vs Sedan)
+    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. Extraction Phase (High-Res Masking)
-    biref_bundle = get_model("birefnet")
-    inp = biref_bundle["transform"](car_pil.convert("RGB")).unsqueeze(0).to(biref_bundle["device"]).float()
+    # B. AI Masking
+    biref = get_model("birefnet")
+    inp = biref["transform"](car_pil.convert("RGB")).unsqueeze(0).to(biref["device"]).float()
     with torch.no_grad():
-        biref_bundle["model"].float()
-        preds = biref_bundle["model"](inp)
-    
+        biref["model"].float()
+        preds = biref["model"](inp)
     mask = cv2.resize(torch.sigmoid(preds[-1]).squeeze().cpu().numpy(), car_pil.size)
     
-    # C. Positioning Phase (Auto-Scale to 95%)
+    # C. Auto-Positioning (Size 95%)
     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)
     
-    # D. Ground Perspective Detection
+    # 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.90)
+    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
-    # Find the lowest tire pixel to anchor the car
     py = ground_y - int(np.where(mask_res > 0.5)[0].max())
-    py = max(5, min(py, bg_pil.size[1] - car_res.size[1] - 5))
-    pos = (px, py)
+    pos = (px, max(5, min(py, bg_pil.size[1] - car_res.size[1] - 5)))
 
-    # E. Color Harmonization
-    bg_np = np.array(bg_pil)
-    sky_tone = bg_np[:int(bg_np.shape[0]*0.2), :].mean(axis=(0,1))
-    ground_tone = bg_np[int(bg_np.shape[0]*0.8):, :].mean(axis=(0,1))
-    env_avg = (sky_tone + ground_tone) / 2
+    # 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")
     
-    car_np = np.array(car_res).astype(np.float32)
-    for c in range(3):
-        car_np[..., c] = (car_np[..., c] * 0.96) + (env_avg[c] * 0.04)
+    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)
     
-    car_ready = Image.fromarray(np.dstack([car_np.astype(np.uint8), (mask_res*255).astype(np.uint8)]), "RGBA")
-
-    # F. Multi-Layer Compositing (Order: Reflection -> Shadow -> Car)
-    final_canvas = bg_pil.convert("RGBA")
-    final_canvas = Image.alpha_composite(final_canvas, generate_subtle_reflection(car_ready, bg_pil.size, pos))
-    final_canvas = Image.alpha_composite(final_canvas, generate_multi_layer_shadow(car_ready.split()[-1], bg_pil.size, pos, v_type))
-    final_canvas.paste(car_ready, pos, car_ready)
-    
-    return final_canvas.convert("RGB")
-
-# ──────────────────────────────────────────────────────────────────
-# 4. API ENDPOINT
-# ──────────────────────────────────────────────────────────────────
+    return result.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")
-        result = run_ensemble_pipeline(c_pil, b_pil)
+        res = run_spyne_pipeline(c_pil, b_pil)
         buf = io.BytesIO()
-        result.save(buf, format="JPEG", quality=95)
+        res.save(buf, format="JPEG", quality=95)
         buf.seek(0)
         return StreamingResponse(buf, media_type="image/jpeg")
     except Exception as e: