Upload 2 files

app.py

@@ -1,20 +1,53 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, File, UploadFile
+from transformers import SamModel, SamProcessor
+import torch
+from PIL import Image
+import io
+import numpy as np
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 
 app = FastAPI(title="SAM_MedTesting")
+MODEL_ID = "facebook/sam-vit-base"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+processor = SamProcessor.from_pretrained(MODEL_ID)
+model     = SamModel.from_pretrained(MODEL_ID).to(device)
 
-class GenerationRequest(BaseModel):
-    prompt: str
-    max_new_tokens: int = 50
+@app.post("/segment")
+async def segment(
+    file: UploadFile = File(...),
+    # an optional list of prompt points [[x1,y1], [x2,y2], …]
+    input_points: list[list[int]] | None = None
+    ):
+    
+    # read & decode image
+    data = await file.read()
+    img  = Image.open(io.BytesIO(data)).convert("RGB")
 
-class GenerationResponse(BaseModel):
-    generated_text: str
+    # prepare inputs for SAM
+    inputs = processor(
+        img,
+        input_points=[input_points] if input_points else None,
+        return_tensors="pt"
+    ).to(device)
 
-@app.post("/generate", response_model=GenerationResponse)
-def generate(req: GenerationRequest):
-    out = f"hello world: {req}"
-    return {"generated_text": out}
+    # forward pass
+    with torch.no_grad():
+        outputs = model(**inputs)
 
-@app.get("/health")
-def health():
-    return {"status": "ok"}
+    # post-process to binary masks
+    masks = processor.image_processor.post_process_masks(
+        outputs.pred_masks.cpu(),
+        inputs["original_sizes"].cpu(),
+        inputs["reshaped_input_sizes"].cpu()
+    )
+    # take the first mask
+    mask_np = (masks[0][0] * 255).astype(np.uint8)
+
+    # turn it into a PNG
+    pil_mask = Image.fromarray(mask_np)
+    buf = io.BytesIO()
+    pil_mask.save(buf, format="PNG")
+    buf.seek(0)
+
+    return StreamingResponse(buf, media_type="image/png")

requirements.txt

@@ -1,4 +1,6 @@
 fastapi
 uvicorn[standard]
 transformers
-torch
+torch
+cv2
+Pillow