Update app.py

app.py

@@ -1,9 +1,8 @@
 import streamlit as st
 from PIL import Image
 import torch
-from torchvision import transforms, utils
+from torchvision import transforms
 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):
@@ -21,7 +20,7 @@ def preprocess_image(image, mtcnn, device):
         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 = to_pil_image(cropped_faces[0].cpu())
+            processed_image = cropped_faces[0].cpu()
     except Exception as e:
         st.write(f"Exception in face detection: {e}")
         processed_image = image
@@ -40,8 +39,8 @@ 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)
+        # Adjust for your model's output structure
+        probabilities = torch.nn.functional.softmax(outputs, dim=1)
         predicted_class = torch.argmax(probabilities, dim=1)
     return predicted_class, probabilities
 
@@ -61,4 +60,4 @@ if uploaded_file is not None:
     
     st.write(f"Predicted class: {predicted_class.item()}")
     # Display the final processed image
-    st.image(final_image, caption='Processed Image', use_column_width=True)
+    st.image(final_image, caption='Processed Image', use_column_width=True)