Upload app.py

app.py

@@ -1,53 +1,48 @@
-from fastapi import FastAPI, File, UploadFile
-from transformers import SamModel, SamProcessor
-import torch
+from fastapi import FastAPI
+from pydantic import BaseModel
+from typing import Optional, List
+import base64, io, json
 from PIL import Image
-import io
-import numpy as np
+import torch
+from transformers import SamModel, SamProcessor
 from fastapi.responses import StreamingResponse
-from pydantic import BaseModel
 
-app = FastAPI(title="SAM_MedTesting")
+class SegmentRequest(BaseModel):
+    file_b64: str
+    input_points: Optional[List[List[int]]] = None
+
+app = FastAPI(title="SAM MedTesting")
+
 MODEL_ID = "facebook/sam-vit-base"
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device   = "cuda" if torch.cuda.is_available() else "cpu"
 processor = SamProcessor.from_pretrained(MODEL_ID)
 model     = SamModel.from_pretrained(MODEL_ID).to(device)
 
 @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")
-
-    # prepare inputs for SAM
+async def segment(req: SegmentRequest):
+    # decode image
+    img_bytes = base64.b64decode(req.file_b64)
+    img       = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+
+    # prepare inputs
+    pts = req.input_points
     inputs = processor(
         img,
-        input_points=[input_points] if input_points else None,
+        input_points=[pts] if pts else None,
         return_tensors="pt"
     ).to(device)
 
-    # forward pass
     with torch.no_grad():
         outputs = model(**inputs)
 
-    # 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)
+    mask_np = (masks[0][0] * 255).astype("uint8")
     buf = io.BytesIO()
-    pil_mask.save(buf, format="PNG")
+    Image.fromarray(mask_np).save(buf, format="PNG")
     buf.seek(0)
 
     return StreamingResponse(buf, media_type="image/png")