Update app.py

app.py

@@ -1,62 +1,22 @@
-import torch
-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)
+    # Resize the image to a fixed size (e.g., 1024x768)
     image = image.resize((1024, 768))
 
-    # Convert to RGB (without adding extra dimensions)
+    # Convert to RGB and normalize pixel values
     input_tensor = np.array(image.convert("RGB")) / 255.0
 
-    # Add the batch dimension
-    input_tensor = input_tensor[np.newaxis, :]
-
-    # Convert the NumPy array to a PyTorch tensor
-    input_tensor = torch.from_numpy(input_tensor)
-
-    return input_tensor
-
-def predict(image):
-    try:
-        print("Predict function called")
+    # 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))
 
-        # Preprocess the image to match the expected input shape
-        input_tensor = preprocess_image(image)
+    # Flatten the patches
+    input_tensor = input_tensor.reshape(-1, patch_size * patch_size * 3)
 
-        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}")
-
-        # Post-process the output (if necessary)
-        # ...
-
-        return output  # Return the output tensor directly
-
-    except Exception as e:
-        print(f"Error during prediction: {str(e)}")
-        return None
+    # Add batch dimension
+    input_tensor = input_tensor[np.newaxis, :]
 
-# Define the Gradio interface
-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."
-)
+    # Convert to PyTorch tensor
+    input_tensor = torch.from_numpy(input_tensor)
 
-if __name__ == "__main__":
-    iface.launch(share=True)
+    return input_tensor