Spaces:
Sleeping
Sleeping
| import os | |
| import cv2 | |
| import gradio as gr | |
| from ultralytics import YOLO | |
| import numpy as np | |
| model_options = ["yolov8n_bird_100.pt", "yolov8s_bird_100.pt", "yolov8m_bird_100.pt", "yolov8l_bird_100.pt"] | |
| model_names = ["Nano", "Small", "Medium", "Large"] | |
| print("before loading examples") | |
| example_list = [["./examples/" + example] for example in os.listdir("examples")] | |
| print("before loading models") | |
| models = [YOLO(os.path.join("./saved_model", option)) for option in model_options] | |
| print("finish preparation") | |
| def process_image(input_image, model_name, conf): | |
| print('start processing: ') | |
| if input_image is None: | |
| return None, "No objects detected." | |
| print(f"model_name : {model_name}") | |
| print(f"conf : {conf}") | |
| if model_name is None: | |
| model_name = model_names[0] | |
| if conf is None: | |
| conf = 0.6 | |
| model_index = model_names.index(model_name) | |
| print('model_index: ') | |
| print(model_index) | |
| model = models[model_index] | |
| input_image = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR) | |
| # cv2.imwrite('./input_image.jpg', input_image) | |
| results = model.predict(input_image, conf=conf) | |
| class_counts = {} | |
| class_counts_str = "Class Counts:\n" | |
| for r in results: | |
| im_array = r.plot() | |
| im_array = im_array.astype(np.uint8) | |
| im_array = cv2.cvtColor(im_array, cv2.COLOR_BGR2RGB) | |
| if len(r.boxes) == 0: # If no objects are detected | |
| return None, "No objects detected." | |
| for box in r.boxes: | |
| class_name = r.names[box.cls[0].item()] | |
| class_counts[class_name] = class_counts.get(class_name, 0) + 1 | |
| for cls, count in class_counts.items(): | |
| class_counts_str += f"\n{cls}: {count}" | |
| return im_array, class_counts_str | |
| iface = gr.Interface( | |
| fn=process_image, | |
| inputs=[ | |
| gr.Image(), | |
| gr.Radio(model_names, label="Choose model", value=model_names[0]), | |
| gr.Slider(minimum=0.2, maximum=1.0, step=0.1, label="Confidence Threshold", value=0.6) | |
| ], | |
| outputs=["image", gr.Textbox(label="More info")], | |
| title="YOLO Object detection. Trained on bird_1_3_dataset (millitary vehicles).", | |
| description='''The bird_1_3_dataset dataset is composed of millitary vehicles.\n | |
| Lost names of objects: BMP-BMD, BMPBMD, BTR, KAMAZ, MTLB, RLS, RSZO, SAU, TANK, URAL, ZRK. | |
| https://universe.roboflow.com/arnold-gudima-t7mdj/bird_1_3 ''', | |
| live=True, | |
| examples=example_list | |
| ) | |
| iface.launch() | |