Update app.py

app.py

@@ -1,32 +1,29 @@
 import torch
-from transformers import AutoImageProcessor
+from transformers import AutoImageProcessor, AutoModelForImageSegmentation
 from PIL import Image
 import requests
-from io import BytesIO
-
-# آدرس مدل
-model_url = "https://huggingface.co/facebook/sapiens/tree/main/sapiens_lite_host/torchscript/normal/checkpoints/sapiens_0.3b"
-
-model = torch.jit.load(model_url)
+import gradio as gr
 
-# بارگذاری پردازشگر تصویر
+# لود پردازشگر تصویر و مدل
 processor = AutoImageProcessor.from_pretrained("facebook/sapiens")
+model = AutoModelForImageSegmentation.from_pretrained("facebook/sapiens")
 
+# تابع برای پردازش تصویر ورودی و اعمال مدل بر روی آن
 def segment_image(image):
     # پردازش تصویر
     inputs = processor(images=image, return_tensors="pt")
     
     # اجرای مدل روی تصویر پردازش شده
     with torch.no_grad():
-        outputs = model(inputs['pixel_values'])
+        outputs = model(**inputs)
     
-    # فرض می‌کنیم خروجی یک ماسک است
-    segmentation = outputs.argmax(dim=1).detach().cpu().numpy()[0]
+    # فرض می‌کنیم خروجی یک ماسک است (مثل کلاس‌بندی پیکسل‌ها)
+    segmentation = outputs.logits.argmax(dim=1).detach().cpu().numpy()[0]
     
+    # بازگرداندن نتیجه به صورت تصویر
     return Image.fromarray(segmentation)
 
-# رابط Gradio
-import gradio as gr
+# ایجاد رابط Gradio برای بارگذاری تصویر و نمایش نتیجه
 interface = gr.Interface(
     fn=segment_image,
     inputs=gr.Image(type="pil"),
@@ -34,4 +31,5 @@ interface = gr.Interface(
     title="Sapiens Body Part Segmentation"
 )
 
+# اجرای برنامه
 interface.launch()