AbdullahMd12 commited on
Commit
0e2ca4b
1 Parent(s): bf23a2d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +95 -2
app.py CHANGED
@@ -1,3 +1,96 @@
1
- import gradio as gr
 
 
 
 
 
 
 
 
2
 
3
- gr.Interface.load("models/openai/clip-vit-large-patch14").launch()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import torch
3
+ import torchvision.transforms as transforms
4
+ from transformers import ViTFeatureExtractor, ViTForImageClassification
5
+ from PIL import Image
6
+ import requests
7
+ from PIL import Image
8
+ import cv2
9
+ import numpy as np
10
 
11
+ # Download the pretrained ViT model
12
+ feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-large-patch16-384')
13
+ vit_model = ViTForImageClassification.from_pretrained('google/vit-large-patch16-384')
14
+
15
+ # Download the pretrained CNN model
16
+ cnn_model = torch.hub.load('pytorch/vision:v0.9.0', 'resnet50', pretrained=True)
17
+
18
+ # Define the image transform
19
+ transform = transforms.Compose([
20
+ transforms.Resize((224, 224)),
21
+ transforms.ToTensor(),
22
+ transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
23
+ ])
24
+
25
+ # Define the function to predict whether an image is genuine or morphed
26
+ def predict(image):
27
+ # Convert the numpy array to PIL Image object
28
+ image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
29
+
30
+ # Preprocess the image
31
+ image = transform(image)
32
+ image = image.unsqueeze(0)
33
+
34
+ # Predict the class using ViT
35
+ with torch.no_grad():
36
+ viT_output = vit_model(image)
37
+ viT_probs = torch.nn.functional.softmax(viT_output, dim=1)
38
+ viT_score, viT_pred = torch.max(viT_probs, 1)
39
+
40
+ # Predict the class using CNN
41
+ with torch.no_grad():
42
+ cnn_output = cnn_model(image)
43
+ cnn_probs = torch.nn.functional.softmax(cnn_output, dim=1)
44
+ cnn_score, cnn_pred = torch.max(cnn_probs, 1)
45
+
46
+ # Combine the predictions using a weighted average
47
+ combined_score = 0.7 * viT_score.item() + 0.3 * cnn_score.item()
48
+ combined_pred = viT_pred.item() if viT_score.item() > cnn_score.item() else cnn_pred.item()
49
+ return combined_pred, combined_score
50
+
51
+ # Define the function to restore an image
52
+ def restore(image):
53
+ # Apply a median blur to the image
54
+ image = cv2.medianBlur(image, 5)
55
+ return image
56
+
57
+ # Define the function to enhance an image
58
+ def enhance(image):
59
+ # Increase the contrast of the image
60
+ lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
61
+ l, a, b = cv2.split(lab)
62
+ clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
63
+ l = clahe.apply(l)
64
+ lab = cv2.merge((l, a, b))
65
+ image = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
66
+ return image
67
+
68
+ # Define the Streamlit app
69
+ def app():
70
+ st.title("Advanced Face Morphing Detection and Restoration")
71
+
72
+ # Upload an image
73
+ uploaded_file = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"])
74
+
75
+ # Display the uploaded image and perform predictions, restoration, and enhancement
76
+ if uploaded_file is not None:
77
+ image = cv2.imdecode(np.fromstring(uploaded_file.read(), np.uint8), cv2.IMREAD_COLOR)
78
+
79
+ # Predict whether the image is genuine or morphed and show the prediction score
80
+ prediction, score = predict(image)
81
+ if prediction == 0:
82
+ st.write("The image is genuine with a score of {:.2f}".format(score))
83
+ else:
84
+ st.write("The image is morphed with a score of {:.2f}".format(score))
85
+
86
+ # Restore the image
87
+ restored_image = restore(image)
88
+ st.image(restored_image, caption="Restored Image", use_column_width=True)
89
+
90
+ # Enhance the image
91
+ enhanced_image = enhance(image)
92
+ st.image(enhanced_image, caption="Enhanced Image", use_column_width=True)
93
+
94
+ # Run the Streamlit app
95
+ if __name__ == '__main__':
96
+ app()