import torch
from transformers import AutoImageProcessor, AutoModelForImageSegmentation
from PIL import Image
import requests
import gradio as gr

# لود پردازشگر تصویر و مدل
processor = AutoImageProcessor.from_pretrained("facebook/sapiens/tree/main/sapiens_lite_host/torchscript/seg/checkpoints/sapiens_0.3b")
model = AutoModelForImageSegmentation.from_pretrained("facebook/sapiens/tree/main/sapiens_lite_host/torchscript/seg/checkpoints/sapiens_0.3b")

# تابع برای پردازش تصویر ورودی و اعمال مدل بر روی آن
def segment_image(image):
    # پردازش تصویر
    inputs = processor(images=image, return_tensors="pt")
    
    # اجرای مدل روی تصویر پردازش شده
    with torch.no_grad():
        outputs = model(**inputs)
    
    # فرض می‌کنیم خروجی یک ماسک است (مثل کلاس‌بندی پیکسل‌ها)
    segmentation = outputs.logits.argmax(dim=1).detach().cpu().numpy()[0]
    
    # بازگرداندن نتیجه به صورت تصویر
    return Image.fromarray(segmentation)

# ایجاد رابط Gradio برای بارگذاری تصویر و نمایش نتیجه
interface = gr.Interface(
    fn=segment_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Sapiens Body Part Segmentation"
)

# اجرای برنامه
interface.launch()