LPX55 commited on
Commit
e88a32d
·
verified ·
1 Parent(s): a9d7990

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +23 -4
app.py CHANGED
@@ -1,19 +1,25 @@
1
  import gradio as gr
2
  from transformers import pipeline, AutoImageProcessor, Swinv2ForImageClassification
3
  from torchvision import transforms
 
 
 
 
4
 
5
  # Load the model and processor
6
  image_processor = AutoImageProcessor.from_pretrained("haywoodsloan/ai-image-detector-deploy")
7
  model = Swinv2ForImageClassification.from_pretrained("haywoodsloan/ai-image-detector-deploy")
8
- clf = pipeline(model=model, task="image-classification", image_processor=image_processor)
 
9
 
10
  # Define class names
11
  class_names = ['artificial', 'real']
12
 
13
- def predict_image(img):
14
  # Convert the image to a PIL Image and resize it
15
  img = transforms.ToPILImage()(img)
16
  img = transforms.Resize((256, 256))(img)
 
17
 
18
  # Get the prediction
19
  prediction = clf(img)
@@ -26,9 +32,22 @@ def predict_image(img):
26
  if class_name not in result:
27
  result[class_name] = 0.0
28
 
29
- return result
 
 
 
 
 
 
30
 
31
  # Define the Gradio interface
32
  image = gr.Image(label="Image to Analyze", sources=['upload'])
 
33
  label = gr.Label(num_top_classes=2)
34
- gr.Interface(fn=predict_image, inputs=image, outputs=label, title="AI Generated Classification").launch()
 
 
 
 
 
 
 
1
  import gradio as gr
2
  from transformers import pipeline, AutoImageProcessor, Swinv2ForImageClassification
3
  from torchvision import transforms
4
+ import torch
5
+
6
+ # Ensure using GPU if available
7
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
8
 
9
  # Load the model and processor
10
  image_processor = AutoImageProcessor.from_pretrained("haywoodsloan/ai-image-detector-deploy")
11
  model = Swinv2ForImageClassification.from_pretrained("haywoodsloan/ai-image-detector-deploy")
12
+ model = model.to(device)
13
+ clf = pipeline(model=model, task="image-classification", image_processor=image_processor, device=device)
14
 
15
  # Define class names
16
  class_names = ['artificial', 'real']
17
 
18
+ def predict_image(img, confidence_threshold):
19
  # Convert the image to a PIL Image and resize it
20
  img = transforms.ToPILImage()(img)
21
  img = transforms.Resize((256, 256))(img)
22
+ img = transforms.ToTensor()(img).unsqueeze(0).to(device) # Add batch dimension and move to GPU
23
 
24
  # Get the prediction
25
  prediction = clf(img)
 
32
  if class_name not in result:
33
  result[class_name] = 0.0
34
 
35
+ # Check if either class meets the confidence threshold
36
+ if result['artificial'] >= confidence_threshold:
37
+ return f"Label: artificial, Confidence: {result['artificial']:.4f}"
38
+ elif result['real'] >= confidence_threshold:
39
+ return f"Label: real, Confidence: {result['real']:.4f}"
40
+ else:
41
+ return "Uncertain Classification"
42
 
43
  # Define the Gradio interface
44
  image = gr.Image(label="Image to Analyze", sources=['upload'])
45
+ confidence_slider = gr.Slider(0.0, 1.0, value=0.5, step=0.01, label="Confidence Threshold")
46
  label = gr.Label(num_top_classes=2)
47
+
48
+ gr.Interface(
49
+ fn=predict_image,
50
+ inputs=[image, confidence_slider],
51
+ outputs=label,
52
+ title="AI Generated Classification"
53
+ ).launch()