Update app.py

app.py

@@ -24,49 +24,59 @@ def load_and_preprocess_image(image):
     return img.astype(np.float32)
 
 def swap_faces(source_image, target_image):
-    # Load the ONNX model
-    session = ort.InferenceSession(model_file)
+    try:
+        # Load the ONNX model
+        session = ort.InferenceSession(model_file)
+        
+        # Get input names
+        input_names = [input.name for input in session.get_inputs()]
+        
+        # Print input shapes for debugging
+        for input in session.get_inputs():
+            print(f"Input '{input.name}' expects shape: {input.shape}")
+        
+        # Preprocess the images
+        source_img = load_and_preprocess_image(source_image)
+        target_img = load_and_preprocess_image(target_image)
+        
+        # Reshape inputs according to model requirements
+        # For the first input (assuming it's the image input)
+        source_input = source_img.transpose(2, 0, 1)[np.newaxis, ...]  # Shape: (1, 3, 512, 512)
+        
+        # For the second input (onnx::Gemm_1), reshape to rank 2 as required by the error message
+        target_features = target_img.transpose(2, 0, 1).reshape(-1, 512)  # Reshape to 2D array
+        
+        # Create input dictionary
+        input_dict = {
+            input_names[0]: source_input.astype(np.float32),
+            input_names[1]: target_features.astype(np.float32)
+        }
+        
+        # Run inference
+        result = session.run(None, input_dict)[0]
+        
+        # Post-process the result
+        result = result[0].transpose(1, 2, 0)  # Convert from NCHW to HWC format
+        result = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)  # Convert back to RGB
+        return np.clip(result * 255, 0, 255).astype(np.uint8)
     
-    # Preprocess the images
-    source_img = load_and_preprocess_image(source_image)
-    target_img = load_and_preprocess_image(target_image)
-    
-    # Prepare input data for the model
-    source_input = np.expand_dims(source_img.transpose(2, 0, 1), 0)  # Shape: (1, 3, 512, 512)
-    target_input = np.expand_dims(target_img.transpose(2, 0, 1), 0)  # Shape: (1, 3, 512, 512)
-    
-    # Get input names dynamically
-    input_names = [input.name for input in session.get_inputs()]
-    
-    # Create input dictionary with correct shapes
-    input_dict = {
-        input_names[0]: source_input,  # First input for source image
-        input_names[1]: target_input   # Second input for target image
-    }
-    
-    # Run inference
-    result = session.run(None, input_dict)[0]
-    
-    # Post-process the result
-    result = result[0].transpose(1, 2, 0)  # Convert from NCHW to HWC format
-    result = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)  # Convert back to RGB
-    return np.clip(result * 255, 0, 255).astype(np.uint8)
+    except Exception as e:
+        print(f"Error during face swapping: {str(e)}")
+        raise
 
 # Create Gradio interface
 interface = gr.Interface(
     fn=swap_faces,
     inputs=[
-        gr.Image(label="Source Face"),
-        gr.Image(label="Target Image")
+        gr.Image(label="Source Face", type="numpy"),
+        gr.Image(label="Target Image", type="numpy")
     ],
     outputs=gr.Image(label="Result"),
     title="Face Swap using SimSwap",
     description="Upload a source face and a target image to swap faces. The source face will be transferred onto the target image.",
-    allow_flagging="never",
-    examples=[
-        ["path_to_example_source.jpg", "path_to_example_target.jpg"]  # Replace with actual example image paths if available
-    ]
+    allow_flagging="never"
 )
 
-# Launch the interface with share=True for public access
-interface.launch(share=True)
+# Launch the interface
+if __name__ == "__main__":
+    interface.launch()