Update app.py
Browse files
app.py
CHANGED
@@ -4,12 +4,16 @@ from transformers import pipeline, AutoImageProcessor, Swinv2ForImageClassificat
|
|
4 |
from torchvision import transforms
|
5 |
import torch
|
6 |
from PIL import Image
|
|
|
|
|
|
|
|
|
7 |
|
8 |
# Ensure using GPU if available
|
9 |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
10 |
|
11 |
# Load the first model and processor
|
12 |
-
image_processor_1 = AutoImageProcessor.from_pretrained("haywoodsloan/ai-image-detector-deploy")
|
13 |
model_1 = Swinv2ForImageClassification.from_pretrained("haywoodsloan/ai-image-detector-deploy")
|
14 |
model_1 = model_1.to(device)
|
15 |
clf_1 = pipeline(model=model_1, task="image-classification", image_processor=image_processor_1, device=device)
|
@@ -22,6 +26,7 @@ clf_2 = pipeline("image-classification", model=model_2_path)
|
|
22 |
class_names_1 = ['artificial', 'real']
|
23 |
class_names_2 = ['AI Image', 'Real Image'] # Adjust if the second model has different classes
|
24 |
|
|
|
25 |
def predict_image(img, confidence_threshold):
|
26 |
# Ensure the image is a PIL Image
|
27 |
if not isinstance(img, Image.Image):
|
@@ -67,10 +72,10 @@ def predict_image(img, confidence_threshold):
|
|
67 |
result_2[class_name] = 0.0
|
68 |
|
69 |
# Check if either class meets the confidence threshold
|
70 |
-
if result_2['
|
71 |
-
label_2 = f"Label:
|
72 |
-
elif result_2['
|
73 |
-
label_2 = f"Label:
|
74 |
else:
|
75 |
label_2 = "Uncertain Classification"
|
76 |
except Exception as e:
|
@@ -94,4 +99,5 @@ gr.Interface(
|
|
94 |
inputs=[image, confidence_slider],
|
95 |
outputs=label,
|
96 |
title="AI Generated Classification",
|
|
|
97 |
).launch()
|
|
|
4 |
from torchvision import transforms
|
5 |
import torch
|
6 |
from PIL import Image
|
7 |
+
import warnings
|
8 |
+
|
9 |
+
# Suppress warnings
|
10 |
+
warnings.filterwarnings("ignore", category=UserWarning, message="Using a slow image processor as `use_fast` is unset")
|
11 |
|
12 |
# Ensure using GPU if available
|
13 |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
14 |
|
15 |
# Load the first model and processor
|
16 |
+
image_processor_1 = AutoImageProcessor.from_pretrained("haywoodsloan/ai-image-detector-deploy", use_fast=True)
|
17 |
model_1 = Swinv2ForImageClassification.from_pretrained("haywoodsloan/ai-image-detector-deploy")
|
18 |
model_1 = model_1.to(device)
|
19 |
clf_1 = pipeline(model=model_1, task="image-classification", image_processor=image_processor_1, device=device)
|
|
|
26 |
class_names_1 = ['artificial', 'real']
|
27 |
class_names_2 = ['AI Image', 'Real Image'] # Adjust if the second model has different classes
|
28 |
|
29 |
+
|
30 |
def predict_image(img, confidence_threshold):
|
31 |
# Ensure the image is a PIL Image
|
32 |
if not isinstance(img, Image.Image):
|
|
|
72 |
result_2[class_name] = 0.0
|
73 |
|
74 |
# Check if either class meets the confidence threshold
|
75 |
+
if result_2['AI Image'] >= confidence_threshold:
|
76 |
+
label_2 = f"Label: AI Image, Confidence: {result_2['AI Image']:.4f}"
|
77 |
+
elif result_2['Real Image'] >= confidence_threshold:
|
78 |
+
label_2 = f"Label: Real Image, Confidence: {result_2['Real Image']:.4f}"
|
79 |
else:
|
80 |
label_2 = "Uncertain Classification"
|
81 |
except Exception as e:
|
|
|
99 |
inputs=[image, confidence_slider],
|
100 |
outputs=label,
|
101 |
title="AI Generated Classification",
|
102 |
+
queue=True # Enable queuing to handle multiple predictions efficiently
|
103 |
).launch()
|