Update app.py

app.py

@@ -1,8 +1,9 @@
 import streamlit as st
 from PIL import Image
 import torch
-from torchvision import transforms
+from torchvision import transforms, utils
 from facenet_pytorch import MTCNN
+from torchvision.transforms.functional import to_pil_image
 
 # Function to load the ViT model and MTCNN
 def load_model_and_mtcnn(model_path):
@@ -15,12 +16,13 @@ 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
-            processed_image = cropped_faces[0].cpu()
+            cropped_image = to_pil_image(cropped_faces[0].cpu())
     except Exception as e:
         st.write(f"Exception in face detection: {e}")
         processed_image = image
@@ -30,17 +32,21 @@ def preprocess_image(image, mtcnn, device):
         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)
-    image_tensor = image_tensor.unsqueeze(0)  # Add a batch dimension
-    return image_tensor, processed_image
+    # 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
 
 # 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 structure
-        probabilities = torch.nn.functional.softmax(outputs, dim=1)
+        # 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