back/main.py - main - 906009/cod

back/main.py
115 строк5 KB
906009
update back/main.py
22 сен 2024, 10:10
22 сен 2024, 10:1047d371b
import cv2
import torch
from datetime import datetime
from ultralytics import YOLO
import math

model = YOLO('best.pt')

def detect_objects(frame):
    results = model(frame)
    return results

def log_violation(violation_id, start_time, current_time, violation_type):
    elapsed_time = current_time - start_time
    start_seconds = elapsed_time.total_seconds()
    end_seconds = (current_time - log_start_time).total_seconds()
    if violation_type != "":
        log_entry = f"{violation_id} | {start_seconds:.2f} | {end_seconds:.2f} | {violation_type}\n"
        with open("log.txt", "a") as log_file:
            log_file.write(log_entry)

def distance(coord1, coord2):
    return math.sqrt((coord1[0] - coord2[0]) ** 2 + (coord1[1] - coord2[1]) ** 2)

with open("log.txt", "w") as log_file:
    log_file.write("ID | Время начала события | Время конца события | Причина\n")

log_start_time = datetime.now()

# Переменные для хранения статусов
robot_position = None  # Координаты робота
helmet_positions = []  # Список координат касок
violation_start_time = None
violation_id = 1

cap = cv2.VideoCapture('0.mp4')

# Задайте границы зоны пола
floor_y_min = 0  # Пример: минимальная координата Y
floor_y_max = 500  # Пример: максимальная координата Y
threshold_distance = 1000  # Пороговое расстояние для оповещения (в пикселях)

# Метки объектов
robot_label = 'robot'
door_label = 'door'
helmet_label = 'helmet'
person_label = 'person'


def process_frame(frame):
    # Детектирование объектов
    results = model(frame)

    # Отрисовка результатов на кадре
    annotated_frame = results[0].plot()  # Отрисовка детекций на изображении

    # Отображение результата

    # Проверка нажатия клавиши 'q' для выхода
    if cv2.waitKey(1) & 0xFF == ord('q'):
        return False
    return annotated_frame

while True:
    ret, frame = cap.read()
    if not ret:
        break

    detections = detect_objects(frame)

    # Инициализация переменных
    human_on_floor = False
    helmet_on_floor = False
    if detections and hasattr(detections, 'boxes'):
        for box in detections.boxes:
            class_id = int(box.cls.item())
            position = (box.xyxy[0].item(), box.xyxy[1].item())
            print(class_id)

            print(f"Detected: Class ID {class_id}, Position {position}")  # Отладочная информация

            if class_id == 3:  # ID для человека
                # Проверка, находится ли человек на полу
                if (position[1] > floor_y_min) and (position[1] < floor_y_max):
                    human_on_floor = True
                    log_violation(violation_id, log_start_time, datetime.now(), "Человек на полу")

            elif class_id == 2:  # ID для каски
                # Проверка, находится ли каска на полу
                if (position[1] > floor_y_min) and (position[1] < floor_y_max):
                    helmet_on_floor = True
                    helmet_positions.append(position)
                    log_violation(violation_id, log_start_time, datetime.now(), "Каска на полу")

            elif class_id == 1:  # ID для робота (предположим, что ID 1 для робота)
                robot_position = position
                print(f"Robot position: {robot_position}")  # Отладочная информация
                for helmet_pos in helmet_positions:
                    print(f"Helmet position: {helmet_pos}")  # Отладочная информация
                    if distance(helmet_pos, robot_position) < threshold_distance:
                        print("Каска близко к роботу!")
                        log_violation(violation_id, log_start_time, datetime.now(), "Каска близко к роботу")

                
    current_time = datetime.now()
    elapsed_time = current_time - log_start_time
    with open("log.txt", "a") as log_file:
        log_file.write(f"Обработка кадра в {elapsed_time.total_seconds():.2f} секунд\n")
    cv2.imshow("Monitoring", process_frame(frame))

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()