sam_server.py

"""
SAM (Segment Anything Model) FastAPI Backend Server
Versione con segmentazione automatica e unione maschere

Run with: uvicorn sam_server:app --host 0.0.0.0 --port 8000 --reload
"""

import os
import io
import base64
import numpy as np
from typing import List, Optional, Dict, Any
from pathlib import Path
import json
import cv2

from fastapi import FastAPI, File, UploadFile, HTTPException, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from PIL import Image

# SAM imports
import torch
from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator

app = FastAPI(
    title="SAM Segmentation API",
    description="API for automatic image segmentation using Segment Anything Model",
    version="2.0.0"
)

# Enable CORS for Next.js frontend
app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "http://localhost:3000",
        "http://127.0.0.1:3000",
        "https://data-visualization-red-six.vercel.app",
        "https://benny2199-sam-server.hf.space",
        "*",
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global model variables
sam_model = None
sam_predictor = None
mask_generator = None
current_model_type = None

# Cache for current image masks
current_image_masks: List[Dict[str, Any]] = []
current_image_shape = None

# Model checkpoint paths
MODEL_CHECKPOINTS = {
    "vit_b": "sam_vit_b_01ec64.pth",
    "vit_l": "sam_vit_l_0b3195.pth",
    "vit_h": "sam_vit_h_4b8939.pth",
}

CHECKPOINTS_DIR = Path(__file__).parent / "checkpoints"

# Color palette for masks
COLOR_PALETTE = [
    [255, 0, 0], [0, 255, 0], [0, 0, 255], [255, 255, 0], [255, 0, 255],
    [0, 255, 255], [255, 128, 0], [128, 0, 255], [0, 255, 128], [255, 0, 128],
    [128, 255, 0], [0, 128, 255], [255, 128, 128], [128, 255, 128], [128, 128, 255],
    [255, 255, 128], [255, 128, 255], [128, 255, 255], [192, 64, 64], [64, 192, 64],
    [64, 64, 192], [192, 192, 64], [192, 64, 192], [64, 192, 192], [160, 32, 240],
    [32, 160, 240], [240, 32, 160], [240, 160, 32], [32, 240, 160], [160, 240, 32],
]


class MaskResult(BaseModel):
    mask_id: int
    mask_base64: str
    score: float
    area: int
    coverage_percent: float
    color: List[int]
    bbox: List[int]  # x, y, width, height


class SegmentationResponse(BaseModel):
    success: bool
    masks: List[MaskResult]
    total_masks: int
    image_width: int
    image_height: int
    message: str = ""


class CombineMasksRequest(BaseModel):
    mask_indices: List[int]
    operation: str = "union"  # union, intersection, difference


class CombinedMaskResponse(BaseModel):
    success: bool
    mask_base64: str
    area: int
    coverage_percent: float
    source_indices: List[int]
    operation: str
    message: str = ""


def get_color(index: int) -> List[int]:
    """Get color from palette"""
    return COLOR_PALETTE[index % len(COLOR_PALETTE)]


def mask_to_base64(mask: np.ndarray, color: List[int], alpha: int = 150) -> str:
    """Convert binary mask to colored RGBA base64 image"""
    h, w = mask.shape
    colored_mask = np.zeros((h, w, 4), dtype=np.uint8)
    colored_mask[mask] = [color[0], color[1], color[2], alpha]
    
    img = Image.fromarray(colored_mask, mode='RGBA')
    buffer = io.BytesIO()
    img.save(buffer, format='PNG')
    return base64.b64encode(buffer.getvalue()).decode('utf-8')


def binary_mask_to_base64(mask: np.ndarray) -> str:
    """Convert binary mask to grayscale base64 image"""
    mask_uint8 = (mask * 255).astype(np.uint8)
    img = Image.fromarray(mask_uint8, mode='L')
    buffer = io.BytesIO()
    img.save(buffer, format='PNG')
    return base64.b64encode(buffer.getvalue()).decode('utf-8')


@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "model_loaded": sam_model is not None,
        "current_model": current_model_type,
        "device": str(torch.device("cuda" if torch.cuda.is_available() else "cpu")),
        "cached_masks": len(current_image_masks)
    }


@app.get("/models")
async def list_models():
    """List available SAM models"""
    available = []
    for model_type, checkpoint in MODEL_CHECKPOINTS.items():
        checkpoint_path = CHECKPOINTS_DIR / checkpoint
        available.append({
            "model_type": model_type,
            "checkpoint": checkpoint,
            "available": checkpoint_path.exists(),
            "path": str(checkpoint_path)
        })
    return {"models": available}


@app.post("/load-model/{model_type}")
async def load_model(model_type: str):
    """Load a SAM model"""
    global sam_model, sam_predictor, mask_generator, current_model_type
    
    if model_type not in MODEL_CHECKPOINTS:
        raise HTTPException(
            status_code=400,
            detail=f"Invalid model type. Choose from: {list(MODEL_CHECKPOINTS.keys())}"
        )
    
    checkpoint_path = CHECKPOINTS_DIR / MODEL_CHECKPOINTS[model_type]
    
    if not checkpoint_path.exists():
        raise HTTPException(
            status_code=404,
            detail=f"Checkpoint not found at {checkpoint_path}. Download from https://github.com/facebookresearch/segment-anything#model-checkpoints"
        )
    
    try:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"Loading SAM model {model_type} on {device}...")
        
        sam_model = sam_model_registry[model_type](checkpoint=str(checkpoint_path))
        sam_model.to(device=device)
        
        sam_predictor = SamPredictor(sam_model)
        
        # Configure automatic mask generator with optimal settings
        mask_generator = SamAutomaticMaskGenerator(
            model=sam_model,
            points_per_side=32,
            pred_iou_thresh=0.88,
            stability_score_thresh=0.92,
            crop_n_layers=1 if model_type == 'vit_h' else 0,
            min_mask_region_area=100,
            box_nms_thresh=0.7
        )
        
        current_model_type = model_type
        
        return {
            "success": True,
            "message": f"Model {model_type} loaded successfully on {device}",
            "model_type": model_type,
            "device": str(device)
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error loading model: {str(e)}")


@app.post("/segment-auto", response_model=SegmentationResponse)
async def segment_automatic(image: UploadFile = File(...)):
    """
    Automatically segment all objects in an image.
    This is the main endpoint - no point prompts needed.
    """
    global current_image_masks, current_image_shape, mask_generator
    
    if mask_generator is None:
        raise HTTPException(status_code=400, detail="No model loaded. Call /load-model first.")
    
    try:
        # Read and convert image
        contents = await image.read()
        img = Image.open(io.BytesIO(contents)).convert("RGB")
        img_array = np.array(img)
        
        current_image_shape = img_array.shape
        h, w = img_array.shape[:2]
        total_area = h * w
        
        print(f"Segmenting image {w}x{h}...")
        
        # Generate all masks automatically
        masks = mask_generator.generate(img_array)
        
        # Sort by area (largest first)
        masks = sorted(masks, key=lambda x: x['area'], reverse=True)
        
        # Store masks for later operations
        current_image_masks = masks
        
        print(f"Found {len(masks)} masks")
        
        # Convert to response format
        results = []
        for idx, mask_data in enumerate(masks):
            mask = mask_data["segmentation"]
            area = mask_data["area"]
            coverage = (area / total_area) * 100
            bbox = mask_data["bbox"]  # x, y, w, h
            score = mask_data.get("stability_score", mask_data.get("predicted_iou", 0.9))
            
            color = get_color(idx)
            mask_b64 = mask_to_base64(mask, color)
            
            results.append(MaskResult(
                mask_id=idx,
                mask_base64=mask_b64,
                score=float(score),
                area=int(area),
                coverage_percent=float(coverage),
                color=color,
                bbox=[int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])]
            ))
        
        return SegmentationResponse(
            success=True,
            masks=results,
            total_masks=len(results),
            image_width=w,
            image_height=h,
            message=f"Found {len(results)} objects automatically"
        )
        
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Error during segmentation: {str(e)}")


@app.post("/combine-masks", response_model=CombinedMaskResponse)
async def combine_masks(request: CombineMasksRequest):
    """
    Combine multiple masks using union, intersection, or difference operations.
    NOTE: Does NOT update the cache - frontend manages mask list separately.
    """
    global current_image_masks, current_image_shape
    
    if not current_image_masks:
        raise HTTPException(status_code=400, detail="No masks available. Run /segment-auto first.")
    
    mask_indices = request.mask_indices
    operation = request.operation.lower()
    
    if not mask_indices:
        raise HTTPException(status_code=400, detail="No mask indices provided")
    
    # Validate indices
    valid_indices = [i for i in mask_indices if 0 <= i < len(current_image_masks)]
    
    if not valid_indices:
        raise HTTPException(status_code=400, detail=f"No valid mask indices. Available: 0-{len(current_image_masks)-1}, requested: {mask_indices}")
    
    if operation not in ["union", "intersection", "difference"]:
        raise HTTPException(status_code=400, detail="Operation must be 'union', 'intersection', or 'difference'")
    
    try:
        h, w = current_image_shape[:2]
        
        print(f"Combining masks at indices: {valid_indices} using {operation}")
        
        if operation == "union":
            # Start with empty mask and OR all together
            combined = np.zeros((h, w), dtype=bool)
            for idx in valid_indices:
                mask_data = current_image_masks[idx]["segmentation"]
                combined = np.logical_or(combined, mask_data)
                print(f"  Added mask {idx}, area: {current_image_masks[idx]['area']}")
        elif operation == "intersection":
            combined = current_image_masks[valid_indices[0]]["segmentation"].copy().astype(bool)
            for idx in valid_indices[1:]:
                combined = np.logical_and(combined, current_image_masks[idx]["segmentation"])
        else:  # difference
            combined = current_image_masks[valid_indices[0]]["segmentation"].copy().astype(bool)
            for idx in valid_indices[1:]:
                combined = np.logical_and(combined, np.logical_not(current_image_masks[idx]["segmentation"]))
        
        # Calculate stats
        area = int(np.sum(combined))
        total_area = h * w
        coverage = (area / total_area) * 100
        
        # Convert to base64
        color = [255, 200, 0]  # Orange/Yellow for combined mask
        mask_b64 = mask_to_base64(combined, color, alpha=180)
        
        print(f"Combined {len(valid_indices)} masks. Result area: {area}")
        
        return CombinedMaskResponse(
            success=True,
            mask_base64=mask_b64,
            area=area,
            coverage_percent=float(coverage),
            source_indices=valid_indices,
            operation=operation,
            message=f"Combined {len(valid_indices)} masks using {operation}"
        )
        
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Error combining masks: {str(e)}")


@app.post("/smooth-mask")
async def smooth_mask(
    mask_indices: str = Form(...),
    kernel_size: int = Form(5),
    iterations: int = Form(1)
):
    """
    Apply morphological smoothing to combined masks.
    """
    global current_image_masks
    
    if not current_image_masks:
        raise HTTPException(status_code=400, detail="No masks available")
    
    try:
        indices = json.loads(mask_indices)
        
        # Combine masks first
        if len(indices) == 1:
            combined = current_image_masks[indices[0]]["segmentation"].copy()
        else:
            combined = current_image_masks[indices[0]]["segmentation"].copy()
            for idx in indices[1:]:
                combined = np.logical_or(combined, current_image_masks[idx]["segmentation"])
        
        # Apply morphological operations
        mask_uint8 = combined.astype(np.uint8)
        kernel = np.ones((kernel_size, kernel_size), np.uint8)
        
        # Close (fill holes)
        smoothed = cv2.morphologyEx(mask_uint8, cv2.MORPH_CLOSE, kernel, iterations=iterations)
        # Open (remove noise)
        smoothed = cv2.morphologyEx(smoothed, cv2.MORPH_OPEN, kernel, iterations=iterations)
        
        smoothed_bool = smoothed.astype(bool)
        
        # Stats
        area = int(np.sum(smoothed_bool))
        total_area = current_image_shape[0] * current_image_shape[1]
        coverage = (area / total_area) * 100
        
        # Convert to base64
        color = [0, 255, 128]  # Green for smoothed
        mask_b64 = mask_to_base64(smoothed_bool, color, alpha=180)
        
        return {
            "success": True,
            "mask_base64": mask_b64,
            "area": area,
            "coverage_percent": coverage,
            "kernel_size": kernel_size,
            "iterations": iterations
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error smoothing mask: {str(e)}")


@app.post("/remove-object")
async def remove_object(
    image: UploadFile = File(...),
    mask_indices: str = Form(...),
    operation: str = Form("union")
):
    """
    Remove selected objects from image (make them transparent).
    Returns the image with objects removed.
    """
    global current_image_masks
    
    if not current_image_masks:
        raise HTTPException(status_code=400, detail="No masks available")
    
    try:
        indices = json.loads(mask_indices)
        
        # Read image
        contents = await image.read()
        img = Image.open(io.BytesIO(contents)).convert("RGBA")
        img_array = np.array(img)
        
        # Combine selected masks
        combined = current_image_masks[indices[0]]["segmentation"].copy()
        for idx in indices[1:]:
            if operation == "union":
                combined = np.logical_or(combined, current_image_masks[idx]["segmentation"])
            else:
                combined = np.logical_and(combined, current_image_masks[idx]["segmentation"])
        
        # Make selected areas transparent
        img_array[combined, 3] = 0  # Set alpha to 0
        
        # Convert to base64
        result_img = Image.fromarray(img_array, mode='RGBA')
        buffer = io.BytesIO()
        result_img.save(buffer, format='PNG')
        result_b64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
        
        return {
            "success": True,
            "image_base64": result_b64,
            "removed_indices": indices,
            "message": f"Removed {len(indices)} objects"
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error removing objects: {str(e)}")


@app.get("/cached-masks")
async def get_cached_masks():
    """Get info about currently cached masks"""
    if not current_image_masks:
        return {"masks": [], "count": 0}
    
    masks_info = []
    for idx, mask_data in enumerate(current_image_masks):
        masks_info.append({
            "mask_id": idx,
            "area": mask_data["area"],
            "bbox": mask_data["bbox"],
            "score": mask_data.get("stability_score", mask_data.get("predicted_iou", 0))
        })
    
    return {
        "masks": masks_info,
        "count": len(masks_info),
        "image_shape": list(current_image_shape) if current_image_shape else None
    }


if __name__ == "__main__":
    import uvicorn

    # Create checkpoints directory
    CHECKPOINTS_DIR.mkdir(exist_ok=True)

    print(f"Checkpoints directory: {CHECKPOINTS_DIR}")
    print("Download SAM model checkpoints from:")
    print("https://github.com/facebookresearch/segment-anything#model-checkpoints")
    print()
    print("Available checkpoints:")
    for model_type, checkpoint in MODEL_CHECKPOINTS.items():
        path = CHECKPOINTS_DIR / checkpoint
        status = "✓" if path.exists() else "✗"
        print(f"  {status} {model_type}: {checkpoint}")
    print()

    port = int(os.environ.get("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)