app.py

"""import gradio as gr

def greet(name):
    return "Hello " + name + "!!"

demo = gr.Interface(fn=greet, inputs="text", outputs="text")
demo.launch()
"""


import PIL.Image as Image
import gradio as gr
import webbrowser
import threading
import time
from ultralytics import ASSETS, YOLO

# Загрузим модель из файла yolov8m-worldv2
model = YOLO("yolov8m-worldv2.pt")

# Cписок транспорта
transport_russian = ["Автомобиль", "Грузовик", "Автобус", "Мотоцикл"]

# Создание словаря с соответствием русских и английских слов
translation_dict = {
    "Автомобиль": "car",
    "Грузовик": "truck",
    "Автобус": "bus",
    "Мотоцикл": "motorcycle",
}


def predict_image(img, conf_threshold, iou_threshold, class_flags):
    """
    Функция для предсказания изображения с учетом выбранных транспортных средств
    :param img: изображение
    :param conf_threshold: порог достоверности для обнаружения объектов
    :param iou_threshold: порог пересечения по объединению для обнаружения объектов
    :param class_flags: флаги классов объектов для обнаружения
    :return: изображение с распознанными объектами
    """
    # Получаем английские названия классов транспорта
    transport_classes = [translation_dict[transport] for transport in class_flags]

    # Задаем пользовательские классы для модели
    model.set_classes(transport_classes)

    # Предсказываем объекты на изображении
    results = model(img, conf=conf_threshold, iou=iou_threshold)

    # Отображаем изображение с распознанными объектами
    for r in results:
        im_array = r.plot()
        im = Image.fromarray(im_array[..., ::-1])  # Преобразование массива в изображение PIL

    return im

# Создаем интерфейс Gradio
iface = gr.Interface(
    fn=predict_image,
    inputs=[
        gr.Image(type="pil", label="Загрузка изображения"),  # Ввод изображения
        gr.Slider(minimum=0, maximum=1, value=0.25, label="Порог достоверности"),  # Порог достоверности
        gr.Slider(minimum=0, maximum=1, value=0.45, label="Порог пересечения по объединению"),  # Порог пересечения по объединению
        gr.CheckboxGroup(transport_russian, label="Транспорт", info="Какой транспорт распознать?"),  # Флаги классов объектов
    ],
    outputs=gr.Image(type="pil", label="Изображение с распознанными объектами"),  # Вывод изображения с распознанными объектами
    title="Ultralytics Gradio",  # Заголовок интерфейса
    description="Загрузите изображения для вывода. Используется модедль Ultralytics YOLOv8m-worldv2."  # Описание интерфейса
)

if __name__ == "__main__":
    iface.launch()