MLops_ONNX

import os
import glob
import cv2
import time
import numpy as np
from ultralytics import YOLO

# Установка устройства для инференса ('cpu' или 'cuda')
DEVICE = 'cuda'  # Измените на 'cpu', если хотите использовать CPU

DATA_YAML_PATH = 'C:/Users/edimv/Desktop/stenosis/data.yaml'
model_path = 'best.onnx'

# Загрузка модели ONNX
model = YOLO(model_path, task="detect")
# Не используем model.to(DEVICE), так как для ONNX-моделей это не поддерживается

def load_annotations(annotation_path, img_width, img_height):
    """
    Загрузка аннотаций из TXT-файла в формате YOLO и преобразование в абсолютные пиксельные координаты.
    """
    annotations = []
    with open(annotation_path, 'r') as f:
        for line in f:
            parts = line.strip().split()
            label = int(parts[0])
            x_center_norm, y_center_norm, width_norm, height_norm = map(float, parts[1:])
            # Преобразование нормализованных координат в абсолютные пиксельные координаты
            x_center = x_center_norm * img_width
            y_center = y_center_norm * img_height
            width = width_norm * img_width
            height = height_norm * img_height
            annotations.append([label, x_center, y_center, width, height])
    return annotations

def calculate_iou(box1, box2):
    """
    Вычисление IoU между двумя боксами в формате (x_center, y_center, width, height).
    """
    # Преобразование боксов в (x1, y1, x2, y2)
    box1_x1 = box1[0] - box1[2] / 2
    box1_y1 = box1[1] - box1[3] / 2
    box1_x2 = box1[0] + box1[2] / 2
    box1_y2 = box1[1] + box1[3] / 2

    box2_x1 = box2[0] - box2[2] / 2
    box2_y1 = box2[1] - box2[3] / 2
    box2_x2 = box2[0] + box2[2] / 2
    box2_y2 = box2[1] + box2[3] / 2

    # Вычисление координат пересечения
    inter_x1 = max(box1_x1, box2_x1)
    inter_y1 = max(box1_y1, box2_y1)
    inter_x2 = min(box1_x2, box2_x2)
    inter_y2 = min(box1_y2, box2_y2)

    # Вычисление площади пересечения
    inter_area = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)

    # Вычисление площадей боксов
    box1_area = (box1_x2 - box1_x1) * (box1_y2 - box1_y1)
    box2_area = (box2_x2 - box2_x1) * (box2_y2 - box2_y1)

    # Вычисление IoU
    union_area = box1_area + box2_area - inter_area
    if union_area == 0:
        return 0  # Избегаем деления на ноль
    iou = inter_area / union_area
    return iou

def draw_boxes(img, pred_boxes, gt_boxes):
    """
    Отрисовка предсказанных и истинных боксов на изображении.
    Предсказанные боксы — зелёные, истинные боксы — красные.
    """
    # Отрисовка предсказанных боксов
    for pred in pred_boxes:
        box = pred['box']
        label = pred['label']
        confidence = pred['conf']
        # Преобразование из [x_center, y_center, width, height] в [x1, y1, x2, y2]
        x_center, y_center, width, height = box
        x1 = int(x_center - width / 2)
        y1 = int(y_center - height / 2)
        x2 = int(x_center + width / 2)
        y2 = int(y_center + height / 2)
        # Отрисовка прямоугольника
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)  # Зелёный цвет для предсказаний
        # Добавление метки и уверенности
        cv2.putText(img, f'Pred {label}:{confidence:.2f}', (x1, y1 - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)

    # Отрисовка истинных боксов
    for gt in gt_boxes:
        label = gt[0]
        box = gt[1:]  # [x_center, y_center, width, height]
        x_center, y_center, width, height = box
        x1 = int(x_center - width / 2)
        y1 = int(y_center - height / 2)
        x2 = int(x_center + width / 2)
        y2 = int(y_center + height / 2)
        # Отрисовка прямоугольника
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2)  # Красный цвет для истинных боксов
        # Добавление метки
        cv2.putText(img, f'GT {label}', (x1, y1 - 25),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
    return img

def test_iou(test_folder):
    """
    Тестирование IoU на наборе изображений и аннотаций с использованием предсказаний модели YOLO.
    Также вычисляет средние времена, метрики и сохраняет примеры.
    """
    iou_results = []
    preprocessing_times = []
    processing_times = []
    postprocessing_times = []
    per_image_data = []  # Для хранения данных по каждому изображению
    low_iou_count = 0
    total_TP = 0
    total_FP = 0
    total_FN = 0

    image_paths = glob.glob(os.path.join(test_folder, "*.bmp"))
    if not os.path.exists('samples1'):
        os.makedirs('samples1')

    for img_path in image_paths:
        # Предобработка
        t1 = time.time()
        annotation_path = img_path.replace(".bmp", ".txt")
        img_filename = os.path.basename(img_path)

        if not os.path.exists(annotation_path):
            print(f"Annotation missing for {img_path}")
            continue

        img = cv2.imread(img_path)
        if img is None:
            print(f"Failed to read image {img_path}")
            continue

        # Сохранение оригинальных размеров изображения (если потребуется)
        original_img_height, original_img_width = img.shape[:2]

        # Изменение размера изображения до 800x800
        img_resized = img
        img_height, img_width = img_resized.shape[:2]

        # Преобразование изображения в RGB
        img_rgb = cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB)

        # Загрузка аннотаций с использованием новых размеров изображения
        gt_boxes = load_annotations(annotation_path, img_width, img_height)
        t2 = time.time()
        preprocessing_times.append(t2 - t1)

        # Инференс
        t3 = time.time()
        results = model.predict(img_rgb, imgsz = 800, device=DEVICE)
        t4 = time.time()
        processing_times.append(t4 - t3)

        # Постобработка
        t5 = time.time()
        if results[0].boxes is None or len(results[0].boxes) == 0:
            print(f"No predictions for {img_path}")
            iou_results.append(0)
            per_image_data.append({
                'img_path': img_path,
                'max_iou': 0,
                'ious': [0],
                'TP': 0,
                'FP': 0,
                'FN': len(gt_boxes),
                'pred_boxes': [],
                'gt_boxes': gt_boxes
            })
            total_FN += len(gt_boxes)
            if 0 < 0.3:
                low_iou_count += 1
            t6 = time.time()
            postprocessing_times.append(t6 - t5)
            continue

        # Получение предсказанных боксов и вероятностей
        predictions = results[0].boxes.xywh.cpu().numpy()  # [x_center, y_center, width, height]
        confidences = results[0].boxes.conf.cpu().numpy()  # Уверенности
        labels = results[0].boxes.cls.cpu().numpy()  # Классы

        pred_boxes = []
        for i in range(len(predictions)):
            pred_box = predictions[i]
            confidence = confidences[i]
            label = int(labels[i])
            pred_boxes.append({'box': pred_box, 'conf': confidence, 'label': label})

        # Сортировка предсказаний по уверенности
        pred_boxes.sort(key=lambda x: x['conf'], reverse=True)

        matched_gt = []
        ious = []
        TP = 0
        FP = 0
        FN = 0
        for pred in pred_boxes:
            pred_box = pred['box']
            pred_conf = pred['conf']
            pred_label = pred['label']
            best_iou = 0
            best_gt_idx = -1
            for idx, gt_box in enumerate(gt_boxes):
                if idx in matched_gt:
                    continue  # Уже сопоставлено
                gt_label = gt_box[0]
                gt_box_coords = gt_box[1:]
                iou = calculate_iou(pred_box, gt_box_coords)
                if iou > best_iou:
                    best_iou = iou
                    best_gt_idx = idx
            if best_iou >= 0.5:
                TP += 1
                total_TP += 1
                matched_gt.append(best_gt_idx)
            else:
                FP += 1
                total_FP += 1
            ious.append(best_iou)

        FN = len(gt_boxes) - len(matched_gt)
        total_FN += FN

        max_iou = max(ious) if ious else 0
        if max_iou < 0.3:
            low_iou_count += 1

        iou_results.append(max_iou)
        per_image_data.append({
            'img_path': img_path,
            'max_iou': max_iou,
            'ious': ious,
            'TP': TP,
            'FP': FP,
            'FN': FN,
            'pred_boxes': pred_boxes,
            'gt_boxes': gt_boxes
        })
        t6 = time.time()
        postprocessing_times.append(t6 - t5)

    # Вычисление средних времен
    avg_preprocessing_time = sum(preprocessing_times) / len(preprocessing_times) if preprocessing_times else 0
    avg_processing_time = sum(processing_times) / len(processing_times) if processing_times else 0
    avg_postprocessing_time = sum(postprocessing_times) / len(postprocessing_times) if postprocessing_times else 0

    # Вычисление метрик
    precision = total_TP / (total_TP + total_FP) if (total_TP + total_FP) > 0 else 0
    recall = total_TP / (total_TP + total_FN) if (total_TP + total_FN) > 0 else 0

    # Вычисление mAP
    iou_thresholds = np.arange(0.5, 1.0, 0.05)
    APs = []
    for iou_thresh in iou_thresholds:
        ap = calculate_ap(per_image_data, iou_thresh)
        APs.append(ap)
    mAP50 = APs[0]
    mAP50_95 = np.mean(APs)

    # Сохранение лучших и худших примеров с отрисованными боксами
    per_image_data.sort(key=lambda x: x['max_iou'], reverse=True)
    best_samples = per_image_data[:3]
    worst_samples = per_image_data[-3:]

    for sample in best_samples:
        img = cv2.imread(sample['img_path'])
        # Изменение размера изображения до 800x800
        img_resized = img
        pred_boxes = sample['pred_boxes']
        gt_boxes = sample['gt_boxes']
        img_with_boxes = draw_boxes(img_resized, pred_boxes, gt_boxes)
        img_name = os.path.basename(sample['img_path'])
        save_path = os.path.join('samples1', f'good_{img_name}')
        cv2.imwrite(save_path, img_with_boxes)

    for sample in worst_samples:
        img = cv2.imread(sample['img_path'])
        # Изменение размера изображения до 800x800
        img_resized = img
        pred_boxes = sample['pred_boxes']
        gt_boxes = sample['gt_boxes']
        img_with_boxes = draw_boxes(img_resized, pred_boxes, gt_boxes)
        img_name = os.path.basename(sample['img_path'])
        save_path = os.path.join('samples1', f'bad_{img_name}')
        cv2.imwrite(save_path, img_with_boxes)

    # Возврат результатов
    results = {
        'average_iou': sum(iou_results) / len(iou_results) if iou_results else 0,
        'avg_preprocessing_time': avg_preprocessing_time,
        'avg_processing_time': avg_processing_time,
        'avg_postprocessing_time': avg_postprocessing_time,
        'precision': precision,
        'recall': recall,
        'mAP50': mAP50,
        'mAP50_95': mAP50_95,
        'low_iou_count': low_iou_count
    }
    return results

def calculate_ap(per_image_data, iou_threshold=0.5):
    """
    Вычисление Average Precision (AP) при заданном пороге IoU.
    """
    # Собираем все предсказания и аннотации
    detections = []
    annotations = []

    for idx, data in enumerate(per_image_data):
        preds = data['pred_boxes']
        gts = data['gt_boxes']
        # Добавляем предсказания
        for pred in preds:
            detections.append([idx, pred['label'], pred['conf'], *pred['box']])
        # Добавляем аннотации
        for gt in gts:
            annotations.append([idx, gt[0], *gt[1:]])

    if len(annotations) == 0:
        return 0

    # Сортировка предсказаний по уверенности
    detections.sort(key=lambda x: x[2], reverse=True)

    TP = np.zeros(len(detections))
    FP = np.zeros(len(detections))

    detected_annotations = []

    for d_idx, detection in enumerate(detections):
        image_idx = detection[0]
        detection_label = detection[1]
        detection_conf = detection[2]
        detection_box = detection[3:]

        gt_annotations = [ann for ann in annotations if ann[0] == image_idx and ann[1] == detection_label]
        max_iou = 0
        matched_ann_idx = -1
        for ann_idx, ann in enumerate(gt_annotations):
            ann_box = ann[2:]
            iou = calculate_iou(detection_box, ann_box)
            if iou > max_iou:
                max_iou = iou
                matched_ann_idx = ann_idx
        if max_iou >= iou_threshold and (image_idx, matched_ann_idx) not in detected_annotations:
            TP[d_idx] = 1
            detected_annotations.append((image_idx, matched_ann_idx))
        else:
            FP[d_idx] = 1

    cumulative_TP = np.cumsum(TP)
    cumulative_FP = np.cumsum(FP)
    recalls = cumulative_TP / len(annotations)
    precisions = cumulative_TP / (cumulative_TP + cumulative_FP)

    # Избегаем деления на ноль
    recalls = np.concatenate(([0], recalls))
    precisions = np.concatenate(([1], precisions))

    # Вычисление AP
    AP = 0
    for i in range(1, len(recalls)):
        AP += (recalls[i] - recalls[i - 1]) * precisions[i]
    return AP

if __name__ == '__main__':
    TEST_FOLDER = "test"
    results = test_iou(TEST_FOLDER)
    print("Average IoU on test set:", results['average_iou'])
    print("Average preprocessing time per image:", results['avg_preprocessing_time'])
    print("Average processing time per image:", results['avg_processing_time'])
    print("Average post-processing time per image:", results['avg_postprocessing_time'])
    print("Precision:", results['precision'])
    print("Recall:", results['recall'])
    print("mAP@0.5:", results['mAP50'])
    print("mAP@0.5:0.95:", results['mAP50_95'])
    print("Number of examples with IoU below 0.3:", results['low_iou_count'])