Update app.py

app.py

@@ -4,52 +4,65 @@ from transformers import AutoFeatureExtractor, AutoModelForImageClassification,
 
 models=[
     "Nahrawy/AIorNot",
-    "RishiDarkDevil/ai-image-det-resnet152",
     "arnolfokam/ai-generated-image-detector",
     "umm-maybe/AI-image-detector",
 ]
-#pipe = pipeline("image-classification", "umm-maybe/AI-image-detector")
 
-def image_classifier(image):
-    outputs = pipe(image)
-    results = {}
-    for result in outputs:
-        results[result['label']] = result['score']
-    return results
-
-
-
-#demo = gr.Interface(fn=image_classifier, inputs=gr.Image(type="pil"), outputs="label", title=title, description=description)
-#demo.launch(show_api=False)
-
-
-
-def aiornot(image,mod_choose):    
+def aiornot0(image):    
     labels = ["Real", "AI"]
-    
-    #feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/resnet-50")
-    mod=models[int(mod_choose)]
+    mod=models[0]
     feature_extractor = AutoFeatureExtractor.from_pretrained(mod)
     model = AutoModelForImageClassification.from_pretrained(mod)
-    
     input = feature_extractor(image, return_tensors="pt")
     with torch.no_grad():
       outputs = model(**input)
-      print(outputs)
-      print(dir(outputs))
+      print (outputs)
       logits = outputs.logits
       print (logits)
     prediction = logits.argmax(-1).item()
-    print(prediction)
     label = labels[prediction] 
     return label
-
+def aiornot1(image):    
+    labels = ["Real", "AI"]
+    mod=models[1]
+    feature_extractor = AutoFeatureExtractor.from_pretrained(mod)
+    model = AutoModelForImageClassification.from_pretrained(mod)
+    input = feature_extractor(image, return_tensors="pt")
+    with torch.no_grad():
+      outputs = model(**input)
+      print (outputs)
+      logits = outputs.logits
+      print (logits)
+    prediction = logits.argmax(-1).item()
+    label = labels[prediction] 
+    return label
+def aiornot2(image):    
+    labels = ["Real", "AI"]
+    mod=models[2]
+    feature_extractor = AutoFeatureExtractor.from_pretrained(mod)
+    model = AutoModelForImageClassification.from_pretrained(mod)
+    input = feature_extractor(image, return_tensors="pt")
+    with torch.no_grad():
+      outputs = model(**input)
+      print (outputs)
+      logits = outputs.logits
+      print (logits)
+    prediction = logits.argmax(-1).item()
+    label = labels[prediction] 
+    return label    
 with gr.Blocks() as app:
     with gr.Row():
         with gr.Column():
             inp = gr.Image()
             mod_choose=gr.Number(value=0)
             btn = gr.Button()
-    outp = gr.Textbox()
-    btn.click(aiornot,[inp,mod_choose],outp)
+            
+        with gr.Column():
+            outp0 = gr.Textbox()
+            outp1 = gr.Textbox()
+            outp2 = gr.Textbox()
+    btn.click(aiornot0,[inp],outp0)
+    btn.click(aiornot1,[inp],outp1)
+    btn.click(aiornot2,[inp],outp2)
+    
 app.launch()