Update app.py

app.py

@@ -3,53 +3,57 @@ import gradio as gr
 from PIL import Image
 import numpy as np
 
-# Load the model from the local file
+# بارگذاری مدل
 model_path = "sapiens_0.3b_render_people_epoch_100_torchscript.pt2"
 model = torch.jit.load(model_path, map_location=torch.device('cpu'))
+model.eval()
 
-# Define the prediction function
 def predict(image):
     try:
         print("Predict function called")
-
-        # تغییر اندازه تصویر به اندازه‌ای که مدل نیاز دارد
-        target_size = (3072, 3072)  # یا هر اندازه‌ای که مدل نیاز دارد
+        
+        # تعیین اندازه ورودی مدل
+        target_size = (512, 512)  # این اندازه را می‌توانید تغییر دهید
         image = image.resize(target_size)
-
-        # Preprocess image
+        
+        # پیش‌پردازش تصویر
         image = image.convert("RGB")
-        input_tensor = torch.from_numpy(np.array(image)).permute(2, 0, 1).unsqueeze(0).float() / 255.0
-
-        print("Image preprocessed")
-
-        # Run the model
+        input_tensor = np.array(image)
+        input_tensor = input_tensor.transpose(2, 0, 1)  # تبدیل از HWC به CHW
+        input_tensor = input_tensor[np.newaxis, :]  # افزودن بعد batch
+        input_tensor = input_tensor / 255.0  # نرمال‌سازی
+        input_tensor = torch.from_numpy(input_tensor).float()
+        
+        print(f"Input tensor shape: {input_tensor.shape}")
+        
+        # اجرای مدل
         with torch.no_grad():
             output = model(input_tensor)
         
-        print("Model executed")
-
-        # Postprocess the output
-        output_image = output.squeeze().permute(1, 2, 0).numpy()
+        print(f"Output tensor shape: {output.shape}")
+        
+        # پس‌پردازش خروجی
+        output_image = output.squeeze().cpu().numpy()
+        output_image = output_image.transpose(1, 2, 0)  # تبدیل از CHW به HWC
         output_image = (output_image * 255).astype(np.uint8)
+        output_image = Image.fromarray(output_image)
         
-        print("Output generated")
-        return Image.fromarray(output_image)
+        print("Output image generated successfully")
+        
+        return output_image
 
     except Exception as e:
         print(f"Error during prediction: {str(e)}")
         return None
 
-
-
-# Gradio Interface
+# تعریف رابط Gradio
 iface = gr.Interface(
-    fn=predict, 
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="Sapiens Body Part Segmentation",
-    description="Upload an image to segment body parts using the Sapiens model."
+    fn=predict,
+    inputs=gr.Image(type="pil", label="Input Image"),
+    outputs=gr.Image(type="pil", label="Output Image"),
+    title="Sapiens Model Inference",
+    description="Upload an image to process with the Sapiens model."
 )
 
-# Launch the interface
 if __name__ == "__main__":
     iface.launch()