ZappY-AI commited on
Commit
f88be9d
1 Parent(s): e342c49

some font changes

Browse files
Files changed (5) hide show
  1. app.py +9 -6
  2. images/cat.jpg +0 -0
  3. images/horse.jpg +0 -0
  4. images/mountain.jpg +0 -0
  5. 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/Non-AI-Generated Image"]
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/horse.jpg"),
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=" \nBased on the paper:'An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale', Alexey et al.\nDataset: 'Fake or Real competition dataset' at https://huggingface.co/datasets/mncai/Fake_or_Real_Competition_Dataset"
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