Update app.py

app.py

@@ -1,14 +1,12 @@
 import streamlit as st
 from PIL import Image
 import torch
-from torchvision import transforms, utils
-from facenet_pytorch import MTCNN
-from torchvision.transforms.functional import to_pil_image
+from torchvision import transforms
 import cv2
 import numpy as np
-import io
+from facenet_pytorch import MTCNN
 
-# Function to load the ViT model and MTCNN
+# Function to load the ViT model
 def load_model(model_path):
     model = torch.load(model_path, map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -22,7 +20,7 @@ mtcnn = MTCNN(keep_all=True, device=device)
 # Function to preprocess the image using MTCNN for face detection
 def preprocess_image(image, device):
     # Convert PIL image to OpenCV format
-    open_cv_image = np.array(image)
+    open_cv_image = np.array(image.convert("RGB"))
     # Convert RGB to BGR for OpenCV
     open_cv_image = cv2.cvtColor(open_cv_image, cv2.COLOR_RGB2BGR)
 
@@ -55,8 +53,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 specifics
+        probabilities = torch.nn.functional.softmax(outputs, dim=1)
         predicted_class = torch.argmax(probabilities, dim=1)
     return predicted_class, probabilities
 
@@ -65,15 +63,17 @@ st.title("Face Detection and Classification with ViT")
 st.write("Upload an image, and the model will detect faces and classify the image.")
 
 model_path = "model_v1.0.pt"  # Adjust this path as necessary
-model, device, mtcnn = load_model_and_mtcnn(model_path)
+model, device = load_model(model_path)
 
 uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 if uploaded_file is not None:
     image = Image.open(uploaded_file).convert("RGB")
     st.image(image, caption='Uploaded Image', use_column_width=True)
-    image_tensor, final_image = preprocess_image(image, mtcnn, device)
+
+    # Preprocess the image and perform inference
+    image_tensor, final_image = preprocess_image(image, device)
     predicted_class, probabilities = predict(image_tensor, model, device)
     
     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)