src/app.py

"""
SAM3 Static Image Segmentation - Correct Implementation

Uses Sam3Model (not Sam3VideoModel) for text-prompted static image segmentation.
"""
import base64
import io
import asyncio
import torch
import numpy as np
from PIL import Image
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoProcessor, AutoModel
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Load SAM3 model for STATIC IMAGES
processor = AutoProcessor.from_pretrained("./model", trust_remote_code=True)
model = AutoModel.from_pretrained(
    "./model",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    trust_remote_code=True
)

model.eval()
if torch.cuda.is_available():
    model.cuda()
    logger.info(f"GPU: {torch.cuda.get_device_name()}")
    logger.info(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")

logger.info(f"✓ Loaded {model.__class__.__name__} for static image segmentation")

# Simple concurrency control
class VRAMManager:
    def __init__(self):
        self.semaphore = asyncio.Semaphore(2)
        self.processing_count = 0

    def get_vram_status(self):
        if not torch.cuda.is_available():
            return {}
        return {
            "total_gb": torch.cuda.get_device_properties(0).total_memory / 1e9,
            "allocated_gb": torch.cuda.memory_allocated() / 1e9,
            "free_gb": (torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_reserved()) / 1e9,
            "processing_now": self.processing_count
        }

    async def acquire(self, rid):
        await self.semaphore.acquire()
        self.processing_count += 1

    def release(self, rid):
        self.processing_count -= 1
        self.semaphore.release()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

vram_manager = VRAMManager()
app = FastAPI(title="SAM3 Static Image API")

class Request(BaseModel):
    inputs: str
    parameters: dict


def run_inference(image_b64: str, classes: list, request_id: str):
    """
    Sam3Model inference for static images with text prompts.

    Uses official SAM3 processor post-processing for correct mask generation.
    """
    try:
        # Decode image
        image_bytes = base64.b64decode(image_b64)
        pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        logger.info(f"[{request_id}] Image: {pil_image.size}, Classes: {classes}")

        # Process with Sam3Processor
        # Sam3Model expects: batch of images matching text prompts
        # For multiple objects in ONE image, repeat the image for each class
        images_batch = [pil_image] * len(classes)
        inputs = processor(
            images=images_batch,  # Repeat image for each text prompt
            text=classes,  # List of text prompts
            return_tensors="pt"
        )

        # Store original sizes for post-processing
        # Format: [[height, width]] for EACH image in batch
        # Since we repeat the image for each class, repeat the size too
        original_size = [pil_image.size[1], pil_image.size[0]]  # [height, width]
        original_sizes = torch.tensor([original_size] * len(classes))
        inputs["original_sizes"] = original_sizes

        logger.info(f"[{request_id}] Processing {len(classes)} classes with batched images")
        logger.info(f"[{request_id}] Original size: {pil_image.size} (W x H)")

        # Move to GPU and match model dtype
        if torch.cuda.is_available():
            model_dtype = next(model.parameters()).dtype
            inputs = {
                k: v.cuda().to(model_dtype) if isinstance(v, torch.Tensor) and v.dtype.is_floating_point else v.cuda() if isinstance(v, torch.Tensor) else v
                for k, v in inputs.items()
            }
            logger.info(f"[{request_id}] Moved inputs to GPU (float tensors to {model_dtype})")

        # Sam3Model Inference
        with torch.no_grad():
            outputs = model(**inputs)
            logger.info(f"[{request_id}] Forward pass successful!")

        logger.info(f"[{request_id}] Output type: {type(outputs)}")

        # Use processor's official post-processing method
        # This handles:
        # - Logit to probability conversion (sigmoid)
        # - Proper thresholding (default 0.5)
        # - Resizing to original image dimensions
        # - Score extraction
        logger.info(f"[{request_id}] Using processor.post_process_instance_segmentation()")

        try:
            processed = processor.post_process_instance_segmentation(
                outputs,
                threshold=0.3,          # Score threshold for detections (lowered to detect road cracks)
                mask_threshold=0.5,     # Probability threshold for mask pixels
                target_sizes=original_sizes.tolist()
            )
            # Returns a LIST of results, one per image in batch (one per class in our case)

            logger.info(f"[{request_id}] Post-processing successful!")
            logger.info(f"[{request_id}] Number of batched results: {len(processed)}")

        except Exception as proc_error:
            logger.error(f"[{request_id}] Post-processing failed: {proc_error}")
            logger.info(f"[{request_id}] Falling back to manual processing")

            # Fallback to manual processing with sigmoid fix
            results = []

            # Extract masks from outputs
            if hasattr(outputs, 'pred_masks'):
                pred_masks = outputs.pred_masks
            elif hasattr(outputs, 'masks'):
                pred_masks = outputs.masks
            elif isinstance(outputs, dict) and 'pred_masks' in outputs:
                pred_masks = outputs['pred_masks']
            else:
                raise ValueError("Cannot find masks in model output")

            logger.info(f"[{request_id}] pred_masks shape: {pred_masks.shape}")

            for i, cls in enumerate(classes):
                if i < pred_masks.shape[1]:
                    mask_tensor = pred_masks[0, i]

                    # Resize to original size
                    if mask_tensor.shape[-2:] != pil_image.size[::-1]:
                        mask_tensor = torch.nn.functional.interpolate(
                            mask_tensor.unsqueeze(0).unsqueeze(0),
                            size=pil_image.size[::-1],
                            mode='bilinear',
                            align_corners=False
                        ).squeeze()

                    # CRITICAL FIX: Convert logits to probabilities THEN threshold
                    probs = torch.sigmoid(mask_tensor)
                    binary_mask = (probs > 0.5).float().cpu().numpy().astype("uint8") * 255
                else:
                    binary_mask = np.zeros(pil_image.size[::-1], dtype="uint8")

                # Convert to PNG
                pil_mask = Image.fromarray(binary_mask, mode="L")
                buf = io.BytesIO()
                pil_mask.save(buf, format="PNG")
                mask_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")

                # Extract score
                score = 1.0
                if hasattr(outputs, 'pred_logits') and i < outputs.pred_logits.shape[1]:
                    # Convert logits to probability
                    score = float(torch.sigmoid(outputs.pred_logits[0, i]).cpu())

                results.append({
                    "label": cls,
                    "mask": mask_b64,
                    "score": score
                })

            logger.info(f"[{request_id}] Completed (fallback): {len(results)} masks generated")
            return results

        # Extract results from processor output
        # CRITICAL: processor returns one result dict per class (batched)
        # Each result dict contains MULTIPLE instances of that class
        results = []

        total_instances = 0
        for i, cls in enumerate(classes):
            class_result = processed[i]  # Results for this specific class

            num_instances = len(class_result['masks']) if 'masks' in class_result else 0
            total_instances += num_instances

            if num_instances > 0:
                logger.info(f"[{request_id}] {cls}: {num_instances} instance(s) detected")

                # Loop through ALL instances of this class
                for j in range(num_instances):
                    # Get mask (already binary, resized to original size)
                    mask_np = class_result['masks'][j].cpu().numpy().astype("uint8") * 255

                    # Convert to PNG
                    pil_mask = Image.fromarray(mask_np, mode="L")
                    buf = io.BytesIO()
                    pil_mask.save(buf, format="PNG")
                    mask_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")

                    # Get score (already converted to probability by processor)
                    score = float(class_result['scores'][j]) if 'scores' in class_result else 1.0

                    # Calculate coverage for logging
                    coverage = (mask_np > 0).sum() / mask_np.size * 100

                    results.append({
                        "label": cls,
                        "mask": mask_b64,
                        "score": score,
                        "instance_id": j
                    })

                    logger.info(f"[{request_id}]   └─ Instance {j}: score={score:.3f}, coverage={coverage:.2f}%")
            else:
                logger.info(f"[{request_id}] {cls}: No instances detected")

        logger.info(f"[{request_id}] Completed: {total_instances} instance(s) across {len(classes)} class(es)")
        return results

    except Exception as e:
        logger.error(f"[{request_id}] Failed: {str(e)}")
        import traceback
        traceback.print_exc()
        raise


@app.post("/")
async def predict(req: Request):
    request_id = str(id(req))[:8]
    try:
        await vram_manager.acquire(request_id)
        try:
            results = await asyncio.to_thread(
                run_inference,
                req.inputs,
                req.parameters.get("classes", []),
                request_id
            )
            return results
        finally:
            vram_manager.release(request_id)
    except Exception as e:
        logger.error(f"[{request_id}] Error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/health")
async def health():
    return {
        "status": "healthy",
        "model": model.__class__.__name__,
        "gpu_available": torch.cuda.is_available(),
        "vram": vram_manager.get_vram_status()
    }


@app.get("/metrics")
async def metrics():
    return vram_manager.get_vram_status()


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860, workers=1)