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

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."}