some changes

app.py

@@ -5,12 +5,12 @@ import torchvision.models as models
 from torchvision.transforms import v2 as transforms
 import os
 
+# Define the class names
 class_names = ['AI-Generated Image', "Real/Non-AI-Generated Image"]
 
-# Downloading the model
-# model = models.vit_b_16()
+# Load the model
 weights_path = "FaKe-ViT-B16.pth"
-model = torch.load(weights_path).to("cpu")
+model = torch.load(weights_path, map_location=torch.device('cpu'))
 model.eval()
 # Preprocessing the image
 preprocess = transforms.Compose([
@@ -22,7 +22,6 @@ preprocess = transforms.Compose([
 
 # Define the prediction function
 def predict_image(image):
-    # inp = Image.fromarray(inp.astype('uint8'), 'RGB')
     image = preprocess(image)
     if image.shape[0] != 3:
         image = image[:3, :, :]
@@ -32,8 +31,6 @@ def predict_image(image):
     output1 = torch.argmax(torch.softmax(output,dim=1),dim=1).item()
     return class_names[output1]
 
-# def image_mod(image):
-#     return image.rotate(45)
 
 
 demo = gr.Interface(
@@ -41,17 +38,13 @@ demo = gr.Interface(
     gr.Image(image_mode="RGB",type="pil"),
     "text",
     flagging_options=["incorrect prediction"],
-    # examples=[
-    #     os.path.join(os.path.dirname(__file__), "images/cheetah1.jpg"),
-    #     os.path.join(os.path.dirname(__file__), "images/lion.jpg"),
-    #     os.path.join(os.path.dirname(__file__), "images/logo.png"),
-    #     os.path.join(os.path.dirname(__file__), "images/tower.jpg"),
-    # ],
+     examples=[
+         os.path.join(os.path.dirname(__file__), "images/cheetah.jpg"),
+         os.path.join(os.path.dirname(__file__), "images/horse.jpg"),
+         os.path.join(os.path.dirname(__file__), "images/astronaut.png"),
+     ],
     title="FaKe-ViT-B/16: AI-Generated Image Detection using Vision Transformer(ViT-B/16)",
     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",
-    css=""".gr-header, .gr-text {
-        font-size: 20px;
-    }""",
     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"
 )