Update app.py

app.py

@@ -1,96 +1,53 @@
 import streamlit as st
-import torch
-import torchvision.transforms as transforms
-from transformers import ViTFeatureExtractor, ViTForImageClassification
-from PIL import Image
-import requests
-from PIL import Image
 import cv2
 import numpy as np
+from transformers import pipeline
 
-# Download the pretrained ViT model
-feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-large-patch16-384')
-vit_model = ViTForImageClassification.from_pretrained('google/vit-large-patch16-384')
+# Set up the CLIP classifier
+model_name = "openai/clip-vit-large-patch14-336"
+classifier = pipeline("zero-shot-image-classification", model=model_name)
+labels_for_classification = ["genuine face", "morphed face"]
 
-# Download the pretrained CNN model
-cnn_model = torch.hub.load('pytorch/vision:v0.9.0', 'resnet50', pretrained=True)
-
-# Define the image transform
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
-])
-
-# Define the function to predict whether an image is genuine or morphed
-def predict(image):
-    # Convert the numpy array to PIL Image object
-    image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
-
-    # Preprocess the image
-    image = transform(image)
-    image = image.unsqueeze(0)
-
-    # Predict the class using ViT
-    with torch.no_grad():
-        viT_output = vit_model(image)
-        viT_probs = torch.nn.functional.softmax(viT_output, dim=1)
-        viT_score, viT_pred = torch.max(viT_probs, 1)
-
-    # Predict the class using CNN
-    with torch.no_grad():
-        cnn_output = cnn_model(image)
-        cnn_probs = torch.nn.functional.softmax(cnn_output, dim=1)
-        cnn_score, cnn_pred = torch.max(cnn_probs, 1)
-
-    # Combine the predictions using a weighted average
-    combined_score = 0.7 * viT_score.item() + 0.3 * cnn_score.item()
-    combined_pred = viT_pred.item() if viT_score.item() > cnn_score.item() else cnn_pred.item()
-    return combined_pred, combined_score
-
-# Define the function to restore an image
+# Define the restoration function
 def restore(image):
+    # Convert the image to float32
+    image = np.float32(image)
     # Apply a median blur to the image
     image = cv2.medianBlur(image, 5)
     return image
 
-# Define the function to enhance an image
+# Define the enhancement function
 def enhance(image):
-    # Increase the contrast of the image
-    lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
-    l, a, b = cv2.split(lab)
-    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
-    l = clahe.apply(l)
-    lab = cv2.merge((l, a, b))
-    image = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
-    return image
+    # Convert the image to grayscale
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    # Apply histogram equalization to the image
+    equalized = cv2.equalizeHist(gray)
+    return equalized
 
 # Define the Streamlit app
 def app():
-    st.title("Advanced Face Morphing Detection and Restoration")
-
-    # Upload an image
-    uploaded_file = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"])
-
-    # Display the uploaded image and perform predictions, restoration, and enhancement
+    # Create a file uploader
+    uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
     if uploaded_file is not None:
-        image = cv2.imdecode(np.fromstring(uploaded_file.read(), np.uint8), cv2.IMREAD_COLOR)
-
-        # Predict whether the image is genuine or morphed and show the prediction score
-        prediction, score = predict(image)
-        if prediction == 0:
-            st.write("The image is genuine with a score of {:.2f}".format(score))
+        # Read the image file
+        image = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
+        image = cv2.imdecode(image, cv2.IMREAD_COLOR)
+
+        # Classify the image using CLIP
+        scores = classifier(image, candidate_labels=labels_for_classification)
+        if scores[0]['label'] == "genuine face":
+            st.write("The image contains a genuine face")
         else:
-            st.write("The image is morphed with a score of {:.2f}".format(score))
+            st.write("The image contains a morphed face")
 
         # Restore the image
-        restored_image = restore(image)
-        st.image(restored_image, caption="Restored Image", use_column_width=True)
+        restored = restore(image)
+        st.image(restored, caption="Restored Image")
 
         # Enhance the image
-        enhanced_image = enhance(image)
-        st.image(enhanced_image, caption="Enhanced Image", use_column_width=True)
+        enhanced = enhance(restored)
+        st.image(enhanced, caption="Enhanced Image")
 
-# Run the Streamlit app
-if __name__ == '__main__':
-    app()
+# Run the app
+if __name__ == "__main__":
+    app()