some font changes
Browse files- app.py +9 -6
- images/cat.jpg +0 -0
- images/horse.jpg +0 -0
- images/mountain.jpg +0 -0
- images/unicorn.jpg +0 -0
app.py
CHANGED
@@ -6,7 +6,7 @@ from torchvision.transforms import v2 as transforms
|
|
6 |
import os
|
7 |
|
8 |
# Define the class names
|
9 |
-
class_names = ['AI-Generated Image', "Real/
|
10 |
|
11 |
# Load the model
|
12 |
weights_path = "FaKe-ViT-B16.pth"
|
@@ -24,7 +24,8 @@ preprocess = transforms.Compose([
|
|
24 |
def predict_image(image):
|
25 |
image = preprocess(image)
|
26 |
if image.shape[0] != 3:
|
27 |
-
image = image[:3, :, :]
|
|
|
28 |
image = image.unsqueeze(0)
|
29 |
with torch.inference_mode():
|
30 |
output = model(image)
|
@@ -40,12 +41,14 @@ demo = gr.Interface(
|
|
40 |
flagging_options=["incorrect prediction"],
|
41 |
examples=[
|
42 |
("images/cheetah.jpg"),
|
43 |
-
( "images/
|
44 |
("images/astronaut.jpg"),
|
|
|
|
|
45 |
],
|
46 |
-
title="FaKe-ViT-B/16: AI-Generated Image Detection using Vision Transformer(ViT-B/16)",
|
47 |
-
description="This is a demo to detect AI-Generated images using Vision Transformer(ViT-B/16). Upload an image and the model will predict whether the image is AI-Generated or Real",
|
48 |
-
article="
|
49 |
)
|
50 |
|
51 |
if __name__ == "__main__":
|
|
|
6 |
import os
|
7 |
|
8 |
# Define the class names
|
9 |
+
class_names = ['Fake/AI-Generated Image', "Real/Not an AI-Generated Image"]
|
10 |
|
11 |
# Load the model
|
12 |
weights_path = "FaKe-ViT-B16.pth"
|
|
|
24 |
def predict_image(image):
|
25 |
image = preprocess(image)
|
26 |
if image.shape[0] != 3:
|
27 |
+
# image = image[:3, :, :]
|
28 |
+
return "Invalid Image: Image should be in RGB format. Please upload a valid image."
|
29 |
image = image.unsqueeze(0)
|
30 |
with torch.inference_mode():
|
31 |
output = model(image)
|
|
|
41 |
flagging_options=["incorrect prediction"],
|
42 |
examples=[
|
43 |
("images/cheetah.jpg"),
|
44 |
+
( "images/cat.jpg"),
|
45 |
("images/astronaut.jpg"),
|
46 |
+
("images/mountain.jpg"),
|
47 |
+
("images/unicorn.jpg")
|
48 |
],
|
49 |
+
title="<u>FaKe-ViT-B/16: Robust and Fast AI-Generated Image Detection using Vision Transformer(ViT-B/16):</u>",
|
50 |
+
description="<p style='font-size: 20px;'>This is a demo to detect AI-Generated images using a fine-tuned Vision Transformer(ViT-B/16). Upload an image and the model will predict whether the image is AI-Generated or Real",
|
51 |
+
article="<p style='font-size: 20px;'><b>Paper</b>: 'An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale', Alexey et al.<br/><b>Dataset</b>: 'Fake or Real competition dataset' at <a href='https://huggingface.co/datasets/mncai/Fake_or_Real_Competition_Dataset'>Fake or Real competition dataset</a>"
|
52 |
)
|
53 |
|
54 |
if __name__ == "__main__":
|
images/cat.jpg
ADDED
images/horse.jpg
DELETED
Binary file (47.8 kB)
|
|
images/mountain.jpg
ADDED
images/unicorn.jpg
ADDED