Spaces:
Runtime error
Runtime error
| 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() | |