Update app.py

app.py

@@ -6,7 +6,7 @@ def greet(name):
 demo = gr.Interface(fn=greet, inputs="text", outputs="text")
 demo.launch()
 """
-import gradio as gr
+'''import gradio as gr
 
 
 def sentence_builder(quantity, animal, countries, place, activity_list, morning):
@@ -38,3 +38,68 @@ demo = gr.Interface(
 
 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()