Update app.py

app.py

@@ -17,34 +17,39 @@ def load_model_and_mtcnn(model_path):
 def preprocess_image(image, mtcnn, device):
     processed_image = image  # Initialize with the original image
     try:
-        # The return_image parameter of MTCNN's forward method can return the original image along with detected faces, but here we directly pass the image
-        cropped_faces, _ = mtcnn(image, return_image=True)
+        # 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:
             processed_image = cropped_faces[0]  # Use the first detected face
-        # No else clause needed; if no faces detected, processed_image remains the original
     except Exception as e:
         st.write(f"Exception in face detection: {e}")
         processed_image = image
     
     transform = transforms.Compose([
-        transforms.Resize((224, 224)),  
+        transforms.Resize((224, 224)),
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
     ])
-    image_tensor = transform(processed_image).to(device)
+    # Ensure processed_image is a PIL Image for the transformation
+    if not isinstance(processed_image, Image.Image):
+        processed_image_pil = Image.fromarray(processed_image.cpu().numpy().astype('uint8'), 'RGB')
+    else:
+        processed_image_pil = processed_image
+    image_tensor = transform(processed_image_pil).to(device)
     image_tensor = image_tensor.unsqueeze(0)  # Add a batch dimension
-    return image_tensor, processed_image
+    return image_tensor, processed_image_pil
 
 # Function for inference
 def predict(image_tensor, model, device):
     model.eval()
     with torch.no_grad():
         outputs = model(image_tensor)
+        # Adjust for your model's output if it does not have a 'logits' attribute
         probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
         predicted_class = torch.argmax(probabilities, dim=1)
     return predicted_class, probabilities
 
-# Streamlit UI
+# Streamlit UI setup
 st.title("Face Detection and Classification with ViT")
 st.write("Upload an image, and the model will detect faces and classify the image.")
 
@@ -58,6 +63,7 @@ if uploaded_file is not None:
     image_tensor, final_image = preprocess_image(image, mtcnn, device)
     predicted_class, probabilities = predict(image_tensor, model, device)
     
+    # Here, customize the display of predicted_class and probabilities based on your model's specifics
     st.write(f"Predicted class: {predicted_class.item()}")
     # Display the final processed image
     st.image(final_image, caption='Processed Image', use_column_width=True)