Update app.py

app.py

@@ -1,35 +1,37 @@
 import torch
-from transformers import AutoImageProcessor, AutoModelForImageSegmentation
+import gradio as gr
 from PIL import Image
 import requests
-import gradio as gr
+from io import BytesIO
 
-# لود پردازشگر تصویر و مدل
-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")
+# Load the model
+model_url = "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/normal/checkpoints/sapiens_0.3b/sapiens_0.3b.pt"
+model = torch.jit.load(model_url, map_location=torch.device('cpu'))
 
-# تابع برای پردازش تصویر ورودی و اعمال مدل بر روی آن
-def segment_image(image):
-    # پردازش تصویر
-    inputs = processor(images=image, return_tensors="pt")
+# Define inference function
+def predict(image):
+    # Preprocess image
+    image = image.convert("RGB")
+    input_tensor = torch.from_numpy(np.array(image)).permute(2, 0, 1).unsqueeze(0).float() / 255.0
     
-    # اجرای مدل روی تصویر پردازش شده
+    # Run model
     with torch.no_grad():
-        outputs = model(**inputs)
-    
-    # فرض می‌کنیم خروجی یک ماسک است (مثل کلاس‌بندی پیکسل‌ها)
-    segmentation = outputs.logits.argmax(dim=1).detach().cpu().numpy()[0]
+        output = model(input_tensor)
     
-    # بازگرداندن نتیجه به صورت تصویر
-    return Image.fromarray(segmentation)
+    # Postprocess output
+    output_image = output.squeeze().permute(1, 2, 0).numpy()
+    output_image = (output_image * 255).astype(np.uint8)
+    return Image.fromarray(output_image)
 
-# ایجاد رابط Gradio برای بارگذاری تصویر و نمایش نتیجه
-interface = gr.Interface(
-    fn=segment_image,
+# Gradio Interface
+iface = gr.Interface(
+    fn=predict, 
     inputs=gr.Image(type="pil"),
     outputs=gr.Image(type="pil"),
-    title="Sapiens Body Part Segmentation"
+    title="Sapiens Body Part Segmentation",
+    description="Upload an image to segment body parts using the Sapiens model."
 )
 
-# اجرای برنامه
-interface.launch()
+# Launch the interface
+if __name__ == "__main__":
+    iface.launch()