import io import os import sys import uuid from pathlib import Path from PIL import Image, ImageFilter import requests from app.config import settings class RembgOnnxEngine: """ Apparel-Specialized Background Removal Engine with Collar Badge & Inner Label Protection. Preserves fine inner details: collar labels, brand badges, neck linings, buttons, and pockets while cleanly stripping studio backgrounds and shadows. """ _session = None MODEL_URL = "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_cloth_seg.onnx" MODEL_PATH = Path.home() / ".u2net" / "u2net_cloth_seg.onnx" @classmethod def get_session(cls): if cls._session is not None: return cls._session try: import onnxruntime as ort cls.MODEL_PATH.parent.mkdir(parents=True, exist_ok=True) if not cls.MODEL_PATH.exists() or os.path.getsize(cls.MODEL_PATH) < 100_000_000: print(f"[RembgEngine] Downloading apparel model to {cls.MODEL_PATH}...") r = requests.get(cls.MODEL_URL, stream=True, timeout=120) with open(cls.MODEL_PATH, "wb") as f: for chunk in r.iter_content(chunk_size=65536): f.write(chunk) print("[RembgEngine] Download complete.") cls._session = ort.InferenceSession(str(cls.MODEL_PATH), providers=['CPUExecutionProvider']) return cls._session except Exception as e: print(f"[RembgEngine] Could not initialize ONNX session: {e}") return None @classmethod def remove_background(cls, image_bytes: bytes) -> tuple[Image.Image, Image.Image]: import numpy as np import cv2 orig_img = Image.open(io.BytesIO(image_bytes)).convert("RGB") w_orig, h_orig = orig_img.size session = cls.get_session() if session is not None: try: # 1. Preprocess for 768x768 cloth segmentation resized = orig_img.resize((768, 768), Image.BILINEAR) img_np = np.array(resized, dtype=np.float32) / 255.0 mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) std = np.array([0.229, 0.224, 0.225], dtype=np.float32) img_np = (img_np - mean) / std img_np = np.transpose(img_np, (2, 0, 1))[np.newaxis, :, :, :] # 2. Neural Network Forward Pass input_name = session.get_inputs()[0].name outputs = session.run(None, {input_name: img_np}) pred = outputs[0][0] # shape (4, 768, 768) if pred.shape[0] >= 4: exp_pred = np.exp(pred - np.max(pred, axis=0, keepdims=True)) softmax_pred = exp_pred / np.sum(exp_pred, axis=0, keepdims=True) cloth_prob = np.sum(softmax_pred[1:], axis=0) mask_raw = (np.clip(cloth_prob * 255, 0, 255)).astype(np.uint8) else: d1 = pred[0] ma, mi = np.max(d1), np.min(d1) d1 = (d1 - mi) / (ma - mi + 1e-8) mask_raw = (d1 * 255).astype(np.uint8) # Resize mask to original dimensions mask_raw = cv2.resize(mask_raw, (w_orig, h_orig), interpolation=cv2.INTER_CUBIC) # 3. Collar Badge & Inner Label Protection Pass # Prevent neural network from falsely cutting out brand labels and inner collar linings _, binary_mask = cv2.threshold(mask_raw, 75, 255, cv2.THRESH_BINARY) contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) if contours: # Find dominant garment silhouette largest_contour = max(contours, key=cv2.contourArea) filled_mask = np.zeros((h_orig, w_orig), dtype=np.uint8) cv2.drawContours(filled_mask, [largest_contour], -1, 255, thickness=cv2.FILLED) # Estimate background color from outer perimeter corners corners = [ orig_img.getpixel((0, 0)), orig_img.getpixel((w_orig - 1, 0)), orig_img.getpixel((0, h_orig - 1)), orig_img.getpixel((w_orig - 1, h_orig - 1)), ] bg_color = np.median(corners, axis=0) img_cv = np.array(orig_img, dtype=np.float32) diff_from_bg = np.linalg.norm(img_cv - bg_color, axis=2) # Keep all non-background pixels inside the garment silhouette (collar badge, neck tag, inner lining) protected_mask = np.where((filled_mask > 0) & (diff_from_bg > 14), 255, 0).astype(np.uint8) # Combine neural mask with protected inner collar mask combined_mask = np.maximum(mask_raw, protected_mask) # Morphological closing to seal fine label stitches and seams kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)) combined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_CLOSE, kernel) # Smooth edges with Gaussian anti-aliasing final_mask = cv2.GaussianBlur(combined_mask, (3, 3), 0) else: final_mask = mask_raw # 4. Composite Cutout Image mask_pil = Image.fromarray(final_mask, mode="L") cutout = Image.new("RGBA", (w_orig, h_orig), (0, 0, 0, 0)) cutout.paste(orig_img, (0, 0), mask=mask_pil) return cutout, mask_pil except Exception as e: print(f"[RembgEngine] ONNX inference error: {e}") # Fallback to CV segmentation return RembgService.cv_segmentation_fallback(orig_img) class RembgService: DEFAULT_PROMPT = "Replace everyone's outfit with the outfit from the reference image" @staticmethod def remove_background(image_bytes: bytes) -> tuple[Image.Image, Image.Image]: return RembgOnnxEngine.remove_background(image_bytes) @staticmethod def cv_segmentation_fallback(input_image: Image.Image) -> tuple[Image.Image, Image.Image]: import numpy as np import cv2 w, h = input_image.size img_cv = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR) corners = [ input_image.getpixel((0, 0)), input_image.getpixel((w - 1, 0)), input_image.getpixel((0, h - 1)), input_image.getpixel((w - 1, h - 1)), input_image.getpixel((w // 2, 0)), input_image.getpixel((0, h // 2)), ] bg_r = np.median([c[0] for c in corners]) bg_g = np.median([c[1] for c in corners]) bg_b = np.median([c[2] for c in corners]) bg_color = np.array([bg_r, bg_g, bg_b], dtype=np.float32) diff = np.linalg.norm(np.array(input_image, dtype=np.float32) - bg_color, axis=2) threshold = max(25.0, float(np.percentile(diff, 20))) mask_np = np.where(diff > threshold, 255, 0).astype(np.uint8) center_x, center_y = w // 2, h // 2 y_indices, x_indices = np.indices((h, w)) dist_from_center = np.sqrt(((x_indices - center_x) / (w / 2)) ** 2 + ((y_indices - center_y) / (h / 2)) ** 2) mask_np = np.where((dist_from_center < 0.94) & (mask_np > 0), 255, 0).astype(np.uint8) mask_img = Image.fromarray(mask_np, mode="L").filter(ImageFilter.GaussianBlur(radius=1.5)) cutout = Image.new("RGBA", (w, h), (0, 0, 0, 0)) cutout.paste(input_image, (0, 0), mask=mask_img) return cutout, mask_img @staticmethod def process_and_save(image_bytes: bytes, filename_prefix: str = "cloth") -> dict: unique_id = uuid.uuid4().hex[:8] orig_filename = f"{filename_prefix}_{unique_id}_orig.png" cutout_filename = f"{filename_prefix}_{unique_id}_cutout.png" mask_filename = f"{filename_prefix}_{unique_id}_mask.png" orig_path = settings.CLOTH_UPLOAD_DIR / orig_filename cutout_path = settings.CUTOUT_UPLOAD_DIR / cutout_filename mask_path = settings.CUTOUT_UPLOAD_DIR / mask_filename orig_img = Image.open(io.BytesIO(image_bytes)).convert("RGBA") orig_img.save(orig_path, format="PNG") cutout_img, mask_img = RembgService.remove_background(image_bytes) cutout_img.save(cutout_path, format="PNG") mask_img.save(mask_path, format="PNG") return { "original_url": f"/static/uploads/clothes/{orig_filename}", "cutout_url": f"/static/uploads/cutouts/{cutout_filename}", "mask_url": f"/static/uploads/cutouts/{mask_filename}", "width": orig_img.width, "height": orig_img.height, } @staticmethod def generate_vton_prompt(cloth_name: str = "", category: str = "", brand: str = "MAP", gender: str = "Unisex") -> str: return RembgService.DEFAULT_PROMPT