sam / app.py
sezer91's picture
ds
7410721
raw
history blame
1.74 kB
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from PIL import Image
import numpy as np
from transformers import SamModel, SamProcessor
import io
import base64
import uvicorn
app = FastAPI(title="SAM-ViT-Base API")
# SAM modelini ve işlemciyi yükle
model = SamModel.from_pretrained("facebook/sam-vit-base")
processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
@app.post("/segment/")
async def segment_image(file: UploadFile = File(...)):
try:
# Görüntüyü oku
image_data = await file.read()
image = Image.open(io.BytesIO(image_data)).convert("RGB")
# Görüntüyü işlemciye hazırla
inputs = processor(image, return_tensors="pt")
# Model ile segmentasyon yap
outputs = model(**inputs)
# Maskeyi al
masks = outputs.pred_masks.detach().cpu().numpy()
mask = masks[0][0] # İlk maskeyi al (örnek olarak)
# Maskeyi binary hale getir
mask = (mask > 0).astype(np.uint8) * 255
# Maskeyi görüntü olarak kaydet
mask_image = Image.fromarray(mask)
buffered = io.BytesIO()
mask_image.save(buffered, format="PNG")
mask_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
return JSONResponse(content={"mask": f"data:image/png;base64,{mask_base64}"})
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/")
async def root():
return {"message": "SAM-ViT-Base API çalışıyor. /segment endpoint'ine görüntü yükleyin."}
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860)