init commit

app.py

@@ -0,0 +1,63 @@
+import streamlit as st
+from PIL import Image
+import torch
+from torchvision import transforms
+from facenet_pytorch import MTCNN
+import matplotlib.pyplot as plt
+
+# Function to load the ViT model and MTCNN
+def load_model_and_mtcnn(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')
+    model.to(device)
+    mtcnn = MTCNN(keep_all=True, device=device)
+    return model, device, mtcnn
+
+# 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
+    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)
+        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.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
+
+# Function for inference
+def predict(image_tensor, model, device):
+    model.eval()
+    with torch.no_grad():
+        outputs = model(image_tensor)
+        probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
+        predicted_class = torch.argmax(probabilities, dim=1)
+    return predicted_class, probabilities
+
+# Streamlit UI
+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 = "path_to_your_ViT_model.pt"  # Make sure to upload your model to Hugging Face Spaces or adjust this path
+model, device, mtcnn = load_model_and_mtcnn(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)
+    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)

model_v1.0.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:49b6bc053a64219bc599a332d053e43a46648d761878e59219924dcf7144d07c
+size 343321862