Deploy clothing detection API to HF Spaces

Dockerfile

@@ -0,0 +1,40 @@
+FROM python:3.11-slim
+
+# Create user as required by HF
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    HOST=0.0.0.0 \
+    PORT=7860 \
+    WARMUP_ON_STARTUP=true
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    libgl1 \
+    libglib2.0-0 \
+    libsm6 \
+    libxext6 \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+# Copy requirements and install Python dependencies
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r requirements.txt
+
+# Copy app code
+COPY --chown=user . /app
+
+# Create results directory
+RUN mkdir -p results
+
+EXPOSE 7860
+
+# HF requires port 7860
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+
+

app.py

@@ -0,0 +1,168 @@
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request
+from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from typing import Optional, List
+from process import get_dominant_color_from_base64
+from clothing_detector import (
+    detect_clothing_types,
+    create_clothing_only_image,
+    get_clothing_detector,
+)
+import logging
+import os
+import base64
+from starlette import status
+
+# Logging setup
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+app = FastAPI(title="FashionAI API", description="Clothing analysis & segmentation API")
+
+# CORS (configure with env ALLOWED_ORIGINS="http://localhost:5173,https://your-site")
+allowed_origins_env = os.getenv("ALLOWED_ORIGINS", "*")
+allow_origins: List[str]
+if allowed_origins_env.strip() == "*":
+    allow_origins = ["*"]
+else:
+    allow_origins = [o.strip() for o in allowed_origins_env.split(",") if o.strip()]
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=allow_origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# API settings
+MAX_UPLOAD_MB = int(os.getenv("MAX_UPLOAD_MB", "10"))
+MAX_UPLOAD_BYTES = MAX_UPLOAD_MB * 1024 * 1024
+ALLOWED_CONTENT_TYPES = {
+    c.strip() for c in os.getenv("ALLOWED_CONTENT_TYPES", "image/jpeg,image/png,image/webp").split(",") if c.strip()
+}
+
+
+@app.exception_handler(Exception)
+async def unhandled_exception_handler(request: Request, exc: Exception):
+    logging.exception("Unhandled server error: %s", exc)
+    return JSONResponse(
+        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+        content={"error": "Internal Server Error"},
+    )
+
+@app.on_event("startup")
+async def maybe_warmup_model():
+    if os.getenv("WARMUP_ON_STARTUP", "true").lower() in {"1", "true", "yes"}:
+        # Warm up model on startup to reduce first request latency
+        get_clothing_detector()
+
+
+@app.get("/")
+async def api_root():
+    return JSONResponse({
+        "name": "FashionAI API",
+        "status": "ok",
+        "docs": "/docs",
+        "endpoints": ["/clothing", "/analyze", "/analyze/base64", "/labels", "/healthz"],
+    })
+
+
+@app.get("/healthz")
+async def health_check():
+    return {"status": "ok"}
+
+@app.post("/clothing")
+async def get_clothing_list(file: UploadFile = File(...)):
+    """Detect all clothing types on image and return coordinates."""
+    logger.info(f"Processing clothing detection for file: {file.filename}")
+    # Validation
+    if file.content_type not in ALLOWED_CONTENT_TYPES:
+        raise HTTPException(status_code=415, detail=f"Unsupported content-type: {file.content_type}")
+    # Read with size guard
+    image_bytes = await file.read()
+    if len(image_bytes) > MAX_UPLOAD_BYTES:
+        raise HTTPException(status_code=413, detail=f"File too large. Max {MAX_UPLOAD_MB}MB")
+    clothing_result = detect_clothing_types(image_bytes)
+    logger.info(f"Clothing detection completed. Found {clothing_result.get('total_detected', 0)} items")
+    return clothing_result
+
+@app.post("/analyze")
+async def analyze_image(
+    file: UploadFile = File(...),
+    selected_clothing: Optional[str] = Form(None)
+):
+    """
+    Full image analysis: clothing detection, clothing-only image, dominant color.
+
+    - selected_clothing: Optional clothing type to focus on
+    - color: Dominant color of clothing
+    - clothing_analysis: Detected clothing types with stats
+    - clothing_only_image: Base64 PNG with transparent background
+    """
+    logger.info(f"Processing full analysis for file: {file.filename}, selected_clothing: {selected_clothing}")
+    if file.content_type not in ALLOWED_CONTENT_TYPES:
+        raise HTTPException(status_code=415, detail=f"Unsupported content-type: {file.content_type}")
+    image_bytes = await file.read()
+    if len(image_bytes) > MAX_UPLOAD_BYTES:
+        raise HTTPException(status_code=413, detail=f"File too large. Max {MAX_UPLOAD_MB}MB")
+
+    # Step 1: Detect clothing types (cached segmentation)
+    logger.info("Detecting clothing types...")
+    clothing_result = detect_clothing_types(image_bytes)
+
+    # Step 2: Create clothing-only image (cached segmentation)
+    logger.info("Creating clothing-only image...")
+    clothing_only_image = create_clothing_only_image(image_bytes, selected_clothing)
+
+    # Step 3: Get dominant color from clothing-only image (no background)
+    logger.info("Getting dominant color from clothing-only image...")
+    color = get_dominant_color_from_base64(clothing_only_image)
+
+    logger.info("Full analysis completed successfully")
+    return JSONResponse(content={
+        "dominant_color": color,
+        "clothing_analysis": clothing_result,
+        "clothing_only_image": clothing_only_image,
+        "selected_clothing": selected_clothing
+    })
+
+
+class Base64AnalyzeRequest(BaseModel):
+    image_base64: str
+    selected_clothing: Optional[str] = None
+
+
+@app.post("/analyze/base64")
+async def analyze_image_base64(payload: Base64AnalyzeRequest):
+    """Analyze base64-encoded image (handy for React Native)."""
+    # Decode image from base64
+    if payload.image_base64.startswith("data:image"):
+        base64_data = payload.image_base64.split(",", 1)[1]
+    else:
+        base64_data = payload.image_base64
+
+    image_bytes = base64.b64decode(base64_data)
+
+    # 1) Clothing detection
+    clothing_result = detect_clothing_types(image_bytes)
+
+    # 2) Clothing-only image
+    clothing_only_image = create_clothing_only_image(image_bytes, payload.selected_clothing)
+
+    # 3) Dominant color from clothing-only image
+    color = get_dominant_color_from_base64(clothing_only_image)
+
+    return JSONResponse(content={
+        "dominant_color": color,
+        "clothing_analysis": clothing_result,
+        "clothing_only_image": clothing_only_image,
+        "selected_clothing": payload.selected_clothing,
+    })
+
+
+@app.get("/labels")
+async def get_labels():
+    detector = get_clothing_detector()
+    return {"labels": list(detector.labels.values())}

clothing_detector.py

@@ -0,0 +1,331 @@
+import hashlib
+from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation
+from PIL import Image
+import torch
+import torch.nn as nn
+from io import BytesIO
+import numpy as np
+from collections import Counter
+import logging
+import base64
+
+# Logging setup
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Global cache for segmentation results
+_segmentation_cache = {}
+
+class ClothingDetector:
+    def __init__(self):
+        """Initialize clothing segmentation model."""
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        logger.info(f"Using device: {self.device}")
+        
+        # Load processor and model
+        self.processor = SegformerImageProcessor.from_pretrained("mattmdjaga/segformer_b2_clothes")
+        self.model = AutoModelForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes")
+        self.model.to(self.device)
+        self.model.eval()
+        
+        # Clothing labels mapping
+        self.labels = {
+            0: "Background",
+            1: "Hat", 
+            2: "Hair",
+            3: "Sunglasses",
+            4: "Upper-clothes",
+            5: "Skirt",
+            6: "Pants",
+            7: "Dress",
+            8: "Belt",
+            9: "Left-shoe",
+            10: "Right-shoe",
+            11: "Face",
+            12: "Left-leg",
+            13: "Right-leg",
+            14: "Left-arm",
+            15: "Right-arm",
+            16: "Bag",
+            17: "Scarf"
+        }
+        
+        # Clothing classes (exclude body parts and background)
+        self.clothing_classes = [4, 5, 6, 7, 8, 9, 10, 16, 17]  # Upper-clothes, Skirt, Pants, Dress, Belt, Left-shoe, Right-shoe, Bag, Scarf
+        
+        logger.info("Clothing detector initialized successfully")
+    
+    def _get_image_hash(self, image_bytes: bytes) -> str:
+        """Create image hash to use as cache key."""
+        return hashlib.md5(image_bytes).hexdigest()
+    
+    def _segment_image(self, image_bytes: bytes):
+        """Run image segmentation with caching."""
+        image_hash = self._get_image_hash(image_bytes)
+        
+        # Check cache
+        if image_hash in _segmentation_cache:
+            logger.info("Using cached segmentation result")
+            return _segmentation_cache[image_hash]
+        
+        # Run segmentation
+        logger.info("Performing new segmentation")
+        image = Image.open(BytesIO(image_bytes)).convert("RGB")
+        
+        # Prepare inputs
+        inputs = self.processor(images=image, return_tensors="pt")
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+        
+        # Forward pass
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            logits = outputs.logits.cpu()
+        
+        # Upsample logits to original image size
+        upsampled_logits = nn.functional.interpolate(
+            logits,
+            size=image.size[::-1],  # (height, width)
+            mode="bilinear",
+            align_corners=False,
+        )
+        
+        # Get predicted mask
+        pred_seg = upsampled_logits.argmax(dim=1)[0].numpy()
+        
+        # Save to cache
+        result = {
+            'pred_seg': pred_seg,
+            'image': image,
+            'image_size': image.size
+        }
+        _segmentation_cache[image_hash] = result
+        
+        # Limit cache size (keep last 10)
+        if len(_segmentation_cache) > 10:
+            oldest_key = next(iter(_segmentation_cache))
+            del _segmentation_cache[oldest_key]
+        
+        return result
+    
+    def detect_clothing(self, image_bytes: bytes) -> dict:
+        """
+        Detect clothing types on image and return coordinates.
+
+        Args:
+            image_bytes: Raw image bytes
+
+        Returns:
+            dict: Clothing types with pixel stats and bounding boxes
+        """
+        try:
+            # Get cached segmentation result
+            seg_result = self._segment_image(image_bytes)
+            pred_seg = seg_result['pred_seg']
+            image = seg_result['image']
+            
+            # Count pixels per class and compute bounding boxes
+            clothing_types = {}
+            coordinates = {}
+            total_pixels = pred_seg.size
+            
+            for class_id, label_name in self.labels.items():
+                if label_name not in ["Background", "Face", "Hair", "Left-arm", "Right-arm", "Left-leg", "Right-leg"]:
+                    # Create mask for this class
+                    mask = (pred_seg == class_id)
+                    
+                    if np.any(mask):
+                        # Count pixels
+                        count = np.sum(mask)
+                        percentage = (count / total_pixels) * 100
+                        
+                        clothing_types[label_name] = {
+                            "pixels": int(count),
+                            "percentage": round(percentage, 2)
+                        }
+                        
+                        # Compute bounding box
+                        rows = np.any(mask, axis=1)
+                        cols = np.any(mask, axis=0)
+                        
+                        if np.any(rows) and np.any(cols):
+                            y_min, y_max = np.where(rows)[0][[0, -1]]
+                            x_min, x_max = np.where(cols)[0][[0, -1]]
+                            
+                            # Add padding (10% of clothing size)
+                            clothing_width = x_max - x_min
+                            clothing_height = y_max - y_min
+                            padding_x = int(clothing_width * 0.1)
+                            padding_y = int(clothing_height * 0.1)
+                            
+                            # Apply padding with image bounds
+                            x_min = max(0, x_min - padding_x)
+                            y_min = max(0, y_min - padding_y)
+                            x_max = min(image.width, x_max + padding_x)
+                            y_max = min(image.height, y_max + padding_y)
+                            
+                            coordinates[label_name] = {
+                                "x_min": int(x_min),
+                                "y_min": int(y_min),
+                                "x_max": int(x_max),
+                                "y_max": int(y_max),
+                                "width": int(x_max - x_min),
+                                "height": int(y_max - y_min)
+                            }
+            
+            # Sort by percentage area
+            sorted_clothing = dict(sorted(
+                clothing_types.items(), 
+                key=lambda x: x[1]["percentage"], 
+                reverse=True
+            ))
+            
+            return {
+                "clothing_types": sorted_clothing,
+                "coordinates": coordinates,
+                "total_detected": len(sorted_clothing),
+                "main_clothing": list(sorted_clothing.keys())[:3] if sorted_clothing else []
+            }
+            
+        except Exception as e:
+            logger.error(f"Error in clothing detection: {str(e)}")
+            return {
+                "clothing_types": {},
+                "coordinates": {},
+                "total_detected": 0,
+                "main_clothing": [],
+                "error": str(e)
+            }
+    
+    def create_clothing_only_image(self, image_bytes: bytes, selected_clothing: str = None) -> str:
+        """
+        Create clothing-only image with transparent background.
+
+        Args:
+            image_bytes: Raw image bytes
+            selected_clothing: Optional clothing label to isolate
+
+        Returns:
+            str: Base64-encoded PNG data URL
+        """
+        try:
+            # Get cached segmentation
+            seg_result = self._segment_image(image_bytes)
+            pred_seg = seg_result['pred_seg']
+            image = seg_result['image']
+            
+            # Create clothing-only mask
+            clothing_mask = np.zeros_like(pred_seg, dtype=bool)
+            
+            if selected_clothing:
+                # If specific clothing selected, find its class id
+                selected_class_id = None
+                for class_id, label_name in self.labels.items():
+                    if label_name == selected_clothing:
+                        selected_class_id = class_id
+                        break
+                
+                if selected_class_id is not None:
+                    # Build mask only for the selected class
+                    clothing_mask = (pred_seg == selected_class_id)
+                else:
+                    # If not found, fall back to all clothing classes
+                    for class_id in self.clothing_classes:
+                        clothing_mask |= (pred_seg == class_id)
+            else:
+                # Otherwise, use all clothing classes
+                for class_id in self.clothing_classes:
+                    clothing_mask |= (pred_seg == class_id)
+            
+            # Convert image to numpy array
+            image_array = np.array(image)
+            
+            # Compose RGBA with transparent background
+            clothing_only_rgba = np.zeros((image_array.shape[0], image_array.shape[1], 4), dtype=np.uint8)
+            clothing_only_rgba[..., :3] = image_array  # RGB channels
+            clothing_only_rgba[..., 3] = 255  # Alpha channel (opaque)
+            clothing_only_rgba[~clothing_mask, 3] = 0  # Transparent for non-clothing
+            
+            # Create PIL image
+            clothing_image = Image.fromarray(clothing_only_rgba, 'RGBA')
+            
+            # If a specific clothing selected, crop with padding
+            if selected_clothing and selected_class_id is not None:
+                clothing_image = self._crop_with_padding(clothing_image, clothing_mask)
+            
+            # Encode to base64
+            buffer = BytesIO()
+            clothing_image.save(buffer, format='PNG')
+            img_str = base64.b64encode(buffer.getvalue()).decode()
+            
+            return f"data:image/png;base64,{img_str}"
+            
+        except Exception as e:
+            logger.error(f"Error in creating clothing-only image: {str(e)}")
+            return ""
+    
+    def _crop_with_padding(self, image: Image.Image, mask: np.ndarray, padding_percent: float = 0.1) -> Image.Image:
+        """
+        Crop image around clothing mask with padding.
+
+        Args:
+            image: PIL image
+            mask: Clothing mask
+            padding_percent: Padding percentage relative to clothing size
+
+        Returns:
+            Image.Image: Cropped image
+        """
+        try:
+            # Find clothing bounds
+            rows = np.any(mask, axis=1)
+            cols = np.any(mask, axis=0)
+            
+            if not np.any(rows) or not np.any(cols):
+                return image  # If no clothing found, return original
+            
+            # Get bounds
+            y_min, y_max = np.where(rows)[0][[0, -1]]
+            x_min, x_max = np.where(cols)[0][[0, -1]]
+            
+            # Compute clothing size
+            clothing_width = x_max - x_min
+            clothing_height = y_max - y_min
+            
+            # Compute padding
+            padding_x = int(clothing_width * padding_percent)
+            padding_y = int(clothing_height * padding_percent)
+            
+            # Apply padding within image bounds
+            x_min = max(0, x_min - padding_x)
+            y_min = max(0, y_min - padding_y)
+            x_max = min(image.width, x_max + padding_x)
+            y_max = min(image.height, y_max + padding_y)
+            
+            # Crop
+            cropped_image = image.crop((x_min, y_min, x_max, y_max))
+            
+            return cropped_image
+            
+        except Exception as e:
+            logger.error(f"Error in cropping with padding: {str(e)}")
+            return image
+
+# Global detector singleton (to reuse model)
+_detector = None
+
+def get_clothing_detector():
+    """Get global detector instance (lazy-init)."""
+    global _detector
+    if _detector is None:
+        _detector = ClothingDetector()
+    return _detector
+
+def detect_clothing_types(image_bytes: bytes) -> dict:
+    """Convenience wrapper for clothing detection."""
+    detector = get_clothing_detector()
+    return detector.detect_clothing(image_bytes)
+
+def create_clothing_only_image(image_bytes: bytes, selected_clothing: str = None) -> str:
+    """Convenience wrapper for clothing-only image creation."""
+    detector = get_clothing_detector()
+    return detector.create_clothing_only_image(image_bytes, selected_clothing) 

process.py

@@ -0,0 +1,79 @@
+from rembg import remove
+from PIL import Image
+from io import BytesIO
+from sklearn.cluster import KMeans
+import base64
+
+import os
+import uuid
+import numpy as np
+
+
+def get_dominant_color(processed_bytes, k=3):
+    # Step 1: load transparent image
+    image = Image.open(BytesIO(processed_bytes)).convert("RGBA")
+    image = image.resize((100, 100))  # Resize to speed up
+
+    # Step 2: Filter only visible (non-transparent) pixels
+    np_image = np.array(image)
+    rgb_pixels = np_image[...,:3]    # Ignore alpha channel
+    alpha = np_image[..., 3]
+    rgb_pixels = rgb_pixels[alpha > 0]  # Keep only pixels where alpha > 0
+
+    # Step 3: KMeans clustering
+    kmeans = KMeans(n_clusters=k, n_init='auto')
+    kmeans.fit(rgb_pixels)
+    dominant_color = kmeans.cluster_centers_[0]
+    r, g, b = map(int, dominant_color)
+    return f"rgb({r}, {g}, {b})"
+
+
+def get_dominant_color_from_base64(base64_image, k=3):
+    """Compute dominant color from base64-encoded clothing-only image."""
+    try:
+        # Step 1: Decode base64 to bytes
+        if base64_image.startswith('data:image'):
+            # Remove data URL prefix
+            base64_data = base64_image.split(',')[1]
+        else:
+            base64_data = base64_image
+            
+        image_bytes = base64.b64decode(base64_data)
+        
+        # Step 2: Load image and convert to RGBA
+        image = Image.open(BytesIO(image_bytes)).convert("RGBA")
+        image = image.resize((100, 100))  # Resize to speed up
+
+        # Step 3: Filter only visible (non-transparent) pixels
+        np_image = np.array(image)
+        rgb_pixels = np_image[...,:3]    # Ignore alpha channel
+        alpha = np_image[..., 3]
+        rgb_pixels = rgb_pixels[alpha > 0]  # Keep only pixels where alpha > 0
+
+        # Check if we have any visible pixels
+        if len(rgb_pixels) == 0:
+            return "rgb(0, 0, 0)"  # Fallback to black if no visible pixels
+
+        # Step 4: KMeans clustering
+        kmeans = KMeans(n_clusters=k, n_init='auto')
+        kmeans.fit(rgb_pixels)
+        dominant_color = kmeans.cluster_centers_[0]
+        r, g, b = map(int, dominant_color)
+        return f"rgb({r}, {g}, {b})"
+        
+    except Exception as e:
+        print(f"Error in get_dominant_color_from_base64: {e}")
+        return "rgb(0, 0, 0)"  # Fallback to black on error
+
+
+def remove_background(image_bytes: bytes) -> bytes:
+    result_bytes = remove(image_bytes)
+
+    # Save image to disk
+    output_image = Image.open(BytesIO(result_bytes))
+    file_name = f"{uuid.uuid4().hex[:8]}.png"
+    output_path = os.path.join("results", file_name)
+    output_image.save(output_path)
+
+    print(f"✅ Saved background-removed image to: {output_path}")
+    return result_bytes

requirements.txt

@@ -0,0 +1,9 @@
+fastapi
+uvicorn[standard]
+pillow
+numpy
+transformers
+torch
+torchvision
+scikit-learn
+python-multipart