File size: 3,088 Bytes
de201ad
 
 
 
 
 
 
0e7d3d5
5099c06
e357c23
 
de201ad
 
 
 
 
 
 
e357c23
de201ad
0e7d3d5
a3a33d5
 
 
 
de201ad
0e7d3d5
e357c23
de201ad
0e7d3d5
 
 
 
 
 
de201ad
0e7d3d5
 
f25a954
de201ad
349b8ae
 
0e7d3d5
de201ad
0e7d3d5
3f84f36
0e7d3d5
3f84f36
 
6cf383a
 
349b8ae
6cf383a
81984cb
 
 
 
 
0e7d3d5
 
3f84f36
de201ad
0e7d3d5
 
de201ad
 
a3a33d5
de201ad
0e7d3d5
de201ad
0e7d3d5
de201ad
0e7d3d5
 
de201ad
 
 
 
 
 
 
0653816
 
0e7d3d5
0653816
 
 
cba4e9e
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
import os
import cv2
import requests
import gradio as gr

from ultralytics import YOLO

file_urls = ["https://www.kp-glas.de/wp-content/uploads/2023/02/Bierflaschen.jpeg", 
             "https://www.bevindustry.com/ext/resources/issues/2020/February/Heineken-0-0-Alcohol-Free-Beer.jpg?1581450636",
            "https://spice-world.co.za/cdn/shop/products/3520_648ad722263c64.28810246_drinks_20copy_1db265d5-e818-45d3-973b-3e046a6b0931_800x.jpg?v=1697124785",
            "https://i.redd.it/0e2zml4mwz021.jpg"]

def download_file(url, save_name):
    if not os.path.exists(save_name):
        file = requests.get(url)
        open(save_name, 'wb').write(file.content)

for i, url in enumerate(file_urls):
    download_file(file_urls[i], f"image_{i}.jpg")

# Available model options
model_paths = [
    ("Yolov8n (Accurate)","best.pt"),
    ("Yolov8n half-precision (Faster)","best.onnx")
    ]

# Paths to example images
path = [['image_0.jpg'], ['image_1.jpg'], ['image_2.jpg'],  ['image_3.jpg']]

# Function to run inference and show predictions
def show_preds_image(image_path, selected_model):
    # Load the selected model
    model = YOLO(selected_model)
    
    # Read the image
    image = cv2.imread(image_path)
    
    # Run model prediction
    outputs = model.predict(source=image_path, conf=0.8)
    results = outputs[0].cpu().numpy()

    font_scale = 2
    # Draw bounding boxes and labels
    for i, det in enumerate(results.boxes.xyxy):
        cv2.rectangle(image, (int(det[0]), int(det[1])), (int(det[2]), int(det[3])), color=(0, 0, 255), thickness=2, lineType=cv2.LINE_AA)
        class_id = int(results.boxes.cls[i])
        class_name = model.names[class_id]  # Assuming model.names contains class names
        label = f"{class_name}"
        
        bbox_width = int(det[2] - det[0])
        
        label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, 2)
        
        label_x1 = int(det[0])
        label_y1 = max(int(det[1]), label_size[1] + 10) - label_size[1]
        label_x2 = label_x1 + label_size[0]
        label_y2 = label_y1 + label_size[1] + 5
        
        cv2.rectangle(image, (label_x1, label_y1), (label_x2, label_y2), (255, 255, 255), thickness=-1)
        cv2.putText(image, label, (label_x1, label_y2 - 5), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), 2, lineType=cv2.LINE_AA)
        
    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# Gradio UI components
inputs_image = [
    gr.components.Image(type="filepath", label="Input Image"),
    gr.components.Dropdown(choices=model_paths, label="Select Model", value='best.pt', interactive=True)  # Dropdown for model selection
]

outputs_image = [
    gr.components.Image(type="numpy", label="Output Image")
]

# Gradio Interface
interface_image = gr.Interface(
    fn=show_preds_image,
    inputs=inputs_image,
    outputs=outputs_image,
    title="Beverage Container Detector",
    examples=path,
    cache_examples=False,
)

# Launch Gradio tabbed interface
gr.TabbedInterface(
    [interface_image],
    tab_names=['Image inference']
).queue().launch()