AIEffDetDivided.py - master - Nemo_499/DecodingX-rayImages

AIEffDetDivided.py
1 170 строк41 KB
Вдовин Денис
Конфигурации объектов и ии
28 апр 2026, 09:38
28 апр 2026, 09:3807f4415
import os
import cv2
import torch

torch._C._jit_set_profiling_executor(False)
torch._C._jit_set_profiling_mode(False)
torch._C._jit_override_can_fuse_on_gpu(False)
torch._C._jit_override_can_fuse_on_cpu(False)

import numpy as np
from torch.utils.data import Dataset
from effdet import EfficientDet, create_model, DetBenchTrain, DetBenchPredict, get_efficientdet_config
import onnxruntime as ort
import traceback
import kornia.augmentation as K
import torch.nn as nn

from fastai.data.core import DataLoaders
from fastai.learner import Learner
from fastai.callback.schedule import fit_one_cycle
from fastai.callback.tracker import EarlyStoppingCallback
from fastai.callback.progress import CSVLogger
from fastai.vision.all import *
from fastai.callback.core import Callback
from fastai.optimizer import Adam, SGD, RMSProp, RAdam, QHAdam, Lamb, Larc, Lookahead, OptimWrapper
from fastai.data.load import DataLoader as FDataLoader

from torchcam.methods import GradCAM
from torchcam.utils import overlay_mask
from torchvision.transforms.functional import to_pil_image
import PIL.Image as Image

import matplotlib.pyplot as plt

from NamesClass import names_id as NamesClass
from NamesClass import set_objects

IMG_SIZE = 640
batch_size = 16
batch_size_lr = 1
model_name = 'tf_efficientdet_d1_ap'
model_save_name_buf=""
best_f1 = 0
count_epochs = 0
NUM_CLASSES = len(NamesClass)
print(NUM_CLASSES)
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Устройство: {DEVICE}")

SnapshotCounter = 0

names = NamesClass

Var_dataset = 0

# =========================
# DATASET
# =========================
class XRayDataset(Dataset):
    def __init__(self, img_dir, label_dir):
        self.label_dir = label_dir
        self.img_dir = img_dir
        self.tiles = []

        flag_append_tiles = False

        for img_name in sorted(os.listdir(img_dir)):
            img_path = (os.path.join(img_dir, img_name))
            image = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
            lbl_path = os.path.join(self.label_dir, img_name.replace(".jpg", ".txt"))

            h, w = image.shape[:2]

            x_starts = []
            x = 0
            while x + IMG_SIZE <= w:
                x_starts.append(x)
                x += IMG_SIZE // 2

            y_starts = []
            y = 0 #h // 2 - IMG_SIZE // 2
            while y + IMG_SIZE <= h:
                y_starts.append(y)
                y += IMG_SIZE // 2


            # Добавление последнего тайла, если нужно
            if x_starts[-1] + IMG_SIZE < w:
                x_starts.append(w - IMG_SIZE)
            if y_starts[-1] + IMG_SIZE < h:
                y_starts.append(h - IMG_SIZE)
            if os.path.exists(lbl_path):
                all_labels = []
                with open(lbl_path) as f:
                    for line in f:
                        parts = list(map(float, line.split()))
                        if len(parts) >= 5 and parts[0] in set_objects[Var_dataset][2]: # Проверка, что строка не пустая и id класса входит в список
                            # Перевод нормализованных координат в пиксели
                            x1, y1 = int(parts[1] * w), int(parts[2] * h)
                            x2, y2 = int(parts[3] * w), int(parts[4] * h)
                            all_labels.append([x1, y1, x2, y2])
                    for y0 in y_starts:
                        for x0 in x_starts:
                            for bbox in all_labels:
                                if(clip_box(bbox, x0, y0, x0 + IMG_SIZE, y0 + IMG_SIZE) is not None):
                                    flag_append_tiles = True
                                    break
                            if flag_append_tiles:
                                self.tiles.append((img_name, x0, y0))
                                flag_append_tiles = False


    def __len__(self):
        return len(self.tiles)

    def __getitem__(self, idx):

        img_name, x0, y0 = self.tiles[idx]
        img_path = (os.path.join(self.img_dir, img_name))
        lbl_path = os.path.join(self.label_dir, img_name.replace(".jpg", ".txt"))

        image_orig = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)

        #clahe = cv2.createCLAHE(clipLimit=3.5, tileGridSize=(8,8))# применение clahe
        #image_clahe = clahe.apply(image_orig)

        image_orig = cv2.cvtColor(image_orig, cv2.COLOR_GRAY2RGB)
        h, w = image_orig.shape[:2]

        # --- load boxes ---
        boxes, labels = [], []
        if os.path.exists(lbl_path):
            with open(lbl_path) as f:
                for line in f:
                    parts = list(map(float, line.split()))
                    if parts[0] in set_objects[Var_dataset][2]:
                        cls = int(parts[0]) - set_objects[Var_dataset][1]
                        xs = parts[1::2]
                        ys = parts[2::2]

                        boxes.append([
                            min(xs) * w,
                            min(ys) * h,
                            max(xs) * w,
                            max(ys) * h
                        ])
                        labels.append(cls)
        
        x1 = x0 + IMG_SIZE
        y1 = y0 + IMG_SIZE

        tile_img = image_orig[y0:y1, x0:x1].copy()

        tile_boxes = []
        tile_labels = []

        for box, label in zip(boxes, labels):
            clipped = clip_box(box, x0, y0, x1, y1)
            if clipped:
                tile_boxes.append(clipped)
                tile_labels.append(label)

        image = torch.from_numpy(tile_img).permute(2,0,1).float()/255.
        # нормализация 
        mean = torch.tensor([0.485, 0.456, 0.406]).view(3,1,1)
        std = torch.tensor([0.229, 0.224, 0.225]).view(3,1,1)
        image = (image - mean) / std

        if len(tile_boxes) == 0:
            boxes = torch.zeros((0,4), dtype=torch.float32)
        else:
            boxes = torch.tensor(tile_boxes, dtype=torch.float32)
        labels = torch.tensor(tile_labels, dtype=torch.int64)

        target = {
            "boxes": boxes,
            "labels": labels
        }

        return image, target

def clip_box(box, x1, y1, x2, y2):
    bx1, by1, bx2, by2 = box

    ix1 = max(bx1, x1)
    iy1 = max(by1, y1)
    ix2 = min(bx2, x2)
    iy2 = min(by2, y2)

    if ix1 >= ix2 or iy1 >= iy2:
        return None

    return [
        ix1 - x1,
        iy1 - y1,
        ix2 - x1,
        iy2 - y1
    ]

class KorniaAugment(nn.Module):
    def __init__(self):
        super().__init__()
        self.augs = K.AugmentationSequential(
            K.RandomHorizontalFlip(p=0.5),
            K.RandomRotation(degrees=5, p=0.5),
            data_keys=["input", "bbox"] 
        )

    def forward(self, x, bboxes):
        # bboxes на входе: [B, N, 4] в формате xyxy
        
        # 1. Преобразавание [B, N, 4] -> [B, N, 4, 2] (4 вершины по 2 координаты)
        # Формирование углов: левый-верхний, правый-верхний, правый-нижний, левый-нижний
        x1, y1, x2, y2 = bboxes.unbind(-1)
        
        # v shape: [B, N, 4, 2]
        v = torch.stack([
            torch.stack([x1, y1], dim=-1),
            torch.stack([x2, y1], dim=-1),
            torch.stack([x2, y2], dim=-1),
            torch.stack([x1, y2], dim=-1)
        ], dim=-2)

        # 2. Применение аугментации к изображениям и вершинам
        img_out, v_out = self.augs(x, v)

        # 3. Обратное преобразование [B, N, 4, 2] -> [B, N, 4] (xyxy)
        x_min = v_out[..., 0].min(dim=-1)[0]
        y_min = v_out[..., 1].min(dim=-1)[0]
        x_max = v_out[..., 0].max(dim=-1)[0]
        y_max = v_out[..., 1].max(dim=-1)[0]
        
        bboxes_out = torch.stack([x_min, y_min, x_max, y_max], dim=-1)

        return img_out, bboxes_out

class KorniaAugmentCallback(Callback):
    def __init__(self, aug_pipeline):
        self.aug_pipeline = aug_pipeline
    
    def before_fit(self):
        self.aug_pipeline = self.aug_pipeline.to(self.learn.dls.device)
    
    def before_batch(self):
        # self.learn.xb — это кортеж кортежей: ((imgs, targets),)
        # Извлечение внутреннего кортежа
        if not self.training: return

        imgs, targets = self.learn.xb[0] 
        bboxes = targets['bbox'] # Сейчас они в формате yxyx (из collate_fn)
        
        # 1. Временный перевод yxyx -> xyxy для Kornia
        bboxes_xyxy = bboxes.clone()
        bboxes_xyxy[:, :, [0, 1, 2, 3]] = bboxes[:, :, [1, 0, 3, 2]]
        
        # 2. Применение аугментации
        imgs_aug, bboxes_aug_xyxy = self.aug_pipeline(imgs, bboxes_xyxy)
        
        # 3. Перевод обратно xyxy -> yxyx
        bboxes_aug_yxyx = bboxes_aug_xyxy.clone()
        bboxes_aug_yxyx[:, :, [0, 1, 2, 3]] = bboxes_aug_xyxy[:, :, [1, 0, 3, 2]]
        
        # 4. Обновление данны в таргетах
        targets['bbox'] = bboxes_aug_yxyx
        
        # Перезапись входных данных батча
        self.learn.__dict__['xb'] = ((imgs_aug, targets),)

class EffDetTrain(nn.Module):

    def __init__(self, num_classes, weights_path=None):
        super().__init__()

        self.model_name = model_name
        self.num_classes = num_classes

        model = create_model(
            model_name,
            num_classes=num_classes,
            pretrained = (weights_path is None)
        )

        self.model = DetBenchTrain(model)

        """for name, _ in model.named_modules():
            print(name)"""

        if weights_path is not None:
            self.load_weights(weights_path)
            print(f"Загружены веса для дообучения")

    def load_weights(self, path):

        checkpoint = torch.load(path, map_location=DEVICE)

        state_dict = checkpoint["model_state_dict"]

        self.model.model.load_state_dict(state_dict)

        print(f"Путь загрузки: {path}")

    def forward(self, images, targets=None):
        # Если targets=None, значит FastAI передал (imgs, targets) в первый аргумент
        if targets is None and isinstance(images, (list, tuple)):
            images, targets = images
        
        # Передача в DetBenchTrain
        return self.model(images, targets)

class GUICallback(Callback):

    def __init__(self, gui_callback,save_stats_signal,train_dl, val_dl):
        self.gui_callback = gui_callback
        self.save_stats_signal = save_stats_signal
        self.train_dl = train_dl
        self.val_dl = val_dl

    def after_batch(self):

        #if not self.training:
        #    return

        if self.gui_callback:
            stats = f"Epoch {self.epoch+1} Batch {self.iter}"

            self.gui_callback(
                self.epoch,
                0.0,
                0.0,
                0.0,
                0.0,
                stats,
                False,
                None
            )

    def before_validate(self):
        if self.gui_callback:
            stats = f"Начало валидации..."

            self.gui_callback(
                self.epoch,
                0.0,
                0.0,
                0.0,
                0.0,
                stats,
                False,
                None
            )
    def after_epoch(self):

        # log — это текущая строка метрик: [epoch, train_loss, valid_loss, time]
        log_values = self.learn.recorder.log
        
        try:
            # log_values: индекс 0 - эпоха, 1 - train_loss, 2 - valid_loss
            train_loss = log_values[1]
            val_loss = log_values[2]
        except (IndexError, TypeError):
            # Если лог еще не сформирован, берем из smooth_loss
            train_loss = getattr(self.learn, 'loss', torch.tensor(0)).item()
            val_loss = None
        base_model = self.learn.model.model.model
        model = DetBenchPredict(base_model).to(DEVICE)
        model.eval()

        stats = f"Идёт расчёт метрик точности..."
        self.gui_callback(
            0,
            0,
            0,
            0,
            0,
            stats,
            False,
            None
        )

        train_p, train_r, train_f1, train_matrix = compute_metrics_and_matrix(
            model,
            self.train_dl,
            NUM_CLASSES,
            epoch=self.epoch,
            save_cam=False
        )

        global SnapshotCounter
        SnapshotCounter = 0

        val_p, val_r, val_f1, val_matrix = compute_metrics_and_matrix(
            model,
            self.val_dl,
            NUM_CLASSES,
            epoch=self.epoch,
            save_cam=True
        )

        stats = f"== ИТОГИ ЭПОХИ {self.epoch+1} ==\nTrain Loss: {train_loss:.4f}\nVal Loss: {val_loss:.4f}\nTrain Accuracy: {train_f1:.4f}\nVal Accuracy: {val_f1:.4f}"
        self.gui_callback(
            self.epoch,
            train_loss,
            val_loss,
            train_f1,
            val_f1,
            stats,
            True,
            val_matrix
        )

        global best_f1
        if val_f1 > best_f1:
            best_f1 = val_f1

            global model_save_name_buf
            model_save_name = model_save_name_buf + f"_{val_f1}"

            save_dir = os.path.join('Weights', model_save_name_buf)

            os.makedirs(save_dir, exist_ok=True)

            save_path = os.path.join(save_dir, model_save_name + ".pth")

            export_model = base_model
            export_model.eval()

            checkpoint = {
            "model_name": model_name,
            "num_classes": NUM_CLASSES,
            "image_size": IMG_SIZE,
            "model_state_dict": export_model.state_dict()
            }

            torch.save(checkpoint, save_path)
            self.save_stats_signal(save_path, model_save_name)


def collate_fn(batch):
    images, targets = zip(*batch)
    images = torch.stack(images, 0)
    
    # Проверка структуры первого таргета для отладки
    if 'boxes' not in targets[0]:
        print(f"DEBUG: Ключи в таргете: {targets[0].keys()}")
    
    # Безопасное получение количества объектов
    # Если 'bbox' нет, считаем, что объектов 0
    max_objs = max([t.get('boxes', torch.zeros((0,4))).shape[0] for t in targets])
    max_objs = max(max_objs, 1)
    
    batch_size = len(targets)
    
    padded_bboxes = torch.zeros((batch_size, max_objs, 4), dtype=torch.float32)
    padded_cls = torch.full((batch_size, max_objs), -1, dtype=torch.float32)

    for i, t in enumerate(targets):
        # Используется .get() чтобы не падать при отсутствии ключа
        bboxes = t.get('boxes', None)
        cls = t.get('labels', None)
        
        if bboxes is not None and bboxes.shape[0] > 0:
            num_objs = bboxes.shape[0]
            yxyx_boxes = bboxes.clone()
            yxyx_boxes[:, [0, 1, 2, 3]] = bboxes[:, [1, 0, 3, 2]]

            padded_bboxes[i, :num_objs] = yxyx_boxes
            padded_cls[i, :num_objs] = cls
            
    packed_targets = {
        'bbox': padded_bboxes,
        'cls': padded_cls,
        'img_scale': torch.ones(batch_size, dtype=torch.float32),
        'img_size': torch.tensor([(IMG_SIZE, IMG_SIZE)] * batch_size, dtype=torch.float32)
    }
    
    return (images, packed_targets), torch.tensor([0])

# =========================
# ФУНКЦИЯ ОБУЧЕНИЯ
# =========================
def train_efficientdet_divided(data_dir, epochs, model_save_name, training_mode, model_path, buf_count_epochs,  var_dataset, callback=None, save_stats_signal = None):
    
    global count_epochs
    count_epochs = buf_count_epochs

    global Var_dataset
    Var_dataset = var_dataset

    global NUM_CLASSES
    NUM_CLASSES = set_objects[Var_dataset][0]

    stats = f"Идёт формирование датасета..."
    callback(
        0,
        0,
        0,
        0,
        0,
        stats,
        False,
        None
    )

    train_ds = XRayDataset(
        f"{data_dir}/train/images",
        f"{data_dir}/train/labels"
    )

    val_ds = XRayDataset(
        f"{data_dir}/val/images",
        f"{data_dir}/val/labels"
    )

    train_dl = FDataLoader(
        train_ds,
        bs=batch_size,
        shuffle=True,
        num_workers=4,
        create_batch=collate_fn
    )

    valid_dl = FDataLoader(
        val_ds,
        bs=batch_size,
        shuffle=False,
        num_workers=2,
        create_batch=collate_fn
    )

    train_dl_lr = FDataLoader(
        train_ds,
        bs=batch_size_lr,
        shuffle=True,
        num_workers=4,
        create_batch=collate_fn
    )

    valid_dl_lr = FDataLoader(
        val_ds,
        bs=batch_size_lr,
        shuffle=False,
        num_workers=2,
        create_batch=collate_fn
    )

    dls = DataLoaders(train_dl, valid_dl, device=DEVICE)

    dls_lr = DataLoaders(train_dl_lr, valid_dl_lr, device=DEVICE)

    global model_save_name_buf
    model_save_name_buf = model_save_name

    base_model = EffDetTrain(NUM_CLASSES, (None if training_mode else model_path))

    if training_mode:
        save_path = os.path.join('Weights/', model_save_name + ".pth")         # новый файл
    else:
        dir_path = os.path.dirname(model_path)          # папка
        save_path = os.path.join(dir_path, model_save_name + ".pth")

    kornia_aug = KorniaAugment()

    def effdet_loss(model_output, target):
        # Если модель вернула словарь
        if isinstance(model_output, dict):
            return model_output['loss']
        # Если модель вернула только лосс
        return model_output

    learn_lr = Learner(
        dls_lr,
        base_model,
        opt_func = Adam,
        loss_func=effdet_loss,
        cbs=[
            KorniaAugmentCallback(kornia_aug),
            EarlyStoppingCallback(monitor='train_loss', patience=2),
            CSVLogger(),
            GUICallback(callback, save_stats_signal, train_dl, valid_dl)
        ]
    ).to_fp16()

    stats = f"Идёт поиск оптимальной скорости обучения..."
    callback(
        0,
        0,
        0,
        0,
        0,
        stats,
        False,
        None
    )

    # В функции train_efficientdet
    batch = learn_lr.dls.train.one_batch()

    # Проверка первого элемента
    inputs = batch[0]
    if isinstance(inputs, (list, tuple)):
        # Если это список/кортеж, берем первый тензор из него
        print(f"DEBUG: Вход батча — это {type(inputs)} из {len(inputs)} элементов")
        print(f"DEBUG: Формат первого тензора входа: {inputs[0].shape}")
    else:
        # Если это одиночный тензор
        print(f"DEBUG: Формат тензора входа: {inputs.shape}")

    # Проверка предсказаний
    targets = batch[1]
    if isinstance(targets, (list, tuple)):
        print(f"DEBUG: Цели — это список из {len(targets)} тензоров")
        for i, t in enumerate(targets):
            print(f"DEBUG: Тензор цели {i} имеет форму: {t.shape}")

    # 1. Поиск оптимальной LR
    print("Поиск оптимальной скорости обучения...")
    optimal_lr = 0.000003 
    try:
        # num_it=30 достаточно, чтобы понять тренд
        lr_find_results = learn_lr.lr_find(show_plot=False, start_lr=1e-7, end_lr=0.5, num_it=30)
        
        # Проверка на пустой результат
        if hasattr(lr_find_results, 'valley') and lr_find_results.valley is not None:
            optimal_lr = lr_find_results.valley
            print(f"Успешно: найден оптимальный LR: {optimal_lr:.6f}")
        else:
            print("Предупреждение: Точка valley не найдена. Используется 1e-3.")
            
    except Exception as e:
        print(f"Ошибка в lr_find: {e}")
        print(f"Используется стандартная скорость обучения: {optimal_lr}")

    stats = f"Начинается обучение..."
    callback(
        0,
        0,
        0,
        0,
        0,
        stats,
        False,
        None
    )

    learn = Learner(
        dls,
        base_model,
        opt_func = Adam,
        loss_func=effdet_loss,
        cbs=[
            KorniaAugmentCallback(kornia_aug),
            EarlyStoppingCallback(monitor='train_loss', patience=5),
            CSVLogger(),
            GUICallback(callback, save_stats_signal, train_dl, valid_dl)
        ]
    ).to_fp16()

    learn.fit_one_cycle(epochs, lr_max = optimal_lr)
    # СОХРАНЕНИЕ И ВОЗВРАТ

    export_model = base_model.model.model
    export_model.eval()

    checkpoint = {
    "model_name": model_name,
    "num_classes": NUM_CLASSES,
    "image_size": IMG_SIZE,
    "model_state_dict": export_model.state_dict()
    }

    torch.save(checkpoint, save_path)
    
    return save_path

# =========================
# ФУНКЦИИ ДЛЯ ВЫЧИСЛЕНИЯ ТОЧНОСТИ
# =========================
def calculate_iou(box1, box2):
    """Вычисляет IoU между двумя bounding boxes в формате [x1, y1, x2, y2]"""
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])
    
    intersection = max(0, x2 - x1) * max(0, y2 - y1)
    
    area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
    area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
    union = area1 + area2 - intersection
    
    return intersection / union if union > 0 else 0


def calculate_detection_metrics(
    model,
    data_loader,
    device,
    conf_thresh=0.3,
    iou_thresh=0.5
):
    model.eval()
    TP = FP = FN = 0

    with torch.no_grad():
        # Распаковка согласно структуре из collate_fn
        for (images, targets_dict), _ in data_loader:
            images = images.to(device)

            outputs = model(images) 

            for i in range(images.shape[0]):
                # Извлчение реальных боксов (где класс != -1)
                labels_mask = targets_dict['cls'][i] != -1
                gt_boxes_raw = targets_dict['bbox'][i][labels_mask].cpu().numpy()
                
                # Конвертация обратно в XYXY для сравнения с предсказаниями
                gt_boxes = gt_boxes_raw.copy()
                if len(gt_boxes) > 0:
                    gt_boxes[:, [0, 1, 2, 3]] = gt_boxes_raw[:, [1, 0, 3, 2]]

                # Предсказания для i-го изображения
                pred = outputs[i].detach().cpu().numpy() 

                # Фильтр по порогу уверенности
                scores = pred[:, 4]
                keep = scores >= conf_thresh
                valid_preds = pred[keep]

                if len(valid_preds) == 0:
                    FN += len(gt_boxes)
                    continue

                pred_boxes = valid_preds[:, :4] 
                matched_gt = set()

                for pb in pred_boxes:
                    found_match = False
                    for i_gt, gb in enumerate(gt_boxes):
                        if i_gt in matched_gt:
                            continue
                        
                        if calculate_iou(pb, gb) >= iou_thresh:
                            TP += 1
                            matched_gt.add(i_gt)
                            found_match = True
                            break
                    
                    if not found_match:
                        FP += 1

                FN += len(gt_boxes) - len(matched_gt)
                
    precision = TP / (TP + FP + 1e-6)
    recall    = TP / (TP + FN + 1e-6)
    f1        = 2 * precision * recall / (precision + recall + 1e-6)

    return precision, recall, f1

def compute_metrics_and_matrix(model, data_loader, num_classes=13, iou_threshold=0.2, 
                            conf_threshold=0.3, epoch=0, save_cam=False):
    model.eval()
    cm = np.zeros((num_classes + 1, num_classes + 1))
    cam_saved_this_epoch = False
    min_iou = 1
    max_iou = 0
    with torch.no_grad():
        for (imgs, targets), _ in data_loader:
            imgs = imgs.to(DEVICE)
            
            # --- БЛОК ГЕНЕРАЦИИ GRAD-CAM ---
            if save_cam and not cam_saved_this_epoch:
                # Проверка наличия объектов в батче. 
                # Проверка наличия дефекта или эталона
                has_pore = (targets['cls'] == 1).any()

                if has_pore:
                    # Поиск индексов изображений в батче, на которых есть хотя бы один объект
                    # targets['cls'] имеет форму [batch_size, max_num_bboxes]
                    pos_mask = (targets['cls'] == 1).any(dim=1)
                    pos_indices = torch.where(pos_mask)[0]

                    if len(pos_indices) > 0:
                        with torch.set_grad_enabled(True):
                            # Одбор изображения с дефектами
                            imgs_with_defects = imgs[pos_indices]
                            # Генерация CAM
                            save_gradcam_images(model.model, imgs_with_defects, epoch)
                        
                        #cam_saved_this_epoch = True 
            # -------------------------------
            outputs = model(imgs)
            
            for i in range(imgs.shape[0]):
                p_out = outputs[i].cpu().numpy()
                p_out = p_out[p_out[:, 4] > conf_threshold]
                
                t_cls_all = targets['cls'][i].cpu().numpy()
                valid_mask = t_cls_all != -1
                t_boxes = targets['bbox'][i][valid_mask].cpu().numpy()
                
                # Конвертация yxyx -> xyxy для IoU
                t_boxes_xyxy = t_boxes.copy()
                if len(t_boxes) > 0:
                    t_boxes_xyxy[:, [0, 1, 2, 3]] = t_boxes[:, [1, 0, 3, 2]]
                
                t_classes = t_cls_all[valid_mask].astype(int)
                matched_t_indices = set()
                
                # Проход по предсказаниям (формирование TP и FP)
                for pred in p_out:
                    p_box, p_idx = pred[:4], int(pred[5])
                    best_iou, best_t_idx = 0, -1
                    
                    for t_idx, t_box in enumerate(t_boxes_xyxy):
                        if t_idx in matched_t_indices: continue
                        iou = calculate_iou(p_box, t_box)
                        if(p_idx == 1 and t_classes[t_idx] == 1):
                            min_iou = min(min_iou, iou)
                            max_iou = max(max_iou, iou)
                        if iou > best_iou:
                            best_iou, best_t_idx = iou, t_idx
                    
                    if best_iou >= iou_threshold:
                        t_idx_cls = t_classes[best_t_idx]
                        cm[t_idx_cls-1, p_idx-1] += 1
                        matched_t_indices.add(best_t_idx)
                    else:
                        # FP: Предсказание класса на пустом месте
                        cm[num_classes, p_idx-1] += 1
                
                # Проход по ненайденным объектам (формирование FN)
                for t_idx in range(len(t_boxes)):
                    if t_idx not in matched_t_indices:
                        t_idx_cls = t_classes[t_idx]
                        cm[t_idx_cls-1, num_classes] += 1

    # Извлечение общих метрик из матрицы
    # TP - диагональ (без учета фона)
    tp_total = np.sum(np.diag(cm)[:num_classes])
    # FP - вся таблица с фоном по реальным объектам, без фона по предсказанным и без диагонали
    fp_total = np.sum(cm[:num_classes + 1, :num_classes]) - tp_total
    # FN - вся таблица с фоном по предсказанным объектам, без фона по реальным и без диагонали
    fn_total = np.sum(cm[:num_classes, :num_classes + 1]) - tp_total

    precision = tp_total / (tp_total + fp_total + 1e-6)
    recall = tp_total / (tp_total + fn_total + 1e-6)
    f1 = 2 * (precision * recall) / (precision + recall + 1e-6)

    print(f"Min IoU: {min_iou:.6f}, Max IoU: {max_iou:.6f}")

    return precision, recall, f1, cm

def denormalize(tensor):
    """Функция производит обртаную нормализавцию"""
    mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1).to(tensor.device)
    std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1).to(tensor.device)
    res = tensor * std + mean
    return torch.clamp(res, 0, 1)

def save_gradcam_images(model, images, epoch, output_dir="Interpretations"):
    """
    Генерирует Grad-CAM для батча изображений и сохраняет их.
    """
    global SnapshotCounter
    global count_epochs
    # Создаем папку для текущей эпохи
    epoch_dir = os.path.join(output_dir + "/" + model_save_name_buf + "/gradcam" , f"epoch_{epoch+1+count_epochs}")
    os.makedirs(epoch_dir, exist_ok=True)

    model.eval()
    
    # Выбор последнего слоя 
    cam_extractor = GradCAM(model, target_layer='backbone.blocks.6.1')

    # Првые 3 изображения из батча
    num_to_save = min(len(images), 3)
    
    for i in range(num_to_save):
        input_tensor = images[i].unsqueeze(0).to(next(model.parameters()).device) # (1, 3, IMG_SIZE, IMG_SIZE)
        with torch.set_grad_enabled(True):
            # Прогон через модель
            out = model(input_tensor)
            
            # 1. Если это кортеж или список, распоковка до первого тензора
            if isinstance(out, (list, tuple)):
                temp_out = out[0]
                while isinstance(temp_out, (list, tuple)):
                    temp_out = temp_out[0]
                # Преобразование скаляра в тензор (1, 1), чтобы torchcam мог его индексировать
                score_to_backprop = temp_out.sum().view(1, 1)
                #print("Сработал первый парсер Grad-CAM") !
            
            # 2. Если это тензор с результатами детекции (DetBenchPredict)
            elif hasattr(out, 'shape') and len(out.shape) > 2 and out.shape[1] > 0:
                # Берем максимальную уверенность и делаем её тензором (1, 1)
                score_to_backprop = out[:, :, 4].max().view(1, 1)
                #print("Сработал второй парсер Grad-CAM")
            
            # 3. Резервный вариант
            else:
                res_tensor = out.sum() if torch.is_tensor(out) else torch.tensor(0.0, requires_grad=True)
                score_to_backprop = res_tensor.view(1, 1)
                #print("Сработал третий парсер Grad-CAM")

            # Извлечение карты активации (на основе лосса или максимального выхода)
            activation_map = cam_extractor(class_idx=1, scores=score_to_backprop) 
        
        # Подготовка оригинального изображения
        orig_img_tensor = denormalize(images[i])
        orig_pil = to_pil_image(orig_img_tensor)
        
        # Наложение маски
        result = overlay_mask(orig_pil, to_pil_image(activation_map[0], mode='F'), alpha=0.5)
        
        # Сохранение
        orig_pil.save(os.path.join(epoch_dir, f"sample_{SnapshotCounter}_original.jpg"))
        result.save(os.path.join(epoch_dir, f"sample_{SnapshotCounter}_heatmap.jpg"))

        SnapshotCounter+=1

    # Очистка экстрактора
    cam_extractor.remove_hooks()

# =========================
# ФУНКЦИИ ИНТЕРФЕЙСА
# =========================
def load_model_for_inference(model_path):
    print(f"Загрузка модели из {model_path}...")

    try:
        checkpoint = torch.load(
            model_path,
            map_location=DEVICE,
            weights_only=False
        )

        # 1. Создаётся BASE EfficientDet
        base_model = create_model(
            checkpoint["model_name"],
            pretrained=False,
            num_classes=checkpoint["num_classes"],
            image_size=(checkpoint["image_size"], checkpoint["image_size"])
        )

        # 2. Загружае.тся веса в base_model
        base_model.load_state_dict(checkpoint["model_state_dict"], strict=True)

        # 3. Оборачивается в DetBenchPredict
        model = DetBenchPredict(base_model)
        model = model.to(DEVICE)
        model.eval()

        print("PyTorch модель загружена успешно (DetBenchPredict)")
        return model

    except Exception as e:
        print(f"Ошибка загрузки модели: {e}")
        traceback.print_exc()
        return None



def inference_single_image(model_path, image_path, confidence_threshold=0.3):
    model = load_model_for_inference(model_path)
    if model is None:
        return []
    
    try:
        img_bytes = np.fromfile(image_path, dtype=np.uint8)
        image = cv2.imdecode(img_bytes, cv2.IMREAD_GRAYSCALE)
        if image is None: return []

        original_h, original_w = image.shape[:2]
        image_rgb = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)

        # 1. Расчет сетки тайлов
        number_w = original_w / IMG_SIZE
        number_h = original_h / IMG_SIZE

        # Список координат левых верхних углов (x, y)
        tile_coords = []
        
        # Опредется шаг по X и Y
        # Используется ceil, чтобы покрыть всё изображение
        cols = int(np.ceil(number_w))
        rows = int(np.ceil(number_h))

        for row in range(rows):
            for col in range(cols):
                # Рассчитываем координату X
                curr_x = col * IMG_SIZE
                # Если это последний столбец и он "вылезает" за границу, 
                # он сдвигается влево, чтобы он закончился ровно по краю
                if curr_x + IMG_SIZE > original_w:
                    curr_x = original_w - IMG_SIZE
                
                # Рассчитывается координата Y
                curr_y = row * IMG_SIZE
                if curr_y + IMG_SIZE > original_h:
                    curr_y = original_h - IMG_SIZE
                
                # Добавляется только уникальные координаты (на случай маленьких картинок)
                coord = (max(0, int(curr_x)), max(0, int(curr_y)))
                if coord not in tile_coords:
                    tile_coords.append(coord)

        all_results_buf = []

        # 2. Цикл инференса по каждой части (тайлу)
        for tx, ty in tile_coords:
            # Вырезается часть изображения
            tile = image_rgb[ty:ty+IMG_SIZE, tx:tx+IMG_SIZE]
            
            # Подготовка тензора для конкретного тайла
            tile_prep = tile.astype(np.float32) / 255.0
            mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
            std  = np.array([0.229, 0.224, 0.225], dtype=np.float32)
            tile_prep = (tile_prep - mean) / std
            
            tile_tensor = torch.from_numpy(tile_prep).permute(2, 0, 1).unsqueeze(0).float().to(DEVICE)
            
            # Инференс
            if isinstance(model, ort.InferenceSession):
                input_name = model.get_inputs()[0].name
                outputs = model.run(None, {input_name: tile_tensor.cpu().numpy()})
                detections = outputs[0]
            else:
                with torch.no_grad():
                    detections = model(tile_tensor).cpu().numpy()

            # 3. Обработка результатов тайла
            if detections is not None and detections.size > 0:
                dets = np.squeeze(detections)
                if dets.ndim == 1: dets = dets[np.newaxis, :]

                for det in dets:
                    print(det.shape)
                    print(det)
                    score = det[4].item() if hasattr(det[4], 'item') else float(det[4])
                    if score >= confidence_threshold:
                        # Координаты внутри тайла (0..IMG_SIZE)
                        x1_loc, y1_loc, x2_loc, y2_loc = det[0:4]
                        cls_id = int(det[5])

                        # Пересчет координат в глобальные
                        x1_global = int(x1_loc + tx)
                        y1_global = int(y1_loc + ty)
                        x2_global = int(x2_loc + tx)
                        y2_global = int(y2_loc + ty)

                        all_results_buf.append({
                            'bbox': [x1_global, y1_global, x2_global, y2_global],
                            'score': score,
                            'label': names.get(cls_id, f"class_{cls_id}")
                        })
        all_results = combining_objects(all_results_buf)
        print(f"Обработано тайлов: {len(tile_coords)}, найдено всего объектов: {len( all_results)}")
        return  all_results

    except Exception as e:
        print(f"Ошибка: {e}")
        traceback.print_exc()
        return []

def build_targets(boxes, labels, device):
    batch_size = len(boxes)
    max_targets = max([b.shape[0] for b in boxes]) if len(boxes) > 0 else 0
    max_targets = max(max_targets, 1)

    combined_boxes = torch.zeros((batch_size, max_targets, 4), device=device)
    combined_labels = torch.full((batch_size, max_targets), -1, dtype=torch.float32, device=device)

    for i in range(batch_size):
        num_objs = boxes[i].shape[0]
        if num_objs > 0:
            # Формат y1, x1, y2, x2
            yxyx_boxes = boxes[i][:, [1, 0, 3, 2]]
            combined_boxes[i, :num_objs] = yxyx_boxes
            combined_labels[i, :num_objs] = labels[i].to(device)

    return {
        "bbox": combined_boxes,
        "cls": combined_labels,
        "img_scale": torch.tensor([1.0] * batch_size, device=device, dtype=torch.float32),
        "img_size": torch.tensor([(IMG_SIZE, IMG_SIZE)] * batch_size, device=device, dtype=torch.float32),
    }

def combine_two_boxes(b1, b2):
    return [
        min(b1[0], b2[0]),
        min(b1[1], b2[1]),
        max(b1[2], b2[2]),
        max(b1[3], b2[3])
    ]


def box_distance(b1, b2):
    dx = max(b2[0] - b1[2], b1[0] - b2[2], 0)
    dy = max(b2[1] - b1[3], b1[1] - b2[3], 0)
    return max(dx, dy)


def combining_objects(all_results_buf, iou_thr=0.3, dist_thr=20):
    merged = True
    results = all_results_buf.copy()

    while merged:
        merged = False
        new_results = []
        used = [False] * len(results)

        for i in range(len(results)):
            if used[i]:
                continue

            curr = results[i]
            curr_box = curr["bbox"]

            for j in range(i + 1, len(results)):
                if used[j]:
                    continue

                other = results[j]

                # объединяются одинаковые классы
                if curr["label"] != other["label"]:
                    continue

                iou = calculate_iou(curr_box, other["bbox"])
                dist = box_distance(curr_box, other["bbox"])

                if iou > iou_thr or dist < dist_thr:
                    curr_box = combine_two_boxes(curr_box, other["bbox"])
                    curr["score"] = max(curr["score"], other["score"])
                    used[j] = True
                    merged = True

            curr["bbox"] = curr_box
            used[i] = True
            new_results.append(curr)

        results = new_results

    return results

def load_model_train_further_education(pth_path, device):
    checkpoint = torch.load(pth_path, map_location=device)

    # 1. Создаём базовую модель ТОЧНО КАК ПРИ ОБУЧЕНИИ
    base_model = create_model(
        checkpoint["model_name"],
        pretrained=False,              # ВАЖНО: False
        num_classes=checkpoint["num_classes"],
        image_size=(checkpoint["image_size"], checkpoint["image_size"]),
    )

    # 2. Загружаем веса
    base_model.load_state_dict(checkpoint["model_state_dict"])

    return base_model

def load_model_train_new():
    base_model = create_model(
        model_name,
        pretrained=True,
        num_classes=NUM_CLASSES,
        image_size=(IMG_SIZE, IMG_SIZE),
    )

    return base_model