pokouqiege/models/line_detect/trainer.py

import os
import time
from datetime import datetime

import numpy as np
import torch
from matplotlib import pyplot as plt
from torch.utils.tensorboard import SummaryWriter

from libs.vision_libs.utils import draw_bounding_boxes, draw_keypoints
from models.base.base_model import BaseModel
from models.base.base_trainer import BaseTrainer
from models.config.config_tool import read_yaml
from models.line_detect.line_dataset import LineDataset
from models.line_net.dataset_LD import WirePointDataset
from models.wirenet.postprocess import postprocess
from tools import utils
from torchvision import transforms
import matplotlib as mpl

cmap = plt.get_cmap("jet")
norm = mpl.colors.Normalize(vmin=0.4, vmax=1.0)
sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])
def _loss(losses):
    total_loss = 0
    for i in losses.keys():
        if i != "loss_wirepoint":
            total_loss += losses[i]
        else:
            loss_labels = losses[i]["losses"]
    loss_labels_k = list(loss_labels[0].keys())
    for j, name in enumerate(loss_labels_k):
        loss = loss_labels[0][name].mean()
        total_loss += loss
    return total_loss
def c(x):
    return sm.to_rgba(x)


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

import matplotlib.pyplot as plt
from PIL import ImageDraw
from torchvision.transforms import functional as F
import torch


# 由低到高蓝黄红
def draw_lines_with_scores(tensor_image, lines, scores, width=3, cmap='viridis'):
    """
    根据得分对线段着色并绘制
    :param tensor_image: (3, H, W) uint8 图像
    :param lines: (N, 2, 2) 每条线 [ [x1,y1], [x2,y2] ]
    :param scores: (N,) 每条线的得分，范围 [0, 1]
    :param width: 线宽
    :param cmap: matplotlib colormap 名称，例如 'viridis', 'jet', 'coolwarm'
    :return: (3, H, W) uint8 画好线的图像
    """
    assert tensor_image.dtype == torch.uint8
    assert tensor_image.shape[0] == 3
    assert lines.shape[0] == scores.shape[0]

    # 准备色图
    colormap = plt.get_cmap(cmap)
    colors = (colormap(scores.cpu().numpy())[:, :3] * 255).astype('uint8')  # 去掉 alpha 通道

    # 转为 PIL 画图
    image_pil = F.to_pil_image(tensor_image)
    draw = ImageDraw.Draw(image_pil)

    for line, color in zip(lines, colors):
        start = tuple(map(float, line[0][:2].tolist()))
        end = tuple(map(float, line[1][:2].tolist()))
        draw.line([start, end], fill=tuple(color), width=width)

    return (F.to_tensor(image_pil) * 255).to(torch.uint8)


class Trainer(BaseTrainer):
    def __init__(self, model=None, **kwargs):
        super().__init__(model, device, **kwargs)
        self.model = model
        # print(f'kwargs:{kwargs}')
        self.init_params(**kwargs)

    def init_params(self, **kwargs):
        if kwargs != {}:
            print(f'train_params:{kwargs["train_params"]}')
            self.freeze_config = kwargs['train_params']['freeze_params']
            print(f'freeze_config:{self.freeze_config}')
            self.dataset_path = kwargs['io']['datadir']
            self.batch_size = kwargs['train_params']['batch_size']
            self.num_workers = kwargs['train_params']['num_workers']
            self.logdir = kwargs['io']['logdir']
            self.resume_from = kwargs['train_params']['resume_from']
            self.optim = ''
            self.train_result_ptath = os.path.join(self.logdir, datetime.now().strftime("%Y%m%d_%H%M%S"))
            self.wts_path = os.path.join(self.train_result_ptath, 'weights')
            self.tb_path = os.path.join(self.train_result_ptath, 'logs')
            self.writer = SummaryWriter(self.tb_path)
            self.last_model_path = os.path.join(self.wts_path, 'last.pth')
            self.best_train_model_path = os.path.join(self.wts_path, 'best_train.pth')
            self.best_val_model_path = os.path.join(self.wts_path, 'best_val.pth')
            self.max_epoch = kwargs['train_params']['max_epoch']

    def move_to_device(self, data, device):
        if isinstance(data, (list, tuple)):
            return type(data)(self.move_to_device(item, device) for item in data)
        elif isinstance(data, dict):
            return {key: self.move_to_device(value, device) for key, value in data.items()}
        elif isinstance(data, torch.Tensor):
            return data.to(device)
        else:
            return data  # 对于非张量类型的数据不做任何改变

    def freeze_params(self, model):
        """根据配置冻结模型参数"""
        default_config = {
            'backbone': True,  # 冻结 backbone
            'rpn': False,  # 不冻结 rpn
            'roi_heads': {
                'box_head': False,
                'box_predictor': False,
                'line_head': False,
                'line_predictor': {
                    'fc1': False,
                    'fc2': {
                        '0': False,
                        '2': False,
                        '4': False
                    }
                }
            }
        }

        # 更新默认配置
        default_config.update(self.freeze_config)
        config = default_config

        print("\n===== Parameter Freezing Configuration =====")
        for name, module in model.named_children():
            if name in config:
                if isinstance(config[name], bool):
                    for param in module.parameters():
                        param.requires_grad = not config[name]
                    print(f"{'Frozen' if config[name] else 'Trainable'} module: {name}")

                elif isinstance(config[name], dict):
                    for subname, submodule in module.named_children():
                        if subname in config[name]:
                            if isinstance(config[name][subname], bool):
                                for param in submodule.parameters():
                                    param.requires_grad = not config[name][subname]
                                print(
                                    f"{'Frozen' if config[name][subname] else 'Trainable'} submodule: {name}.{subname}")

                            elif isinstance(config[name][subname], dict):
                                for subsubname, subsubmodule in submodule.named_children():
                                    if subsubname in config[name][subname]:
                                        for param in subsubmodule.parameters():
                                            param.requires_grad = not config[name][subname][subsubname]
                                        print(
                                            f"{'Frozen' if config[name][subname][subsubname] else 'Trainable'} sub-submodule: {name}.{subname}.{subsubname}")

        # 打印参数统计
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        print(f"\nTotal Parameters: {total_params:,}")
        print(f"Trainable Parameters: {trainable_params:,}")
        print(f"Frozen Parameters: {total_params - trainable_params:,}")

    def load_best_model(self, model, optimizer, save_path, device):
        if os.path.exists(save_path):
            checkpoint = torch.load(save_path, map_location=device)
            model.load_state_dict(checkpoint['model_state_dict'])
            if optimizer is not None:
                optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
            epoch = checkpoint['epoch']
            loss = checkpoint['loss']
            print(f"Loaded best model from {save_path} at epoch {epoch} with loss {loss:.4f}")
        else:
            print(f"No saved model found at {save_path}")
        return model, optimizer





    def writer_predict_result(self, img, result, epoch):
        img = img.cpu().detach()
        im = img.permute(1, 2, 0)  # [512, 512, 3]
        self.writer.add_image("z-ori", im, epoch, dataformats="HWC")

        boxed_image = draw_bounding_boxes((img * 255).to(torch.uint8), result["boxes"],
                                          colors="yellow", width=1)

        # plt.imshow(boxed_image.permute(1, 2, 0).detach().cpu().numpy())
        # plt.show()

        self.writer.add_image("z-obj", boxed_image.permute(1, 2, 0), epoch, dataformats="HWC")
        keypoint_img = draw_keypoints(boxed_image, result['lines'], colors='red', width=3)

        self.writer.add_image("z-output", keypoint_img, epoch)
        print("lines shape:", result['lines'].shape)

        # 用自己写的函数画线段
        # line_image = draw_lines(boxed_image, result['lines'], color='red', width=3)
        print(f"shape of linescore:{result['liness_scores'].shape}")
        scores = result['liness_scores'].mean(dim=1)  # shape: [31]

        line_image = draw_lines_with_scores((img * 255).to(torch.uint8),  result['lines'],scores, width=3, cmap='jet')

        self.writer.add_image("z-output_line", line_image.permute(1, 2, 0), epoch, dataformats="HWC")




    def writer_loss(self, losses, epoch, phase='train'):
        try:
            for key, value in losses.items():
                if key == 'loss_wirepoint':
                    for subdict in losses['loss_wirepoint']['losses']:
                        for subkey, subvalue in subdict.items():
                            self.writer.add_scalar(f'{phase}/loss/{subkey}',
                                                   subvalue.item() if hasattr(subvalue, 'item') else subvalue,
                                                   epoch)
                elif isinstance(value, torch.Tensor):
                    self.writer.add_scalar(f'{phase}/loss/{key}', value.item(), epoch)
        except Exception as e:
            print(f"TensorBoard logging error: {e}")

    def train_from_cfg(self, model: BaseModel, cfg, freeze_config=None):  # 新增：支持传入冻结配置
        cfg = read_yaml(cfg)
        # print(f'cfg:{cfg}')
        # self.freeze_config = freeze_config or {}  # 更新冻结配置

        self.train(model, **cfg)

    def train(self, model, **kwargs):

        self.init_params(**kwargs)

        dataset_train = LineDataset(dataset_path=self.dataset_path, dataset_type='train')
        dataset_val = LineDataset(dataset_path=self.dataset_path, dataset_type='val')

        train_sampler = torch.utils.data.RandomSampler(dataset_train)
        val_sampler = torch.utils.data.RandomSampler(dataset_val)
        train_batch_sampler = torch.utils.data.BatchSampler(train_sampler, batch_size=self.batch_size, drop_last=True)
        val_batch_sampler = torch.utils.data.BatchSampler(val_sampler, batch_size=self.batch_size, drop_last=True)
        train_collate_fn = utils.collate_fn
        val_collate_fn = utils.collate_fn

        data_loader_train = torch.utils.data.DataLoader(
            dataset_train, batch_sampler=train_batch_sampler, num_workers=self.num_workers, collate_fn=train_collate_fn
        )
        data_loader_val = torch.utils.data.DataLoader(
            dataset_val, batch_sampler=val_batch_sampler, num_workers=self.num_workers, collate_fn=val_collate_fn
        )

        model.to(device)

        optimizer = torch.optim.Adam(
            filter(lambda p: p.requires_grad, model.parameters()),
            lr=kwargs['train_params']['optim']['lr']
        )

        for epoch in range(self.max_epoch):
            print(f"train epoch:{epoch}")

            model, epoch_train_loss = self.one_epoch(model, data_loader_train, epoch, optimizer)

            # ========== Validation ==========
            with torch.no_grad():
                model, epoch_val_loss = self.one_epoch(model, data_loader_val, epoch, optimizer, phase='val')

            if epoch==0:
                best_train_loss = epoch_train_loss
                best_val_loss = epoch_val_loss

            self.save_last_model(model,self.last_model_path, epoch, optimizer)
            best_train_loss = self.save_best_model(model, self.best_train_model_path, epoch, epoch_train_loss,
                                                   best_train_loss,
                                                   optimizer)
            best_val_loss = self.save_best_model(model, self.best_val_model_path, epoch, epoch_val_loss, best_val_loss,
                                                 optimizer)

    def one_epoch(self, model, data_loader, epoch, optimizer, phase='train'):
        if phase == 'train':
            model.train()
        if phase == 'val':
            model.eval()

        total_loss = 0
        epoch_step = 0
        global_step = epoch * len(data_loader)
        for imgs, targets in data_loader:
            imgs = self.move_to_device(imgs, device)
            targets = self.move_to_device(targets, device)
            if phase== 'val':
                result,loss_dict = model(imgs, targets)
                losses = sum(loss_dict.values()) if loss_dict else torch.tensor(0.0, device=device)
                print(f'val losses:{losses}')
            else:
                loss_dict = model(imgs, targets)
                losses = sum(loss_dict.values()) if loss_dict else torch.tensor(0.0, device=device)
                print(f'train losses:{losses}')

            # loss = _loss(losses)
            loss=losses
            total_loss += loss.item()
            if phase == 'train':
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
            self.writer_loss(loss_dict, global_step, phase=phase)
            global_step += 1

            if epoch_step == 0 and phase == 'val':
                t_start = time.time()
                print(f'start to predict:{t_start}')
                result = model(self.move_to_device(imgs, self.device))
                # print(f'result:{result}')
                t_end = time.time()
                print(f'predict used:{t_end - t_start}')
                self.writer_predict_result(img=imgs[0], result=result[0], epoch=epoch)
                epoch_step+=1

        avg_loss = total_loss / len(data_loader)
        print(f'{phase}/loss epoch{epoch}:{avg_loss:4f}')
        self.writer.add_scalar(f'loss/{phase}', avg_loss, epoch)
        return model, avg_loss

    def save_best_model(self, model, save_path, epoch, current_loss, best_loss, optimizer=None):
        os.makedirs(os.path.dirname(save_path), exist_ok=True)

        if current_loss <= best_loss:
            checkpoint = {
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'loss': current_loss
            }
            if optimizer is not None:
                checkpoint['optimizer_state_dict'] = optimizer.state_dict()

            torch.save(checkpoint, save_path)
            print(f"Saved best model at epoch {epoch} with loss {current_loss:.4f}")

            return current_loss

        return best_loss

    def save_last_model(self, model, save_path, epoch, optimizer=None):

        if os.path.exists(f'{self.wts_path}/last.pt'):
            os.remove(f'{self.wts_path}/last.pt')

        os.makedirs(os.path.dirname(save_path), exist_ok=True)

        checkpoint = {
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
        }

        if optimizer is not None:
            checkpoint['optimizer_state_dict'] = optimizer.state_dict()

        torch.save(checkpoint, save_path)


if __name__ == '__main__':
    print('')