Обнаружение объектов в реальном времени на Python с YOLO и OpenCV

import cv2
import numpy as np
 
# Загрузка YOLO
net = cv2.dnn.readNet("yolov3.weights", "yolov3.cfg")
layer_names = net.getLayerNames()
output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]
 
# Загрузка классов
with open("coco.names", "r") as f:
    classes = [line.strip() for line in f.readlines()]

<class_id> <x_center> <y_center> <width> <height>

transform = A.Compose([
    A.RandomBrightnessContrast(p=0.5),
    A.HorizontalFlip(p=0.5),
    A.Rotate(limit=10, p=0.5)
], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))

print(f"CUDA доступен: {torch.cuda.is_available()}")
print(f"Количество GPU: {torch.cuda.device_count()}")
print(f"Название устройства: {torch.cuda.get_device_name(0)}")

net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)

net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA_FP16)

def benchmark(net, input_size=(416, 416), iterations=100):
    # Создаём случайное изображение для теста
    img = np.random.random((input_size[0], input_size[1], 3))
    img = (img * 255).astype(np.uint8)
    
    # Разогрев
    blob = cv2.dnn.blobFromImage(img, 1/255.0, input_size, swapRB=True)
    net.setInput(blob)
    net.forward(output_layers)
    
    # Измерение времени
    start = time.time()
    for _ in range(iterations):
        blob = cv2.dnn.blobFromImage(img, 1/255.0, input_size, swapRB=True)
        net.setInput(blob)
        net.forward(output_layers)
    elapsed = time.time() - start
    
    fps = iterations / elapsed
    return fps

import cv2
import numpy as np
import time
 
# Загрузка YOLO
net = cv2.dnn.readNet("yolov3.weights", "yolov3.cfg")
with open("coco.names", "r") as f:
    classes = [line.strip() for line in f.readlines()]
layer_names = net.getLayerNames()
output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]
 
# Открываем видеопоток
cap = cv2.VideoCapture(0)  # 0 для веб-камеры, или путь к видеофайлу
 
while True:
    ret, frame = cap.read()
    if not ret:
        break
        
    height, width, channels = frame.shape
    
    # Создаём blob из кадра
    blob = cv2.dnn.blobFromImage(frame, 0.00392, (416, 416), (0, 0, 0), True, crop=False)
    
    # Пропускаем blob через сеть
    net.setInput(blob)
    start = time.time()
    outs = net.forward(output_layers)
    end = time.time()
    
    # Информация для отображения
    class_ids = []
    confidences = []
    boxes = []
    
    # Обработка выходных данных сети
    for out in outs:
        for detection in out:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            
            if confidence > 0.5:
                # Координаты объекта
                center_x = int(detection[0] * width)
                center_y = int(detection[1] * height)
                w = int(detection[2] * width)
                h = int(detection[3] * height)
                
                # Прямоугольные координаты
                x = int(center_x - w / 2)
                y = int(center_y - h / 2)
                
                boxes.append([x, y, w, h])
                confidences.append(float(confidence))
                class_ids.append(class_id)
    
    # Применение Non-Maximum Suppression
    indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)
    
    # Отрисовка результатов
    for i in range(len(boxes)):
        if i in indexes:
            x, y, w, h = boxes[i]
            label = str(classes[class_ids[i]])
            confidence = confidences[i]
            
            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
            cv2.putText(frame, f"{label} {confidence:.2f}", (x, y - 10), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    
    # Отображение FPS
    cv2.putText(frame, f"FPS: {1/(end-start):.2f}", (10, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
                
    cv2.imshow("Object Detection", frame)
    
    if cv2.waitKey(1) == 27:  # ESC для выхода
        break
 
cap.release()
cv2.destroyAllWindows()

def preprocess_image(image):
    # Удаляем шум с помощью билатерального фильтра (сохраняет края)
    denoised = cv2.bilateralFilter(image, 9, 75, 75)
    
    # Улучшаем контраст с помощью CLAHE
    lab = cv2.cvtColor(denoised, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    cl = clahe.apply(l)
    enhanced_lab = cv2.merge((cl, a, b))
    enhanced = cv2.cvtColor(enhanced_lab, cv2.COLOR_LAB2BGR)
    
    return enhanced

def process_detections(outputs, confidence_threshold=0.5, nms_threshold=0.4):
    class_ids = []
    confidences = []
    boxes = []
    
    for out in outputs:
        for detection in out:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            
            if confidence > confidence_threshold:
                center_x = int(detection[0] * width)
                center_y = int(detection[1] * height)
                w = int(detection[2] * width)
                h = int(detection[3] * height)
                
                x = int(center_x - w / 2)
                y = int(center_y - h / 2)
                
                boxes.append([x, y, w, h])
                confidences.append(float(confidence))
                class_ids.append(class_id)
    
    # Применение Non-Maximum Suppression
    indexes = cv2.dnn.NMSBoxes(boxes, confidences, confidence_threshold, nms_threshold)
    
    return boxes, confidences, class_ids, indexes

import threading
import queue
 
class VideoStreamProcessor:
    def __init__(self, source=0):
        self.capture = cv2.VideoCapture(source)
        self.frame_queue = queue.Queue(maxsize=2)
        self.result_queue = queue.Queue(maxsize=2)
        self.stopped = False
        
    def start(self):
        # Поток для захвата кадров
        capture_thread = threading.Thread(target=self._capture_loop)
        capture_thread.daemon = True
        capture_thread.start()
        
        # Поток для обработки кадров
        process_thread = threading.Thread(target=self._process_loop)
        process_thread.daemon = True
        process_thread.start()
        
        return self
    
    def _capture_loop(self):
        while not self.stopped:
            if not self.frame_queue.full():
                ret, frame = self.capture.read()
                if ret:
                    self.frame_queue.put(frame)
                else:
                    self.stopped = True
            else:
                time.sleep(0.01)  # Небольшая задержка, чтобы не загружать CPU
    
    def _process_loop(self):
        while not self.stopped:
            if not self.frame_queue.empty() and not self.result_queue.full():
                frame = self.frame_queue.get()
                # Здесь вся обработка YOLO
                # ...
                # Результат помещаем в очередь вывода
                self.result_queue.put(processed_frame)
            else:
                time.sleep(0.01)
    
    def read(self):
        if self.result_queue.empty():
            return False, None
        return True, self.result_queue.get()
    
    def stop(self):
        self.stopped = True

class BoxFilter:
    def __init__(self, process_noise=0.01, measurement_noise=0.1):
        self.kalman = cv2.KalmanFilter(8, 4)
        # Матрица перехода состояния (F)
        self.kalman.transitionMatrix = np.array([
            [1, 0, 0, 0, 1, 0, 0, 0],  # x
            [0, 1, 0, 0, 0, 1, 0, 0],  # y
            [0, 0, 1, 0, 0, 0, 1, 0],  # w
            [0, 0, 0, 1, 0, 0, 0, 1],  # h
            [0, 0, 0, 0, 1, 0, 0, 0],  # dx
            [0, 0, 0, 0, 0, 1, 0, 0],  # dy
            [0, 0, 0, 0, 0, 0, 1, 0],  # dw
            [0, 0, 0, 0, 0, 0, 0, 1]   # dh
        ], np.float32)
        
        # Матрица измерения (H)
        self.kalman.measurementMatrix = np.array([
            [1, 0, 0, 0, 0, 0, 0, 0],
            [0, 1, 0, 0, 0, 0, 0, 0],
            [0, 0, 1, 0, 0, 0, 0, 0],
            [0, 0, 0, 1, 0, 0, 0, 0]
        ], np.float32)
        
        # Установка шума процесса и измерения
        self.kalman.processNoiseCov = np.eye(8, dtype=np.float32) * process_noise
        self.kalman.measurementNoiseCov = np.eye(4, dtype=np.float32) * measurement_noise
        
        self.initialized = False
        
    def update(self, bbox):
        x, y, w, h = bbox
        
        if not self.initialized:
            self.kalman.statePre = np.array([[x], [y], [w], [h], [0], [0], [0], [0]], np.float32)
            self.kalman.statePost = self.kalman.statePre.copy()
            self.initialized = True
        
        # Предсказание
        prediction = self.kalman.predict()
        
        # Обновление с помощью измерения
        measurement = np.array([[x], [y], [w], [h]], np.float32)
        corrected = self.kalman.correct(measurement)
        
        # Возвращаем сглаженные координаты
        return [int(corrected[0]), int(corrected[1]), int(corrected[2]), int(corrected[3])]

from scipy.optimize import linear_sum_assignment
 
class Sort:
    def __init__(self, max_age=10, min_hits=3, iou_threshold=0.3):
        self.max_age = max_age
        self.min_hits = min_hits
        self.iou_threshold = iou_threshold
        self.trackers = []
        self.frame_count = 0
        self.next_id = 1
        
    def update(self, detections):
        self.frame_count += 1
        
        # Предсказание новых положений трекеров
        predicted_boxes = []
        for tracker in self.trackers:
            predicted_box = tracker.predict()
            predicted_boxes.append(predicted_box)
        
        # Создаём матрицу IoU между детекциями и предсказаниями
        iou_matrix = np.zeros((len(detections), len(predicted_boxes)))
        for d, detection in enumerate(detections):
            for t, prediction in enumerate(predicted_boxes):
                iou_matrix[d, t] = self._iou(detection, prediction)
        
        # Используем Венгерский алгоритм для назначения детекций трекерам
        row_indices, col_indices = linear_sum_assignment(-iou_matrix)
        matched_indices = np.array(list(zip(row_indices, col_indices)))
        
        # Обновление трекеров с соответствующими детекциями
        for m in matched_indices:
            if iou_matrix[m[0], m[1]] < self.iou_threshold:
                continue
            self.trackers[m[1]].update(detections[m[0]])
        
        # Создаём новые трекеры для необработанных детекций
        unmatched_detections = [i for i in range(len(detections))
                               if i not in matched_indices[:, 0]]
        for i in unmatched_detections:
            self.trackers.append(KalmanBoxTracker(detections[i], self.next_id))
            self.next_id += 1
        
        # Удаляем трекеры, которые не обновлялись длительное время
        self.trackers = [t for t in self.trackers if t.time_since_update < self.max_age]
        
        # Формируем и возвращаем результаты
        results = []
        for tracker in self.trackers:
            if tracker.hits >= self.min_hits:
                results.append((tracker.id, tracker.get_state()))
        return results
    
    def _iou(self, box1, box2):
        # Реализация вычисления IoU между двумя боксами
        x1, y1, w1, h1 = box1
        x2, y2, w2, h2 = box2
        
        xi1 = max(x1, x2)
        yi1 = max(y1, y2)
        xi2 = min(x1 + w1, x2 + w2)
        yi2 = min(y1 + h1, y2 + h2)
        
        if xi2 <= xi1 or yi2 <= yi1:
            return 0.0
            
        intersection = (xi2 - xi1) * (yi2 - yi1)
        box1_area = w1 * h1
        box2_area = w2 * h2
        
        return intersection / (box1_area + box2_area - intersection)

# В основном цикле обработки
tracker = Sort()
 
while True:
    ret, frame = cap.read()
    if not ret:
        break
        
    # Обработка YOLO (как в предыдущем коде)
    # ...
    
    # Преобразуем боксы в формат [x, y, w, h]
    detections = []
    for i in indexes:
        i = i[0]
        detections.append(boxes[i])
    
    # Обновляем трекер
    tracked_objects = tracker.update(detections)
    
    # Отрисовка результатов с уникальными ID
    for obj_id, box in tracked_objects:
        x, y, w, h = [int(v) for v in box]
        color = get_color_by_id(obj_id)  # Функция, возвращающая уникальный цвет для ID
        cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
        cv2.putText(frame, f"ID: {obj_id}", (x, y - 10), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
    
    cv2.imshow("Tracking", frame)
    # ...

# Подготовка данных на GPU
if cv2.cuda.getCudaEnabledDeviceCount() > 0:
    gpu_frame = cv2.cuda_GpuMat()
    gpu_frame.upload(frame)
    gpu_frame_resized = cv2.cuda.resize(gpu_frame, (416, 416))
    preprocessed = gpu_frame_resized.download()
else:
    preprocessed = cv2.resize(frame, (416, 416))

frame_counter = 0
process_every_n_frames = 2
 
while True:
    ret, frame = cap.read()
    frame_counter += 1
    
    if frame_counter % process_every_n_frames != 0:
        continue
        
    # Обычная обработка
    # ...

def adaptive_resolution(frame, target_fps=30, current_fps=None):
    if current_fps is None or current_fps >= target_fps:
        return cv2.resize(frame, (416, 416))
    
    # Если FPS ниже целевого, уменьшаем разрешение
    scale = max(0.5, min(1.0, current_fps / target_fps))
    new_size = (int(416 * scale), int(416 * scale))
    return cv2.resize(frame, new_size)

# Определение маски интереса 
roi_mask = np.zeros(frame.shape[:2], dtype=np.uint8)
roi_points = np.array([[100, 100], [100, 400], [500, 400], [500, 100]])
cv2.fillPoly(roi_mask, [roi_points], 255)
 
# Применение маски
masked_frame = cv2.bitwise_and(frame, frame, mask=roi_mask)

class LineCounter:
    def __init__(self, line_start, line_end):
        self.line_start = line_start  # (x1, y1)
        self.line_end = line_end      # (x2, y2)
        self.counts = {}
        self.crossed_ids = set()
        
    def update(self, tracked_objects):
        for obj_id, box in tracked_objects:
            x, y, w, h = box
            center_x = x + w // 2
            center_y = y + h // 2
            
            if self._point_crossed_line((center_x, center_y)) and obj_id not in self.crossed_ids:
                self.crossed_ids.add(obj_id)
                if obj_id in self.counts:
                    self.counts[obj_id] += 1
                else:
                    self.counts[obj_id] = 1
                    
    def _point_crossed_line(self, point):
        # Проверка, пересекает ли точка линию
        # Здесь можно использовать более сложную логику,
        # учитывающую направление движения
        x, y = point
        x1, y1 = self.line_start
        x2, y2 = self.line_end
        
        # Расстояние от точки до линии
        distance = abs((y2-y1)*x - (x2-x1)*y + x2*y1 - y2*x1) / \
                  np.sqrt((y2-y1)[B]2 + (x2-x1)[/B]2)
                  
        return distance < 5  # Порог для определения пересечения

import csv
from datetime import datetime
 
def log_detections(detections, filename="detections_log.csv"):
    with open(filename, 'a', newline='') as csvfile:
        fieldnames = ['timestamp', 'object_id', 'class', 'confidence', 'x', 'y', 'w', 'h']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        
        if csvfile.tell() == 0:  # Если файл пуст, пишем заголовки
            writer.writeheader()
        
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        for obj_id, (x, y, w, h), class_id, confidence in detections:
            writer.writerow({
                'timestamp': timestamp,
                'object_id': obj_id,
                'class': class_names[class_id],
                'confidence': confidence,
                'x': x, 'y': y, 'w': w, 'h': h
            })

class TrajectoryVisualizer:
    def __init__(self, max_history=30):
        self.trajectories = {}
        self.max_history = max_history
        
    def update(self, tracked_objects):
        for obj_id, box in tracked_objects:
            x, y, w, h = box
            center = (x + w // 2, y + h // 2)
            
            if obj_id not in self.trajectories:
                self.trajectories[obj_id] = []
                
            self.trajectories[obj_id].append(center)
            
            # Ограничиваем историю
            if len(self.trajectories[obj_id]) > self.max_history:
                self.trajectories[obj_id].pop(0)
                
    def draw(self, frame):
        for obj_id, trajectory in self.trajectories.items():
            if len(trajectory) < 2:
                continue
                
            color = get_color_by_id(obj_id)
            
            # Рисуем линии траектории
            for i in range(1, len(trajectory)):
                thickness = int(np.sqrt(self.max_history / float(i + 1)) * 2.5)
                cv2.line(frame, trajectory[i-1], trajectory[i], color, thickness)
                
        return frame

# Адаптивная настройка контраста для сложных сцен
def enhance_scene(frame, brightness_threshold=120):
    # Оцениваем среднюю яркость
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    avg_brightness = np.mean(gray)
    
    # Для тёмных сцен повышаем яркость и контраст
    if avg_brightness < brightness_threshold:
        # Конвертируем в HSV для работы с яркостью
        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
        h, s, v = cv2.split(hsv)
        
        # Увеличиваем яркость 
        v = cv2.add(v, 30)
        # Повышаем контраст
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
        v = clahe.apply(v)
        
        # Собираем изображение обратно
        hsv = cv2.merge([h, s, v])
        return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
    
    return frame

# Пример дообучения YOLOv5 (с использованием Ultralytics)
!python train.py --img 640 --batch 16 --epochs 50 --data custom.yaml \
--weights yolov5s.pt --cache

# Минималистичный REST API для обработки изображений
from flask import Flask, request, jsonify
import base64
import cv2
import numpy as np
 
app = Flask(__name__)
detector = YoloDetector()  # Наш класс детектора
 
@app.route('/detect', methods=['POST'])
def detect_objects():
    # Получаем изображение из запроса
    image_data = request.json.get('image')
    image_bytes = base64.b64decode(image_data)
    
    # Конвертируем в формат OpenCV
    nparr = np.frombuffer(image_bytes, np.uint8)
    image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    
    # Обнаруживаем объекты
    detections = detector.detect(image)
    
    # Возвращаем результаты в JSON
    return jsonify({'detections': detections})
 
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)

# Конвертация модели в TensorRT
import tensorrt as trt
 
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)
 
# Парсим ONNX-модель
with open('yolov3.onnx', 'rb') as model:
    if not parser.parse(model.read()):
        print('ERROR: Failed to parse ONNX model')
        for error in range(parser.num_errors):
            print(parser.get_error(error))
 
# Создаём оптимизированный TensorRT-движок
config = builder.create_builder_config()
config.max_workspace_size = 1 << 28  # 256 МБ
engine = builder.build_engine(network, config)