LPX55 commited on
Commit
b1387d5
·
verified ·
1 Parent(s): 4639d7c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -5
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['artificial'] >= confidence_threshold:
71
- label_2 = f"Label: artificial, Confidence: {result_2['artificial']:.4f}"
72
- elif result_2['real'] >= confidence_threshold:
73
- label_2 = f"Label: real, Confidence: {result_2['real']:.4f}"
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()