Update app.py

app.py

@@ -0,0 +1,50 @@
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+import gradio as gr
+from timm import create_model
+import torch.nn as nn
+import os
+
+class VisionTransformer(nn.Module):
+    def __init__(self, num_classes, model_name):
+        super(VisionTransformer, self).__init__()
+        self.model = create_model(model_name, pretrained=False, num_classes=num_classes)
+        self.model.head = nn.Sequential(
+            nn.Linear(self.model.num_features, 512),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.Linear(512, num_classes)
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+model_path = "./models/vit_small_patch16_224_final.pth"
+device = torch.device("cpu")
+
+model = VisionTransformer(num_classes=2, model_name="vit_small_patch16_224")
+model.load_state_dict(torch.load(model_path, map_location=device))
+model.eval()
+
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.5]*3, std=[0.5]*3)
+])
+
+def classify_image(img: Image.Image):
+    img_tensor = transform(img).unsqueeze(0)
+    with torch.no_grad():
+        outputs = model(img_tensor)
+        _, predicted = torch.max(outputs, 1)
+        label = 'Fake' if predicted.item() == 0 else 'Real'
+    return label
+
+gr.Interface(
+    fn=classify_image,
+    inputs=gr.Image(type="pil"),
+    outputs="label",
+    title="Luxury Item Authenticity Detector",
+    description="Upload an image to check if it's a real or fake item."
+).launch()