app.py

"""import gradio as gr

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

demo = gr.Interface(fn=greet, inputs="text", outputs="text")
demo.launch()
"""
'''import gradio as gr


def sentence_builder(quantity, animal, countries, place, activity_list, morning):
    return f"""The {quantity} {animal}s from {" and ".join(countries)} went to the {place} where they {" and ".join(activity_list)} until the {"morning" if morning else "night"}"""


demo = gr.Interface(
    sentence_builder,
    [
        gr.Slider(2, 20, value=4, label="Count", info="Choose between 2 and 20"),
        gr.Dropdown(
            ["cat", "dog", "bird"], label="Animal", info="Will add more animals later!"
        ),
        gr.CheckboxGroup(["USA", "Japan", "Pakistan"], label="Countries", info="Where are they from?"),
        gr.Radio(["park", "zoo", "road"], label="Location", info="Where did they go?"),
        gr.Dropdown(
            ["ran", "swam", "ate", "slept"], value=["swam", "slept"], multiselect=True, label="Activity", info="Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed auctor, nisl eget ultricies aliquam, nunc nisl aliquet nunc, eget aliquam nisl nunc vel nisl."
        ),
        gr.Checkbox(label="Morning", info="Did they do it in the morning?"),
    ],
    "text",
    examples=[
        [2, "cat", ["Japan", "Pakistan"], "park", ["ate", "swam"], True],
        [4, "dog", ["Japan"], "zoo", ["ate", "swam"], False],
        [10, "bird", ["USA", "Pakistan"], "road", ["ran"], False],
        [8, "cat", ["Pakistan"], "zoo", ["ate"], True],
    ]
)

if __name__ == "__main__":
    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()