Update app.py

app.py

@@ -16,13 +16,12 @@ def load_model_and_mtcnn(model_path):
 # Function to preprocess the image and return both the tensor and the final PIL image for display
 def preprocess_image(image, mtcnn, device):
     processed_image = image  # Initialize with the original image
-    cropped_image = None
     try:
         # Directly call mtcnn with the image to get cropped faces
         cropped_faces = mtcnn(image)
         if cropped_faces is not None and len(cropped_faces) > 0:
             # Convert the first detected face tensor back to PIL Image for further processing
-            cropped_image = to_pil_image(cropped_faces[0].cpu())
+            processed_image = to_pil_image(cropped_faces[0].cpu(),mode='RGB')
     except Exception as e:
         st.write(f"Exception in face detection: {e}")
         processed_image = image
@@ -32,13 +31,9 @@ def preprocess_image(image, mtcnn, device):
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
     ])
-    # Apply the transformation to the cropped image if available
-    if cropped_image is not None:
-        processed_image = transform(cropped_image).to(device)
-        # Add a batch dimension
-        processed_image = processed_image.unsqueeze(0)  
-
-    return processed_image, cropped_image
+    image_tensor = transform(processed_image).to(device)
+    image_tensor = image_tensor.unsqueeze(0)  # Add a batch dimension
+    return image_tensor, processed_image
 
 # Function for inference
 def predict(image_tensor, model, device):