Spaces:

Eminakkoc
/

Scenedrobe

Sleeping

Emin Akkoç

Replace CLIP with a more visual oriented model

c2005ed 7 months ago

21.1 kB

"""
Scenedrobe API – FastAPI server for Grounded SAM 2.
Models are loaded once at startup and kept in memory.
"""

from __future__ import annotations

import base64
import io
import os
import time
import uuid
from contextlib import asynccontextmanager
from typing import Optional

import numpy as np
import torch
from pathlib import Path

from fastapi import FastAPI, File, Form, UploadFile
from fastapi.responses import JSONResponse, Response
from fastapi.staticfiles import StaticFiles
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, PointStruct, VectorParams

from detect import (
    detect,
    draw_results,
    embed_masked_crop,
    load_clip,
    load_dinov2,
    load_gdino,
    load_sam2,
    segment,
    verify_detections,
)

# ── shared state ────────────────────────────────────────────────────────
models = {}

QDRANT_URL = os.getenv("QDRANT_URL", "")  # empty = use local file storage
QDRANT_PATH = os.getenv("QDRANT_PATH", "./qdrant_data")
COLLECTION_NAME = "catalog"
EMBED_DIM = 384  # DINOv2-small output dimension

# ── detection cache (stores masks + image for re-embedding with erosion) ─
_detection_cache = {}


def get_device() -> str:
    if torch.backends.mps.is_available():
        return "mps"
    if torch.cuda.is_available():
        return "cuda"
    return "cpu"


def init_qdrant() -> QdrantClient:
    if QDRANT_URL:
        print(f"Connecting to remote Qdrant at {QDRANT_URL}")
        client = QdrantClient(url=QDRANT_URL)
    else:
        print(f"Using local Qdrant storage at {QDRANT_PATH}")
        client = QdrantClient(path=QDRANT_PATH)
    collections = [c.name for c in client.get_collections().collections]
    if COLLECTION_NAME not in collections:
        client.create_collection(
            collection_name=COLLECTION_NAME,
            vectors_config=VectorParams(size=EMBED_DIM, distance=Distance.COSINE),
        )
        print(f"Created Qdrant collection '{COLLECTION_NAME}'")
    else:
        print(f"Qdrant collection '{COLLECTION_NAME}' already exists")
    return client


@asynccontextmanager
async def lifespan(app: FastAPI):
    device = get_device()
    print(f"Device: {device}")

    print("Loading Grounding DINO...")
    gdino_processor, gdino_model = load_gdino(device)

    print("Loading SAM 2...")
    sam2_processor, sam2_model = load_sam2(device)

    print("Loading CLIP...")
    clip_processor, clip_model = load_clip(device)

    print("Loading DINOv2...")
    dinov2_processor, dinov2_model = load_dinov2(device)

    print("Connecting to Qdrant...")
    qdrant = init_qdrant()

    models["device"] = device
    models["gdino_processor"] = gdino_processor
    models["gdino_model"] = gdino_model
    models["sam2_processor"] = sam2_processor
    models["sam2_model"] = sam2_model
    models["clip_processor"] = clip_processor
    models["clip_model"] = clip_model
    models["dinov2_processor"] = dinov2_processor
    models["dinov2_model"] = dinov2_model
    models["qdrant"] = qdrant

    print("Models loaded. Server ready.\n")
    yield
    models.clear()


app = FastAPI(title="Scenedrobe API", lifespan=lifespan)


@app.post("/detect")
async def detect_endpoint(
    image: UploadFile = File(...),
    prompt: str = Form(...),
    box_threshold: float = Form(0.15),
    text_threshold: float = Form(0.15),
):
    """
    Upload an image + text prompt, get back the annotated image
    with segmentation masks.
    """
    from PIL import Image as PILImage

    start = time.time()

    # read uploaded image
    image_bytes = await image.read()
    pil_image = PILImage.open(io.BytesIO(image_bytes)).convert("RGB")

    # ensure prompt ends with period
    text = prompt.strip()
    if not text.endswith("."):
        text += "."

    device = models["device"]

    # step 1: detect with Grounding DINO (low threshold for recall)
    boxes, scores, labels = detect(
        pil_image, text,
        models["gdino_processor"], models["gdino_model"], device,
        box_threshold=box_threshold,
        text_threshold=text_threshold,
    )

    # step 2: verify with CLIP (filter false positives)
    if len(boxes) > 0:
        boxes, scores, labels = verify_detections(
            pil_image, boxes, scores, labels,
            models["clip_processor"], models["clip_model"], device,
        )

    if len(boxes) == 0:
        buf = io.BytesIO()
        pil_image.save(buf, format="JPEG")
        return Response(
            content=buf.getvalue(),
            media_type="image/jpeg",
            headers={"X-Detections": "0", "X-Time-Ms": str(int((time.time() - start) * 1000))},
        )

    # step 2: segment with SAM 2
    masks = segment(pil_image, boxes, models["sam2_processor"], models["sam2_model"], device)

    # draw masks on image
    annotated = draw_results(pil_image, masks, scores, labels)

    # encode result
    buf = io.BytesIO()
    annotated.save(buf, format="JPEG", quality=90)

    elapsed_ms = int((time.time() - start) * 1000)
    n = len(boxes)
    print(f"  Prompt: {text!r} | {n} detection(s) | {elapsed_ms}ms")

    return Response(
        content=buf.getvalue(),
        media_type="image/jpeg",
        headers={
            "X-Detections": str(n),
            "X-Time-Ms": str(elapsed_ms),
        },
    )


def mask_to_polygon(mask_np: np.ndarray, tolerance: int = 2) -> list[list[int]]:
    """Extract a simplified polygon boundary from a binary mask using basic edge tracing."""
    from PIL import Image as PILImage, ImageFilter

    # find contour by eroding the mask and XOR-ing with the original
    mask_img = PILImage.fromarray(mask_np.astype(np.uint8) * 255, "L")
    eroded = mask_img.filter(ImageFilter.MinFilter(3))
    edge = np.array(mask_img).astype(bool) & ~np.array(eroded).astype(bool)

    ys, xs = np.where(edge)
    if len(ys) == 0:
        return []

    # subsample points for a lighter polygon
    step = max(1, len(ys) // 200)
    points = list(zip(xs[::step].tolist(), ys[::step].tolist()))
    return points


def mask_to_bbox(mask_np: np.ndarray) -> list[int]:
    """Tight bounding box [x_min, y_min, x_max, y_max] from a binary mask."""
    ys, xs = np.where(mask_np)
    return [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())]


def mask_to_base64_png(mask_np: np.ndarray, color: tuple = (255, 0, 0, 120)) -> str:
    """Encode a binary mask as a transparent PNG, returned as a base64 string."""
    from PIL import Image as PILImage

    rgba = np.zeros((*mask_np.shape, 4), dtype=np.uint8)
    rgba[mask_np] = color
    img = PILImage.fromarray(rgba, "RGBA")
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode()


@app.post("/detect/json")
async def detect_json_endpoint(
    image: UploadFile = File(...),
    prompt: str = Form(...),
    box_threshold: float = Form(0.15),
    text_threshold: float = Form(0.15),
    include_mask_png: bool = Form(False),
):
    """
    Upload an image + text prompt, get back JSON with:
    - bbox: tight bounding box around the mask
    - polygon: contour points for SVG/canvas overlay
    - mask_png: (optional) base64-encoded transparent PNG of the mask
    """
    from PIL import Image as PILImage

    start = time.time()

    image_bytes = await image.read()
    pil_image = PILImage.open(io.BytesIO(image_bytes)).convert("RGB")
    width, height = pil_image.size

    text = prompt.strip()
    if not text.endswith("."):
        text += "."

    device = models["device"]

    # step 1: detect with Grounding DINO (low threshold for recall)
    boxes, scores, labels = detect(
        pil_image, text,
        models["gdino_processor"], models["gdino_model"], device,
        box_threshold=box_threshold,
        text_threshold=text_threshold,
    )

    # step 2: verify with CLIP (filter false positives)
    if len(boxes) > 0:
        boxes, scores, labels = verify_detections(
            pil_image, boxes, scores, labels,
            models["clip_processor"], models["clip_model"], device,
        )

    if len(boxes) == 0:
        return JSONResponse({
            "image_size": {"width": width, "height": height},
            "detections": [],
            "time_ms": int((time.time() - start) * 1000),
        })

    # step 3: segment
    masks = segment(pil_image, boxes, models["sam2_processor"], models["sam2_model"], device)

    COLORS = [
        (255, 0, 0, 180), (0, 200, 0, 180), (0, 100, 255, 180),
        (255, 255, 0, 180), (255, 0, 255, 180), (0, 255, 255, 180),
    ]

    MIN_SCORE = 0.3

    # generate a request ID for caching masks
    request_id = str(uuid.uuid4())
    cached_masks = []

    detections = []
    for i, (mask, score, label, box) in enumerate(zip(masks, scores, labels, boxes)):
        if score.item() < MIN_SCORE:
            continue
        mask_np = mask.numpy().astype(bool)
        mask_bbox = mask_to_bbox(mask_np)
        embedding, prepared_crop = embed_masked_crop(
            pil_image, mask_np, mask_bbox,
            models["dinov2_processor"], models["dinov2_model"], device,
        )

        # thumbnail of the prepared crop (what CLIP actually sees)
        crop_thumb = _make_thumbnail(prepared_crop)

        det_index = len(detections)
        det = {
            "label": label,
            "score": round(score.item(), 3),
            "dino_bbox": [round(v.item(), 1) for v in box],
            "mask_bbox": mask_bbox,
            "polygon": mask_to_polygon(mask_np),
            "embedding": embedding,
            "crop_thumbnail": crop_thumb,
            "detection_id": f"{request_id}:{det_index}",
        }
        if include_mask_png:
            det["mask_png"] = mask_to_base64_png(mask_np, COLORS[i % len(COLORS)])
        detections.append(det)
        cached_masks.append({"mask_np": mask_np, "label": label})

    # cache image + masks for re-embedding with erosion
    _detection_cache[request_id] = {
        "image": pil_image,
        "masks": cached_masks,
        "created": time.time(),
    }
    # evict old caches (keep last 5)
    if len(_detection_cache) > 5:
        oldest = sorted(_detection_cache, key=lambda k: _detection_cache[k]["created"])
        for k in oldest[:-5]:
            del _detection_cache[k]

    elapsed_ms = int((time.time() - start) * 1000)
    print(f"  [JSON] Prompt: {text!r} | {len(detections)} detection(s) | {elapsed_ms}ms")

    return JSONResponse({
        "image_size": {"width": width, "height": height},
        "detections": detections,
        "time_ms": elapsed_ms,
    })


@app.post("/detect/re-embed")
async def detect_re_embed(
    detection_id: str = Form(...),
    mask_erosion: int = Form(0),
):
    """Re-compute embedding for a cached detection with a new erosion value.
    Returns new embedding + crop thumbnail."""
    from PIL import Image as PILImage, ImageFilter

    parts = detection_id.split(":")
    if len(parts) != 2:
        return JSONResponse({"status": "error", "message": "Invalid detection_id"}, status_code=400)
    request_id, det_idx = parts[0], int(parts[1])

    cache = _detection_cache.get(request_id)
    if not cache:
        return JSONResponse({"status": "error", "message": "Detection expired — re-run detection"}, status_code=404)

    if det_idx >= len(cache["masks"]):
        return JSONResponse({"status": "error", "message": "Invalid detection index"}, status_code=400)

    pil_image = cache["image"]
    mask_data = cache["masks"][det_idx]
    mask_np = mask_data["mask_np"].copy()
    label = mask_data["label"]
    device = models["device"]

    # apply erosion
    if mask_erosion > 0:
        kernel_size = mask_erosion * 2 + 1
        mask_img = PILImage.fromarray(mask_np.astype(np.uint8) * 255, "L")
        mask_img = mask_img.filter(ImageFilter.MinFilter(kernel_size))
        mask_np = np.array(mask_img) > 127

    ys, xs = np.where(mask_np)
    if len(ys) == 0:
        return JSONResponse({"status": "error", "message": "Mask empty after erosion"}, status_code=400)

    mask_bbox = [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())]

    # re-embed
    embedding, prepared_crop = embed_masked_crop(
        pil_image, mask_np, mask_bbox,
        models["dinov2_processor"], models["dinov2_model"], device,
    )

    crop_thumb = _make_thumbnail(prepared_crop)

    return JSONResponse({
        "status": "ok",
        "embedding": embedding,
        "crop_thumbnail": crop_thumb,
    })


def _make_thumbnail(pil_image, max_size=200, quality=80) -> str:
    """Resize and encode a PIL image as a base64 JPEG thumbnail."""
    thumb = pil_image.copy()
    thumb.thumbnail((max_size, max_size))
    buf = io.BytesIO()
    thumb.save(buf, format="JPEG", quality=quality)
    return base64.b64encode(buf.getvalue()).decode()


def _extract_garment(
    pil_image,
    label: str,
    box_threshold: float = 0.15,
    text_threshold: float = 0.15,
    clip_threshold: float = 0.45,
    mask_erosion: int = 0,
) -> dict:
    """Detect, segment, and embed a garment from an image. Returns extraction data or raises ValueError."""
    from PIL import Image as PILImage

    device = models["device"]
    prompt = label.strip()
    if not prompt.endswith("."):
        prompt += "."

    print(f"  [Extract] label={prompt!r} box_thresh={box_threshold} text_thresh={text_threshold} clip_thresh={clip_threshold}")

    boxes, scores, labels_out = detect(
        pil_image, prompt,
        models["gdino_processor"], models["gdino_model"], device,
        box_threshold=box_threshold, text_threshold=text_threshold,
    )
    print(f"  [Extract] DINO found {len(boxes)} detection(s)")

    if len(boxes) > 0:
        boxes, scores, labels_out = verify_detections(
            pil_image, boxes, scores, labels_out,
            models["clip_processor"], models["clip_model"], device,
            clip_threshold=clip_threshold,
        )

    if len(boxes) == 0:
        raise ValueError(f"Could not detect '{label}' in the image.")

    best = scores.argmax().item()
    masks = segment(pil_image, boxes[best:best+1], models["sam2_processor"], models["sam2_model"], device)
    mask = masks[0]
    mask_np = mask.numpy().astype(bool)

    # erode mask to trim bleeding edges
    if mask_erosion > 0:
        from PIL import ImageFilter
        from PIL import Image as _Img
        kernel_size = mask_erosion * 2 + 1
        mask_img = _Img.fromarray(mask_np.astype(np.uint8) * 255, "L")
        mask_img = mask_img.filter(ImageFilter.MinFilter(kernel_size))
        mask_np = np.array(mask_img) > 127

    ys, xs = np.where(mask_np)
    if len(ys) == 0:
        raise ValueError(f"Mask became empty after erosion — try a lower value.")
    mask_bbox = [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())]

    embedding, prepared_crop = embed_masked_crop(
        pil_image, mask_np, mask_bbox,
        models["dinov2_processor"], models["dinov2_model"], device,
    )

    # extracted_thumbnail = white-background crop (for visual display)
    x1, y1, x2, y2 = mask_bbox
    crop = pil_image.crop((x1, y1, x2, y2))
    crop_mask = mask_np[y1:y2, x1:x2]
    crop_np = np.array(crop)
    crop_np[~crop_mask] = 255
    extracted = PILImage.fromarray(crop_np)

    return {
        "embedding": embedding,
        "original_thumbnail": _make_thumbnail(pil_image),
        "extracted_thumbnail": _make_thumbnail(extracted),
        "prepared_thumbnail": _make_thumbnail(prepared_crop),
        "score": round(scores[best].item(), 3),
    }


@app.post("/catalog/preview")
async def catalog_preview(
    image: UploadFile = File(...),
    label: str = Form(...),
    box_threshold: float = Form(0.15),
    text_threshold: float = Form(0.15),
    clip_threshold: float = Form(0.45),
    mask_erosion: int = Form(0),
):
    """
    Run detection + extraction without saving. Returns thumbnails and embedding
    so the admin can verify before committing.
    """
    from PIL import Image as PILImage

    image_bytes = await image.read()
    pil_image = PILImage.open(io.BytesIO(image_bytes)).convert("RGB")

    try:
        result = _extract_garment(
            pil_image, label,
            box_threshold=box_threshold,
            text_threshold=text_threshold,
            clip_threshold=clip_threshold,
            mask_erosion=mask_erosion,
        )
    except ValueError as e:
        return JSONResponse({"status": "error", "message": str(e)}, status_code=400)

    return JSONResponse({
        "status": "ok",
        "label": label,
        "detection_score": result["score"],
        "original_thumbnail": result["original_thumbnail"],
        "extracted_thumbnail": result["extracted_thumbnail"],
        "prepared_thumbnail": result["prepared_thumbnail"],
        "embedding": result["embedding"],
    })


@app.post("/catalog/add")
async def catalog_add(
    label: str = Form(...),
    product_url: str = Form(""),
    price: Optional[str] = Form(None),
    embedding: str = Form(...),
    original_thumbnail: str = Form(...),
    extracted_thumbnail: str = Form(...),
    prepared_thumbnail: str = Form(""),
):
    """
    Save a pre-extracted garment to the catalog.
    Accepts the embedding and thumbnails from a prior /catalog/preview call.
    """
    import json
    embedding_vec = json.loads(embedding) if isinstance(embedding, str) else embedding

    point_id = str(uuid.uuid4())
    models["qdrant"].upsert(
        collection_name=COLLECTION_NAME,
        points=[
            PointStruct(
                id=point_id,
                vector=embedding_vec,
                payload={
                    "label": label,
                    "product_url": product_url,
                    "price": price,
                    "original_thumbnail": original_thumbnail,
                    "extracted_thumbnail": extracted_thumbnail,
                    "prepared_thumbnail": prepared_thumbnail,
                },
            )
        ],
    )

    return JSONResponse({"id": point_id, "label": label, "status": "added"})


@app.post("/catalog/search")
async def catalog_search(
    embedding: str = Form(...),
    limit: int = Form(5),
    min_score: float = Form(0.30),
):
    """Search the catalog for items similar to the given CLIP embedding vector."""
    import json
    embedding_vec = json.loads(embedding)

    results = models["qdrant"].query_points(
        collection_name=COLLECTION_NAME,
        query=embedding_vec,
        limit=limit,
    )

    items = []
    for point in results.points:
        items.append({
            "id": point.id,
            "score": round(point.score, 3),
            "label": point.payload.get("label", ""),
            "product_url": point.payload.get("product_url", ""),
            "price": point.payload.get("price"),
            "original_thumbnail": point.payload.get("original_thumbnail", ""),
            "extracted_thumbnail": point.payload.get("extracted_thumbnail", ""),
            "prepared_thumbnail": point.payload.get("prepared_thumbnail", ""),
        })

    for point in results.points:
        print(f"  [Search] id={point.id} label={point.payload.get('label','')} score={point.score:.4f}")
    print(f"  [Search] {len(results.points)} result(s) above {min_score}, returned {len(items)}")

    return JSONResponse({"results": items})


@app.get("/catalog/count")
async def catalog_count():
    """Return the number of items in the catalog."""
    info = models["qdrant"].get_collection(collection_name=COLLECTION_NAME)
    return {"count": info.points_count}


@app.get("/catalog/list")
async def catalog_list(
    limit: int = 50,
    offset: Optional[str] = None,
):
    """List all catalog items with pagination."""
    result = models["qdrant"].scroll(
        collection_name=COLLECTION_NAME,
        limit=limit,
        offset=offset,
        with_vectors=False,
        with_payload=True,
    )

    points, next_offset = result
    items = []
    for point in points:
        items.append({
            "id": point.id,
            "label": point.payload.get("label", ""),
            "product_url": point.payload.get("product_url", ""),
            "price": point.payload.get("price"),
            "original_thumbnail": point.payload.get("original_thumbnail", ""),
            "extracted_thumbnail": point.payload.get("extracted_thumbnail", ""),
        })

    return JSONResponse({
        "items": items,
        "next_offset": next_offset,
    })


@app.delete("/catalog/{item_id}")
async def catalog_delete(item_id: str):
    """Delete a single item from the catalog."""
    models["qdrant"].delete(
        collection_name=COLLECTION_NAME,
        points_selector=[item_id],
    )
    return JSONResponse({"id": item_id, "status": "deleted"})


@app.get("/health")
async def health():
    return {"status": "ok", "models_loaded": bool(models)}


# serve frontend — must be last so API routes take priority
WEB_DIR = Path(__file__).resolve().parent / "web"
app.mount("/", StaticFiles(directory=WEB_DIR, html=True), name="web")

	"""
	Scenedrobe API – FastAPI server for Grounded SAM 2.

	Models are loaded once at startup and kept in memory.
	"""

	from __future__ import annotations

	import base64
	import io
	import os
	import time
	import uuid
	from contextlib import asynccontextmanager
	from typing import Optional

	import numpy as np
	import torch
	from pathlib import Path

	from fastapi import FastAPI, File, Form, UploadFile
	from fastapi.responses import JSONResponse, Response
	from fastapi.staticfiles import StaticFiles
	from qdrant_client import QdrantClient
	from qdrant_client.models import Distance, PointStruct, VectorParams

	from detect import (
	detect,
	draw_results,
	embed_masked_crop,
	load_clip,
	load_dinov2,
	load_gdino,
	load_sam2,
	segment,
	verify_detections,
	)

	# ── shared state ────────────────────────────────────────────────────────
	models = {}

	QDRANT_URL = os.getenv("QDRANT_URL", "") # empty = use local file storage
	QDRANT_PATH = os.getenv("QDRANT_PATH", "./qdrant_data")
	COLLECTION_NAME = "catalog"
	EMBED_DIM = 384 # DINOv2-small output dimension

	# ── detection cache (stores masks + image for re-embedding with erosion) ─
	_detection_cache = {}


	def get_device() -> str:
	if torch.backends.mps.is_available():
	return "mps"
	if torch.cuda.is_available():
	return "cuda"
	return "cpu"


	def init_qdrant() -> QdrantClient:
	if QDRANT_URL:
	print(f"Connecting to remote Qdrant at {QDRANT_URL}")
	client = QdrantClient(url=QDRANT_URL)
	else:
	print(f"Using local Qdrant storage at {QDRANT_PATH}")
	client = QdrantClient(path=QDRANT_PATH)
	collections = [c.name for c in client.get_collections().collections]
	if COLLECTION_NAME not in collections:
	client.create_collection(
	collection_name=COLLECTION_NAME,
	vectors_config=VectorParams(size=EMBED_DIM, distance=Distance.COSINE),
	)
	print(f"Created Qdrant collection '{COLLECTION_NAME}'")
	else:
	print(f"Qdrant collection '{COLLECTION_NAME}' already exists")
	return client


	@asynccontextmanager
	async def lifespan(app: FastAPI):
	device = get_device()
	print(f"Device: {device}")

	print("Loading Grounding DINO...")
	gdino_processor, gdino_model = load_gdino(device)

	print("Loading SAM 2...")
	sam2_processor, sam2_model = load_sam2(device)

	print("Loading CLIP...")
	clip_processor, clip_model = load_clip(device)

	print("Loading DINOv2...")
	dinov2_processor, dinov2_model = load_dinov2(device)

	print("Connecting to Qdrant...")
	qdrant = init_qdrant()

	models["device"] = device
	models["gdino_processor"] = gdino_processor
	models["gdino_model"] = gdino_model
	models["sam2_processor"] = sam2_processor
	models["sam2_model"] = sam2_model
	models["clip_processor"] = clip_processor
	models["clip_model"] = clip_model
	models["dinov2_processor"] = dinov2_processor
	models["dinov2_model"] = dinov2_model
	models["qdrant"] = qdrant

	print("Models loaded. Server ready.\n")
	yield
	models.clear()


	app = FastAPI(title="Scenedrobe API", lifespan=lifespan)


	@app.post("/detect")
	async def detect_endpoint(
	image: UploadFile = File(...),
	prompt: str = Form(...),
	box_threshold: float = Form(0.15),
	text_threshold: float = Form(0.15),
	):
	"""
	Upload an image + text prompt, get back the annotated image
	with segmentation masks.
	"""
	from PIL import Image as PILImage

	start = time.time()

	# read uploaded image
	image_bytes = await image.read()
	pil_image = PILImage.open(io.BytesIO(image_bytes)).convert("RGB")

	# ensure prompt ends with period
	text = prompt.strip()
	if not text.endswith("."):
	text += "."

	device = models["device"]

	# step 1: detect with Grounding DINO (low threshold for recall)
	boxes, scores, labels = detect(
	pil_image, text,
	models["gdino_processor"], models["gdino_model"], device,
	box_threshold=box_threshold,
	text_threshold=text_threshold,
	)

	# step 2: verify with CLIP (filter false positives)
	if len(boxes) > 0:
	boxes, scores, labels = verify_detections(
	pil_image, boxes, scores, labels,
	models["clip_processor"], models["clip_model"], device,
	)

	if len(boxes) == 0:
	buf = io.BytesIO()
	pil_image.save(buf, format="JPEG")
	return Response(
	content=buf.getvalue(),
	media_type="image/jpeg",
	headers={"X-Detections": "0", "X-Time-Ms": str(int((time.time() - start) * 1000))},
	)

	# step 2: segment with SAM 2
	masks = segment(pil_image, boxes, models["sam2_processor"], models["sam2_model"], device)

	# draw masks on image
	annotated = draw_results(pil_image, masks, scores, labels)

	# encode result
	buf = io.BytesIO()
	annotated.save(buf, format="JPEG", quality=90)

	elapsed_ms = int((time.time() - start) * 1000)
	n = len(boxes)
	print(f" Prompt: {text!r} \| {n} detection(s) \| {elapsed_ms}ms")

	return Response(
	content=buf.getvalue(),
	media_type="image/jpeg",
	headers={
	"X-Detections": str(n),
	"X-Time-Ms": str(elapsed_ms),
	},
	)


	def mask_to_polygon(mask_np: np.ndarray, tolerance: int = 2) -> list[list[int]]:
	"""Extract a simplified polygon boundary from a binary mask using basic edge tracing."""
	from PIL import Image as PILImage, ImageFilter

	# find contour by eroding the mask and XOR-ing with the original
	mask_img = PILImage.fromarray(mask_np.astype(np.uint8) * 255, "L")
	eroded = mask_img.filter(ImageFilter.MinFilter(3))
	edge = np.array(mask_img).astype(bool) & ~np.array(eroded).astype(bool)

	ys, xs = np.where(edge)
	if len(ys) == 0:
	return []

	# subsample points for a lighter polygon
	step = max(1, len(ys) // 200)
	points = list(zip(xs[::step].tolist(), ys[::step].tolist()))
	return points


	def mask_to_bbox(mask_np: np.ndarray) -> list[int]:
	"""Tight bounding box [x_min, y_min, x_max, y_max] from a binary mask."""
	ys, xs = np.where(mask_np)
	return [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())]


	def mask_to_base64_png(mask_np: np.ndarray, color: tuple = (255, 0, 0, 120)) -> str:
	"""Encode a binary mask as a transparent PNG, returned as a base64 string."""
	from PIL import Image as PILImage

	rgba = np.zeros((*mask_np.shape, 4), dtype=np.uint8)
	rgba[mask_np] = color
	img = PILImage.fromarray(rgba, "RGBA")
	buf = io.BytesIO()
	img.save(buf, format="PNG")
	return base64.b64encode(buf.getvalue()).decode()


	@app.post("/detect/json")
	async def detect_json_endpoint(
	image: UploadFile = File(...),
	prompt: str = Form(...),
	box_threshold: float = Form(0.15),
	text_threshold: float = Form(0.15),
	include_mask_png: bool = Form(False),
	):
	"""
	Upload an image + text prompt, get back JSON with:
	- bbox: tight bounding box around the mask
	- polygon: contour points for SVG/canvas overlay
	- mask_png: (optional) base64-encoded transparent PNG of the mask
	"""
	from PIL import Image as PILImage

	start = time.time()

	image_bytes = await image.read()
	pil_image = PILImage.open(io.BytesIO(image_bytes)).convert("RGB")
	width, height = pil_image.size

	text = prompt.strip()
	if not text.endswith("."):
	text += "."

	device = models["device"]

	# step 1: detect with Grounding DINO (low threshold for recall)
	boxes, scores, labels = detect(
	pil_image, text,
	models["gdino_processor"], models["gdino_model"], device,
	box_threshold=box_threshold,
	text_threshold=text_threshold,
	)

	# step 2: verify with CLIP (filter false positives)
	if len(boxes) > 0:
	boxes, scores, labels = verify_detections(
	pil_image, boxes, scores, labels,
	models["clip_processor"], models["clip_model"], device,
	)

	if len(boxes) == 0:
	return JSONResponse({
	"image_size": {"width": width, "height": height},
	"detections": [],
	"time_ms": int((time.time() - start) * 1000),
	})

	# step 3: segment
	masks = segment(pil_image, boxes, models["sam2_processor"], models["sam2_model"], device)

	COLORS = [
	(255, 0, 0, 180), (0, 200, 0, 180), (0, 100, 255, 180),
	(255, 255, 0, 180), (255, 0, 255, 180), (0, 255, 255, 180),
	]

	MIN_SCORE = 0.3

	# generate a request ID for caching masks
	request_id = str(uuid.uuid4())
	cached_masks = []

	detections = []
	for i, (mask, score, label, box) in enumerate(zip(masks, scores, labels, boxes)):
	if score.item() < MIN_SCORE:
	continue
	mask_np = mask.numpy().astype(bool)
	mask_bbox = mask_to_bbox(mask_np)
	embedding, prepared_crop = embed_masked_crop(
	pil_image, mask_np, mask_bbox,
	models["dinov2_processor"], models["dinov2_model"], device,
	)

	# thumbnail of the prepared crop (what CLIP actually sees)
	crop_thumb = _make_thumbnail(prepared_crop)

	det_index = len(detections)
	det = {
	"label": label,
	"score": round(score.item(), 3),
	"dino_bbox": [round(v.item(), 1) for v in box],
	"mask_bbox": mask_bbox,
	"polygon": mask_to_polygon(mask_np),
	"embedding": embedding,
	"crop_thumbnail": crop_thumb,
	"detection_id": f"{request_id}:{det_index}",
	}
	if include_mask_png:
	det["mask_png"] = mask_to_base64_png(mask_np, COLORS[i % len(COLORS)])
	detections.append(det)
	cached_masks.append({"mask_np": mask_np, "label": label})

	# cache image + masks for re-embedding with erosion
	_detection_cache[request_id] = {
	"image": pil_image,
	"masks": cached_masks,
	"created": time.time(),
	}
	# evict old caches (keep last 5)
	if len(_detection_cache) > 5:
	oldest = sorted(_detection_cache, key=lambda k: _detection_cache[k]["created"])
	for k in oldest[:-5]:
	del _detection_cache[k]

	elapsed_ms = int((time.time() - start) * 1000)
	print(f" [JSON] Prompt: {text!r} \| {len(detections)} detection(s) \| {elapsed_ms}ms")

	return JSONResponse({
	"image_size": {"width": width, "height": height},
	"detections": detections,
	"time_ms": elapsed_ms,
	})


	@app.post("/detect/re-embed")
	async def detect_re_embed(
	detection_id: str = Form(...),
	mask_erosion: int = Form(0),
	):
	"""Re-compute embedding for a cached detection with a new erosion value.
	Returns new embedding + crop thumbnail."""
	from PIL import Image as PILImage, ImageFilter

	parts = detection_id.split(":")
	if len(parts) != 2:
	return JSONResponse({"status": "error", "message": "Invalid detection_id"}, status_code=400)
	request_id, det_idx = parts[0], int(parts[1])

	cache = _detection_cache.get(request_id)
	if not cache:
	return JSONResponse({"status": "error", "message": "Detection expired — re-run detection"}, status_code=404)

	if det_idx >= len(cache["masks"]):
	return JSONResponse({"status": "error", "message": "Invalid detection index"}, status_code=400)

	pil_image = cache["image"]
	mask_data = cache["masks"][det_idx]
	mask_np = mask_data["mask_np"].copy()
	label = mask_data["label"]
	device = models["device"]

	# apply erosion
	if mask_erosion > 0:
	kernel_size = mask_erosion * 2 + 1
	mask_img = PILImage.fromarray(mask_np.astype(np.uint8) * 255, "L")
	mask_img = mask_img.filter(ImageFilter.MinFilter(kernel_size))
	mask_np = np.array(mask_img) > 127

	ys, xs = np.where(mask_np)
	if len(ys) == 0:
	return JSONResponse({"status": "error", "message": "Mask empty after erosion"}, status_code=400)

	mask_bbox = [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())]

	# re-embed
	embedding, prepared_crop = embed_masked_crop(
	pil_image, mask_np, mask_bbox,
	models["dinov2_processor"], models["dinov2_model"], device,
	)

	crop_thumb = _make_thumbnail(prepared_crop)

	return JSONResponse({
	"status": "ok",
	"embedding": embedding,
	"crop_thumbnail": crop_thumb,
	})


	def _make_thumbnail(pil_image, max_size=200, quality=80) -> str:
	"""Resize and encode a PIL image as a base64 JPEG thumbnail."""
	thumb = pil_image.copy()
	thumb.thumbnail((max_size, max_size))
	buf = io.BytesIO()
	thumb.save(buf, format="JPEG", quality=quality)
	return base64.b64encode(buf.getvalue()).decode()


	def _extract_garment(
	pil_image,
	label: str,
	box_threshold: float = 0.15,
	text_threshold: float = 0.15,
	clip_threshold: float = 0.45,
	mask_erosion: int = 0,
	) -> dict:
	"""Detect, segment, and embed a garment from an image. Returns extraction data or raises ValueError."""
	from PIL import Image as PILImage

	device = models["device"]
	prompt = label.strip()
	if not prompt.endswith("."):
	prompt += "."

	print(f" [Extract] label={prompt!r} box_thresh={box_threshold} text_thresh={text_threshold} clip_thresh={clip_threshold}")

	boxes, scores, labels_out = detect(
	pil_image, prompt,
	models["gdino_processor"], models["gdino_model"], device,
	box_threshold=box_threshold, text_threshold=text_threshold,
	)
	print(f" [Extract] DINO found {len(boxes)} detection(s)")

	if len(boxes) > 0:
	boxes, scores, labels_out = verify_detections(
	pil_image, boxes, scores, labels_out,
	models["clip_processor"], models["clip_model"], device,
	clip_threshold=clip_threshold,
	)

	if len(boxes) == 0:
	raise ValueError(f"Could not detect '{label}' in the image.")

	best = scores.argmax().item()
	masks = segment(pil_image, boxes[best:best+1], models["sam2_processor"], models["sam2_model"], device)
	mask = masks[0]
	mask_np = mask.numpy().astype(bool)

	# erode mask to trim bleeding edges
	if mask_erosion > 0:
	from PIL import ImageFilter
	from PIL import Image as _Img
	kernel_size = mask_erosion * 2 + 1
	mask_img = _Img.fromarray(mask_np.astype(np.uint8) * 255, "L")
	mask_img = mask_img.filter(ImageFilter.MinFilter(kernel_size))
	mask_np = np.array(mask_img) > 127

	ys, xs = np.where(mask_np)
	if len(ys) == 0:
	raise ValueError(f"Mask became empty after erosion — try a lower value.")
	mask_bbox = [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())]

	embedding, prepared_crop = embed_masked_crop(
	pil_image, mask_np, mask_bbox,
	models["dinov2_processor"], models["dinov2_model"], device,
	)

	# extracted_thumbnail = white-background crop (for visual display)
	x1, y1, x2, y2 = mask_bbox
	crop = pil_image.crop((x1, y1, x2, y2))
	crop_mask = mask_np[y1:y2, x1:x2]
	crop_np = np.array(crop)
	crop_np[~crop_mask] = 255
	extracted = PILImage.fromarray(crop_np)

	return {
	"embedding": embedding,
	"original_thumbnail": _make_thumbnail(pil_image),
	"extracted_thumbnail": _make_thumbnail(extracted),
	"prepared_thumbnail": _make_thumbnail(prepared_crop),
	"score": round(scores[best].item(), 3),
	}


	@app.post("/catalog/preview")
	async def catalog_preview(
	image: UploadFile = File(...),
	label: str = Form(...),
	box_threshold: float = Form(0.15),
	text_threshold: float = Form(0.15),
	clip_threshold: float = Form(0.45),
	mask_erosion: int = Form(0),
	):
	"""
	Run detection + extraction without saving. Returns thumbnails and embedding
	so the admin can verify before committing.
	"""
	from PIL import Image as PILImage

	image_bytes = await image.read()
	pil_image = PILImage.open(io.BytesIO(image_bytes)).convert("RGB")

	try:
	result = _extract_garment(
	pil_image, label,
	box_threshold=box_threshold,
	text_threshold=text_threshold,
	clip_threshold=clip_threshold,
	mask_erosion=mask_erosion,
	)
	except ValueError as e:
	return JSONResponse({"status": "error", "message": str(e)}, status_code=400)

	return JSONResponse({
	"status": "ok",
	"label": label,
	"detection_score": result["score"],
	"original_thumbnail": result["original_thumbnail"],
	"extracted_thumbnail": result["extracted_thumbnail"],
	"prepared_thumbnail": result["prepared_thumbnail"],
	"embedding": result["embedding"],
	})


	@app.post("/catalog/add")
	async def catalog_add(
	label: str = Form(...),
	product_url: str = Form(""),
	price: Optional[str] = Form(None),
	embedding: str = Form(...),
	original_thumbnail: str = Form(...),
	extracted_thumbnail: str = Form(...),
	prepared_thumbnail: str = Form(""),
	):
	"""
	Save a pre-extracted garment to the catalog.
	Accepts the embedding and thumbnails from a prior /catalog/preview call.
	"""
	import json
	embedding_vec = json.loads(embedding) if isinstance(embedding, str) else embedding

	point_id = str(uuid.uuid4())
	models["qdrant"].upsert(
	collection_name=COLLECTION_NAME,
	points=[
	PointStruct(
	id=point_id,
	vector=embedding_vec,
	payload={
	"label": label,
	"product_url": product_url,
	"price": price,
	"original_thumbnail": original_thumbnail,
	"extracted_thumbnail": extracted_thumbnail,
	"prepared_thumbnail": prepared_thumbnail,
	},
	)
	],
	)

	return JSONResponse({"id": point_id, "label": label, "status": "added"})


	@app.post("/catalog/search")
	async def catalog_search(
	embedding: str = Form(...),
	limit: int = Form(5),
	min_score: float = Form(0.30),
	):
	"""Search the catalog for items similar to the given CLIP embedding vector."""
	import json
	embedding_vec = json.loads(embedding)

	results = models["qdrant"].query_points(
	collection_name=COLLECTION_NAME,
	query=embedding_vec,
	limit=limit,
	)

	items = []
	for point in results.points:
	items.append({
	"id": point.id,
	"score": round(point.score, 3),
	"label": point.payload.get("label", ""),
	"product_url": point.payload.get("product_url", ""),
	"price": point.payload.get("price"),
	"original_thumbnail": point.payload.get("original_thumbnail", ""),
	"extracted_thumbnail": point.payload.get("extracted_thumbnail", ""),
	"prepared_thumbnail": point.payload.get("prepared_thumbnail", ""),
	})

	for point in results.points:
	print(f" [Search] id={point.id} label={point.payload.get('label','')} score={point.score:.4f}")
	print(f" [Search] {len(results.points)} result(s) above {min_score}, returned {len(items)}")

	return JSONResponse({"results": items})


	@app.get("/catalog/count")
	async def catalog_count():
	"""Return the number of items in the catalog."""
	info = models["qdrant"].get_collection(collection_name=COLLECTION_NAME)
	return {"count": info.points_count}


	@app.get("/catalog/list")
	async def catalog_list(
	limit: int = 50,
	offset: Optional[str] = None,
	):
	"""List all catalog items with pagination."""
	result = models["qdrant"].scroll(
	collection_name=COLLECTION_NAME,
	limit=limit,
	offset=offset,
	with_vectors=False,
	with_payload=True,
	)

	points, next_offset = result
	items = []
	for point in points:
	items.append({
	"id": point.id,
	"label": point.payload.get("label", ""),
	"product_url": point.payload.get("product_url", ""),
	"price": point.payload.get("price"),
	"original_thumbnail": point.payload.get("original_thumbnail", ""),
	"extracted_thumbnail": point.payload.get("extracted_thumbnail", ""),
	})

	return JSONResponse({
	"items": items,
	"next_offset": next_offset,
	})


	@app.delete("/catalog/{item_id}")
	async def catalog_delete(item_id: str):
	"""Delete a single item from the catalog."""
	models["qdrant"].delete(
	collection_name=COLLECTION_NAME,
	points_selector=[item_id],
	)
	return JSONResponse({"id": item_id, "status": "deleted"})


	@app.get("/health")
	async def health():
	return {"status": "ok", "models_loaded": bool(models)}


	# serve frontend — must be last so API routes take priority
	WEB_DIR = Path(__file__).resolve().parent / "web"
	app.mount("/", StaticFiles(directory=WEB_DIR, html=True), name="web")