Spaces:
Runtime error
Runtime error
File size: 3,452 Bytes
68a6918 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | import io
import json
import numpy as np
import torch
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import StreamingResponse
from PIL import Image, ImageDraw
from transformers import SamModel, SamProcessor
app = FastAPI(title="MedSAM Image Segmentation Service")
# Global device initialization (Fall back gracefully to CPU on standard HF spaces)
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Loading MedSAM model onto: {device}...")
processor = SamProcessor.from_pretrained("flaviagiammarino/medsam-vit-base")
model = SamModel.from_pretrained("flaviagiammarino/medsam-vit-base").to(device)
model.eval()
print("Model loaded successfully.")
@app.get("/")
def health_check():
return {"status": "healthy", "device": device}
@app.post("/segment")
async def segment_image(
file: UploadFile = File(...),
bbox: str = Form(..., description="JSON string array format: '[xmin, ymin, xmax, ymax]'")
):
# 1. Parse bounding box coordinates
try:
box_coords = json.loads(bbox)
if not isinstance(box_coords, list) or len(box_coords) != 4:
raise ValueError
except ValueError:
raise HTTPException(status_code=400, detail="Bounding box format must be '[xmin, ymin, xmax, ymax]'")
# 2. Read image
try:
image_bytes = await file.read()
raw_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
except Exception:
raise HTTPException(status_code=400, detail="Invalid image file format uploaded.")
# 3. Process image and run inference
try:
inputs = processor(
raw_image,
input_boxes=[[box_coords]],
return_tensors="pt"
).to(device)
with torch.no_grad():
outputs = model(**inputs)
# Post-process low resolution logits into a concrete mask array
masks = processor.image_processor.post_process_masks(
outputs.pred_masks.cpu(),
inputs["original_sizes"].cpu(),
inputs["reshaped_input_sizes"].cpu()
)
# Binary mask array (True/False boolean configuration)
mask_array = masks[0][0][0].numpy()
except Exception as e:
raise HTTPException(
status_code=500,
detail=f"MedSAM core processing exception: {str(e)}"
)
# 4. Generate the composite visual output (matching reference target)
# Create an overlay layout where True items highlight yellow at an alpha state
yellow_overlay = Image.new(
"RGBA",
raw_image.size,
(251, 252, 30, 150)
)
mask_image = Image.fromarray(
(mask_array * 255).astype(np.uint8),
mode="L"
)
# Composite the isolated mask layout with the underlying source picture
segmented_image = Image.composite(
yellow_overlay,
raw_image.convert("RGBA"),
mask_image
)
# Draw the diagnostic bounding box outline border matching your input configuration
draw = ImageDraw.Draw(segmented_image)
draw.rectangle(
box_coords,
outline="blue",
width=3
)
# Stream back final constructed image payload
buffer = io.BytesIO()
segmented_image.convert("RGB").save(
buffer, format="PNG"
)
buffer.seek(0)
return StreamingResponse(
buffer,
media_type="image/png"
) |