Hayloo9838 commited on
Commit
a3d0c64
·
verified ·
1 Parent(s): 0a4514f

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +56 -0
app.py ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import cv2
3
+ import torch
4
+ import numpy as np
5
+ from transformers import CLIPProcessor, CLIPVisionModel
6
+ from PIL import Image
7
+ from torch import nn
8
+ import requests
9
+ import matplotlib.pyplot as plt
10
+ from huggingface_hub import hf_hub_download
11
+
12
+ # ... (rest of your code remains the same)
13
+
14
+ def process_image_classification(image):
15
+ model, processor, reverse_mapping, device = load_model()
16
+
17
+ # Convert image to PIL Image
18
+ image = Image.fromarray(image)
19
+
20
+ inputs = processor(images=image, return_tensors="pt")
21
+ pixel_values = inputs.pixel_values.to(device)
22
+
23
+ with torch.no_grad():
24
+ logits, attentions = model(pixel_values, output_attentions=True)
25
+ probs = torch.nn.functional.softmax(logits, dim=-1)
26
+ prediction = torch.argmax(probs).item()
27
+
28
+ # Generate attention map
29
+ attention_map = get_attention_map(attentions)
30
+
31
+ visualization = apply_heatmap(image, attention_map)
32
+
33
+ card_name = reverse_mapping[prediction]
34
+ confidence = probs[0][prediction].item()
35
+
36
+ # Convert back to RGB for matplotlib display
37
+ visualization_rgb = cv2.cvtColor(visualization, cv2.COLOR_BGR2RGB)
38
+
39
+ return visualization_rgb, card_name, confidence
40
+
41
+ def gradio_interface():
42
+ gr_interface = gr.Interface(
43
+ fn=process_image_classification,
44
+ inputs=gr.inputs.Image(type="numpy"),
45
+ outputs=[
46
+ gr.outputs.Image(label="Heatmap Plot"),
47
+ gr.outputs.Textbox(label="Predicted Card"),
48
+ gr.outputs.Textbox(label="Confidence")
49
+ ],
50
+ title="Uno Card Recognizer",
51
+ description="Upload an image or use your webcam to recognize an Uno card."
52
+ )
53
+ gr_interface.launch()
54
+
55
+ if __name__ == "__main__":
56
+ gradio_interface()