Update app.py

app.py

@@ -3,48 +3,53 @@ import gradio as gr
 from PIL import Image
 import numpy as np
 
-# بارگذاری مدل
+# Load the model
 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 a function to preprocess images to match the expected input shape
+def preprocess_image(image):
+    # Resize the image to a fixed size (e.g., 224x224)
+    image = image.resize((224, 224))
+
+    # Convert to RGB and normalize pixel values
+    image = image.convert("RGB")
+    input_tensor = np.array(image) / 255.0
+
+    # Flatten the image into a 1D vector
+    input_tensor = input_tensor.reshape(-1)
+
+    # Add the batch dimension
+    input_tensor = input_tensor[np.newaxis, :]
+
+    return input_tensor
+
 def predict(image):
     try:
         print("Predict function called")
 
-        # تغییر اندازه تصویر به 224x224
-        image = image.resize((224, 224))  # تغییر اندازه به 224x224
-
-        # پیش‌پردازش تصویر
-        image = image.convert("RGB")
-        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()
+        # Preprocess the image to match the expected input shape
+        input_tensor = preprocess_image(image)
 
         print(f"Input tensor shape: {input_tensor.shape}")
 
-        # اجرای مدل
+        # Run the model
         with torch.no_grad():
             output = model(input_tensor)
 
         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)
+        # Post-process the output (if necessary)
+        # ...
 
-        print("Output image generated successfully")
-        return output_image
+        return output  # Return the output tensor directly
 
     except Exception as e:
         print(f"Error during prediction: {str(e)}")
         return None
 
-# تعریف رابط Gradio
+# Define the Gradio interface
 iface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil", label="Input Image"),
@@ -54,4 +59,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch(share=True)
+    iface.launch(share=True)