Update shadow_generator.py

shadow_generator.py

@@ -229,7 +229,6 @@ from PIL import Image, ImageFilter, ImageChops
 import torch
 from torchvision import transforms as T
 import numpy as np
-import cv2
 
 # -----------------------------
 # Paths & constants
@@ -345,8 +344,7 @@ class SSNWrapper:
         try:
             from SSN_Model import SSN_Model
             self.model = SSN_Model()
-            with torch.serialization.safe_globals([np.core.multiarray.scalar]):
-                state = torch.load(str(WEIGHT_FILE), map_location=self.device, weights_only=False)
+            state = torch.load(str(WEIGHT_FILE), map_location=self.device)
             if isinstance(state, dict):
                 if "model_state_dict" in state:
                     sd = state["model_state_dict"]
@@ -407,7 +405,7 @@ def load_ssn_once() -> SSNWrapper:
     return _ssn_wrapper
 
 # -----------------------------
-# Shadow compositing (Attached Shadow)
+# Shadow compositing (REALISTIC SSN)
 # -----------------------------
 def _hex_to_rgb(color: str) -> Tuple[int, int, int]:
     c = color.strip().lstrip("#")
@@ -417,32 +415,15 @@ def _hex_to_rgb(color: str) -> Tuple[int, int, int]:
         raise ValueError("Invalid color hex.")
     return tuple(int(c[i:i+2],16) for i in (0,2,4))
 
-def _warp_shadow(matte_img: Image.Image, w: int, h: int, direction: float, distance: float) -> Image.Image:
-    matte = np.array(matte_img.resize((w, h)).convert("L"))
-
-    angle = np.deg2rad(direction)
-    dx = int(np.cos(angle) * distance)
-    dy = int(np.sin(angle) * distance)
-
-    src = np.float32([[0,0],[w,0],[0,h],[w,h]])
-    dst = np.float32([[0,0],[w,0],[0+dx,h],[w+dx,h+dy]])
-
-    M = cv2.getPerspectiveTransform(src, dst)
-    warped = cv2.warpPerspective(matte, M, (w,h), borderValue=0)
-
-    return Image.fromarray(warped)
-
 def _composite_shadow_and_image(original_img: Image.Image, matte_img: Image.Image, params: Dict[str, Any]) -> Image.Image:
-    _log("Compositing shadow and image (anchored)...")
+    _log("Compositing shadow and image (REAL SSN)...")
     w, h = original_img.size
     r, g, b = _hex_to_rgb(params.get("color", "#000000"))
 
-    matte_warped = _warp_shadow(matte_img, w, h, params["direction"], params["distance"])
-    obj_alpha = original_img.getchannel("A")
-    matte_clean = ImageChops.multiply(matte_warped, obj_alpha)
+    matte_resized = matte_img.resize((w, h)).convert("L")
 
     shadow_rgba = Image.new("RGBA", (w, h), (r, g, b, 0))
-    alpha = matte_clean.point(lambda p: int(p * float(params.get("opacity", 0.7))))
+    alpha = matte_resized.point(lambda p: int(p * float(params.get("opacity", 0.7))))
     shadow_rgba.putalpha(alpha)
 
     softness = max(0.0, float(params.get("softness", 20.0)))
@@ -453,22 +434,14 @@ def _composite_shadow_and_image(original_img: Image.Image, matte_img: Image.Imag
     final_rgba.alpha_composite(shadow_rgba)
     final_rgba.alpha_composite(original_img)
 
-    # --- NEW: Option to flatten background ---
-    if params.get("background", "transparent") != "transparent":
-        bg_color = params.get("background", "#FFFFFF")
-        br, bgc, bb = _hex_to_rgb(bg_color)
-        flattened = Image.new("RGB", (w, h), (br, bgc, bb))
-        flattened.paste(final_rgba, mask=final_rgba.getchannel("A"))
-        return flattened.convert("RGBA")
-
-    _log("Composition complete with anchored shadow.")
+    _log("✅ Composition complete with realistic SSN shadow.")
     return final_rgba
 
 # -----------------------------
 # Public API
 # -----------------------------
 def _apply_params_defaults(params: Optional[Dict[str, Any]]) -> Dict[str, Any]:
-    defaults = dict(softness=20.0, opacity=0.7, color="#000000", direction=45.0, distance=80.0, background="transparent")
+    defaults = dict(softness=20.0, opacity=0.7, color="#000000", direction=45.0, distance=80.0)
     merged = {**defaults, **(params or {})}
     merged["opacity"] = max(0.0, min(1.0, float(merged["opacity"])))
     merged["softness"] = max(0.0, float(merged["softness"]))
@@ -487,10 +460,10 @@ def generate_shadow_rgba(rgba_file_bytes: bytes, params: Optional[Dict[str, Any]
         raise RuntimeError("SSN model not available.")
     
     matte_img = ssn.infer_shadow_matte(img, params["direction"])
-
     if matte_img is None:
         _log("❌ Failed to generate shadow matte")
         raise RuntimeError("Failed to generate shadow matte.")
+    
     final_img = _composite_shadow_and_image(img, matte_img, params)
     buf = io.BytesIO()
     final_img.save(buf, format="PNG")