Update app.py

app.py

@@ -1,22 +1,57 @@
-def preprocess_image(image):
-    # Resize the image to a fixed size (e.g., 1024x768)
-    image = image.resize((1024, 768))
+import torch
+import gradio as gr
+from PIL import Image
+import numpy as np
 
-    # Convert to RGB and normalize pixel values
-    input_tensor = np.array(image.convert("RGB")) / 255.0
+# بارگذاری مدل
+model_path = "sapiens_0.3b_render_people_epoch_100_torchscript.pt2"
+model = torch.jit.load(model_path, map_location=torch.device('cpu'))
+model.eval()
 
-    # Divide the image into patches (adjust patch size as needed)
-    patch_size = 16  # Assuming a patch size of 16 based on model information
-    num_patches = (1024 // patch_size) * (768 // patch_size)
-    input_tensor = input_tensor.reshape((num_patches, patch_size, patch_size, 3))
+def predict(image):
+    try:
+        print("Predict function called")
 
-    # Flatten the patches
-    input_tensor = input_tensor.reshape(-1, patch_size * patch_size * 3)
+        # تغییر اندازه تصویر به 224x224
+        image = image.resize((224, 224))  # تغییر اندازه به 224x224
 
-    # Add batch dimension
-    input_tensor = input_tensor[np.newaxis, :]
+        # پیش‌پردازش تصویر
+        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()
 
-    # Convert to PyTorch tensor
-    input_tensor = torch.from_numpy(input_tensor)
+        print(f"Input tensor shape: {input_tensor.shape}")
 
-    return input_tensor
+        # اجرای مدل
+        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)
+
+        print("Output image generated successfully")
+        return output_image
+
+    except Exception as e:
+        print(f"Error during prediction: {str(e)}")
+        return None
+
+# تعریف رابط Gradio
+iface = gr.Interface(
+    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."
+)
+
+if __name__ == "__main__":
+    iface.launch(share=True)