Update app.py

app.py

@@ -1,7 +1,19 @@
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+import torch
+from datasets import load_dataset
 
-unmasker = pipeline('fill-mask', model='bert-large-uncased')
-user_input = input("Enter a sentence with a masked word: ")
-results = unmasker(user_input)
+dataset = load_dataset("competitions/aiornot")
+image = dataset["test"][0]["image"]
+
+processor = AutoImageProcessor.from_pretrained("microsoft/resnet-50")
+model = AutoModelForImageClassification.from_pretrained("arnolfokam/ai-generated-image-detector", trust_remote_code=True)
+
+inputs = processor(image, return_tensors="pt")
+
+with torch.no_grad():
+    logits = model(**inputs).logits
+
+# model predicts one of the 2 classes
+predicted_label = logits.argmax(-1).item()
+print(model.config.id2label[predicted_label])
 
-for result in results:
-    print(result['sequence'])