from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import Response
from PIL import Image
import torch
from torchvision import transforms
from transformers import AutoModelForImageSegmentation
import io

app = FastAPI()

# ========= تحميل الموديل =========
device = "cuda" if torch.cuda.is_available() else "cpu"
birefnet = AutoModelForImageSegmentation.from_pretrained(
    "ZhengPeng7/BiRefNet", trust_remote_code=True
).to(device)
birefnet.eval()

# ========= معالجة الصورة =========
transform_image = transforms.Compose([
    transforms.Resize((1024, 1024)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

@app.post("/process")
async def process_image(data: UploadFile = File(...), mode: str = Form("r")):
    # 1. قراءة الصورة
    content = await data.read()
    im = Image.open(io.BytesIO(content)).convert("RGB")
    original_size = im.size

    # 2. إزالة الخلفية
    input_tensor = transform_image(im).unsqueeze(0).to(device)
    with torch.no_grad():
        preds = birefnet(input_tensor)[-1].sigmoid().cpu()
    mask = transforms.ToPILImage()(preds[0].squeeze()).resize(original_size)

    # 3. التطبيق (شفاف)
    transparent = im.convert("RGBA")
    transparent.putalpha(mask)

    # 4. لو المطلوب أبيض (Mode w)
    if mode.lower() == "w":
        bg = Image.new("RGBA", transparent.size, (255, 255, 255, 255))
        bg.paste(transparent, (0, 0), transparent)
        final_img = bg.convert("RGB")
        format_type = "JPEG"
    else:
        final_img = transparent
        format_type = "PNG"

    # 5. تحويل الصورة لـ Bytes عشان ترجع لـ n8n
    img_byte_arr = io.BytesIO()
    final_img.save(img_byte_arr, format=format_type)
    return Response(content=img_byte_arr.getvalue(), media_type=f"image/{format_type.lower()}")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)