Update app.py

app.py

@@ -3,24 +3,58 @@ from transformers import DetrImageProcessor, DetrForObjectDetection
 import torch
 from PIL import Image
 
+# Загрузка модели и процессора (кешируется при первом запуске)
 processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
 model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
 
 def detect_objects(image):
+    # Преобразуем входное изображение
+    if isinstance(image, str):  # если путь к файлу
+        image = Image.open(image)
+    
+    # Детекция объектов
     inputs = processor(images=image, return_tensors="pt")
-    outputs = model(**inputs)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    
+    # Постобработка результатов
     target_sizes = torch.tensor([image.size[::-1]])
-    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.7)[0]
+    results = processor.post_process_object_detection(
+        outputs, 
+        target_sizes=target_sizes, 
+        threshold=0.7
+    )[0]
+    
+    # Форматирование результатов для AnnotatedImage
+    annotations = []
+    for box, label, score in zip(results["boxes"], results["labels"], results["scores"]):
+        box = [round(i, 2) for i in box.tolist()]  # округляем координаты
+        label_text = f"{model.config.id2label[label.item()]} ({round(score.item(), 2)})"
+        annotations.append((box, label_text))
     
-    # Подготовка данных для AnnotatedImage
-    boxes = results["boxes"].tolist()
-    labels = [model.config.id2label[label.item()] for label in results["labels"]]
-    return (image, boxes, labels)  # формат: (изображение, bbox, метки)
+    # Проверка на наличие драк (если обнаружено >1 человека)
+    people_count = sum(1 for label in results["labels"] if label.item() == 1)
+    if people_count >= 2:
+        annotations.append(([0, 0, 100, 30], "⚠️ Potential fight!"))
+    
+    return (image, annotations)
 
+# Создание интерфейса
 demo = gr.Interface(
     fn=detect_objects,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.AnnotatedImage(),  # Показывает bbox на изображении
-    title="Детектор драк"
+    inputs=gr.Image(type="pil", label="Input Image"),
+    outputs=gr.AnnotatedImage(
+        label="Detection Results",
+        show_legend=True
+    ),
+    title="Fight Detection with DETR",
+    description="Upload an image to detect people and potential fights. Model: facebook/detr-resnet-50",
+    examples=[
+        ["example1.jpg"],  # добавьте свои примеры во вкладке Files
+        ["example2.jpg"]
+    ],
+    allow_flagging="never"
 )
-demo.launch()
+
+# Запуск приложения
+demo.launch(debug=True)