feat: add flag

.gitignore

@@ -0,0 +1 @@
+best_model.pth

app.py

@@ -24,11 +24,11 @@ def set_seed(seed):
 
 set_seed(42)
 
+flag = os.environ["FLAG"] if "FLAG" in os.environ else "fakeflag{placeholder}"
 device = "cpu"
 config = ViTConfig.from_pretrained("google/vit-base-patch16-224")
 config.num_labels = 2  # Binary classification
 
-# Download the model file
 model_url = "https://huggingface.co/spuun/yummy-paws/resolve/main/model.pth"
 model_path = "best_model.pth"
 
@@ -37,16 +37,14 @@ if not os.path.exists(model_path):
     with open(model_path, "wb") as f:
         f.write(response.content)
 
-# Load the trained model
 model = ViTForImageClassification.from_pretrained(
     model_path,
     config=config,
-    ignore_mismatched_sizes=True,  # weights_only=False
+    ignore_mismatched_sizes=True,
 )
 model.classifier = nn.Linear(model.config.hidden_size, 2)
 model.to(device)
 
-# Download the reference image
 reference_image_url = (
     "https://huggingface.co/spuun/yummy-paws/resolve/main/images%20(15).jpeg"
 )
@@ -57,7 +55,6 @@ if not os.path.exists(reference_image_path):
     with open(reference_image_path, "wb") as f:
         f.write(response.content)
 
-# Load the reference image for SSIM comparison
 reference_image = Image.open(reference_image_path)
 
 
@@ -92,15 +89,21 @@ def predict_and_compare(image):
         predicted_class = class_names[predicted_class_index]
         probability = probabilities[predicted_class_index].item()
 
-    return f"Predicted: {predicted_class}\nProbability: {probability:.4f}\nSSIM with reference: {ssim_value:.4f}"
+    return f"""{'SUCCESSFULLY AUTHENTICATED!!\nFLAG: '+flag if ssim_value>0.9 and predicted_class == "True" else "FAILED TO AUTHENTICATE :("}
+=====================
+
+Predicted: {predicted_class}
+Probability: {probability:.4f}
+SSIM with reference: {ssim_value:.4f}"""
 
 
 iface = gr.Interface(
     fn=predict_and_compare,
     inputs=gr.Image(type="pil"),
     outputs="text",
-    title="Image Classification and Comparison",
-    description="Upload an image to classify it and compare with a reference image.",
+    title="Image authentication",
+    description="Upload your image here to be authenticated!",
+    allow_flagging="never",
 )
 
 iface.launch()