pokouqiege/models/line_detect/trainer.py

import os
import time
from datetime import datetime

import cv2
import numpy as np
import torch
from PIL.ImageDraw import ImageDraw
from matplotlib import pyplot as plt
from scipy.ndimage import gaussian_filter
from torch.optim.lr_scheduler import ReduceLROnPlateau
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
import torch.nn.functional as F


from tools import utils

import matplotlib as mpl

from utils.data_process.show_prams import print_params

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')



def draw_ellipses_on_image(image, masks_pred, threshold=0.5, color=(0, 255, 0), thickness=2):
    """
    在单张原始图像上绘制从 masks 拟合出的椭圆。
    自动将 masks resize 到 image 的空间尺寸。

    Args:
        image: Tensor [3, H_img, W_img] —— 原始图像（如 [3, 2000, 2000]）
        masks_pred: Tensor [N, 1, H_mask, W_mask] or [N, H_mask, W_mask] —— 模型输出 mask（如 [2, 1, 672, 672]）
        threshold: 二值化阈值
        color: BGR color for OpenCV
        thickness: ellipse line thickness

    Returns:
        drawn_image: numpy array [H_img, W_img, 3] in RGB
    """
    # Step 1: 标准化 masks_pred to [N, H, W]
    if masks_pred.ndim == 4:
        if masks_pred.shape[1] == 1:
            masks_pred = masks_pred.squeeze(1)  # [N, 1, H, W] -> [N, H, W]
        else:
            raise ValueError(f"Expected channel=1 in masks_pred, got shape {masks_pred.shape}")
    elif masks_pred.ndim != 3:
        raise ValueError(f"masks_pred must be 3D (N, H, W) or 4D (N, 1, H, W), got {masks_pred.shape}")

    N, H_mask, W_mask = masks_pred.shape
    C, H_img, W_img = image.shape

    # Step 2: Resize masks to original image size using bilinear interpolation
    masks_resized = F.interpolate(
        masks_pred.unsqueeze(1).float(),  # [N, 1, H_mask, W_mask]
        size=(H_img, W_img),
        mode='bilinear',
        align_corners=False
    ).squeeze(1)  # [N, H_img, W_img]

    # Step 3: Convert image to numpy RGB
    img_tensor = image.detach().cpu()
    if img_tensor.max() <= 1.0:
        img_np = (img_tensor * 255).byte().numpy()  # [3, H, W]
    else:
        img_np = img_tensor.byte().numpy()
    img_rgb = np.transpose(img_np, (1, 2, 0))  # [H, W, 3]
    img_out = img_rgb.copy()

    # Step 4: Process each mask
    for mask in masks_resized:
        mask_cpu = mask.detach().cpu()
        mask_prob = torch.sigmoid(mask_cpu) if mask_cpu.min() < 0 else mask_cpu
        binary = (mask_prob > threshold).numpy().astype(np.uint8) * 255  # [H_img, W_img]

        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if contours:
            largest_contour = max(contours, key=cv2.contourArea)
            if len(largest_contour) >= 5:
                try:
                    ellipse = cv2.fitEllipse(largest_contour)
                    img_bgr = cv2.cvtColor(img_out, cv2.COLOR_RGB2BGR)
                    cv2.ellipse(img_bgr, ellipse, color=color, thickness=thickness)
                    img_out = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
                except cv2.error as e:
                    print(f"Warning: Failed to fit ellipse: {e}")

    return img_out


def fit_circle(points):
    """
    Fit a circle to a set of points (at least 3).

    Args:
        points: torch.Tensor 或 numpy array, shape (N, 2)

    Returns:
        center (cx, cy), radius r
    """
    # 如果是 torch.Tensor，先转为 numpy
    if isinstance(points, torch.Tensor):
        if points.dim() == 3:
            points = points[0]  # 去掉 batch 维度
        points = points.detach().cpu().numpy()

    if not (isinstance(points, np.ndarray) and points.ndim == 2 and points.shape[1] == 2):
        raise ValueError(f"Expected points shape (N, 2), got {points.shape}")

    x = points[:, 0].astype(float)
    y = points[:, 1].astype(float)

    # 确保 A 是二维数组
    A = np.column_stack((x, y, np.ones_like(x)))  # 使用 column_stack 代替 stack 可能更清晰
    b = -(x ** 2 + y ** 2)

    try:
        sol, residuals, rank, s = np.linalg.lstsq(A, b, rcond=None)
    except np.linalg.LinAlgError as e:
        print(f"Linear algebra error occurred: {e}")
        raise ValueError("Could not fit circle to points.")

    D, E, F = sol

    cx = -D / 2.0
    cy = -E / 2.0
    r = np.sqrt(cx ** 2 + cy ** 2 - F)

    return (cx, cy), r
from PIL import ImageDraw, Image
import io

def draw_el(all, background_img):
    """
    all = [x_center, y_center, a, b, theta, x1, y1, x2, y2]
        theta: ellipse rotation (degrees)
        (x1, y1): start point
        (x2, y2): end point
    """

    if isinstance(all, torch.Tensor):
        all = all.cpu().numpy()

    # Unpack parameters
    cx, cy, a, b, theta_deg, x1, y1, x2, y2 = all
    theta = np.radians(theta_deg)

    # ====== Draw ellipse ======
    phi = np.linspace(0, np.pi * 2, 500)
    x_ellipse = cx + a * np.cos(phi) * np.cos(theta) - b * np.sin(phi) * np.sin(theta)
    y_ellipse = cy + a * np.cos(phi) * np.sin(theta) + b * np.sin(phi) * np.cos(theta)

    # ====== Draw image ======
    plt.figure(figsize=(10, 10))
    plt.imshow(background_img)

    # Ellipse
    plt.plot(x_ellipse, y_ellipse, 'b-', linewidth=2)

    # Center
    plt.plot(cx, cy, 'ko', markersize=8)

    # Start & End points (now real coordinates)
    plt.plot(x1, y1, 'ro', markersize=10)
    plt.plot(x2, y2, 'go', markersize=10)

    # ====== Convert to tensor ======
    buf = io.BytesIO()
    plt.savefig(buf, format='png', bbox_inches='tight')
    buf.seek(0)
    result_img = Image.open(buf).convert('RGB')
    img_tensor = torch.from_numpy(np.array(result_img)).permute(2, 0, 1)
    plt.close()

    return img_tensor

# from PIL import ImageDraw, Image
# import io
# # 绘制椭圆
# def draw_el(all, background_img):
#     # 解析椭圆参数
#     if isinstance(all, torch.Tensor):
#         all = all.cpu().numpy()
#     print_params(all)
#     x, y, a, b, q, q1, q2 = all
#     theta = np.radians(q)
#     phi1 = np.radians(q1)  # 第一个点的参数角
#     phi2 = np.radians(q2)  # 第二个点的参数角
#
#     # 生成椭圆上的点
#     phi = np.linspace(0, 2 * np.pi, 500)
#     x_ellipse = x + a * np.cos(phi) * np.cos(theta) - b * np.sin(phi) * np.sin(theta)
#     y_ellipse = y + a * np.cos(phi) * np.sin(theta) + b * np.sin(phi) * np.cos(theta)
#
#     # 计算两个指定点的坐标
#     def param_to_point(phi, xc, yc, a, b, theta):
#         x = xc + a * np.cos(phi) * np.cos(theta) - b * np.sin(phi) * np.sin(theta)
#         y = yc + a * np.cos(phi) * np.sin(theta) + b * np.sin(phi) * np.cos(theta)
#         return x, y
#
#     P1 = param_to_point(phi1, x, y, a, b, theta)
#     P2 = param_to_point(phi2, x, y, a, b, theta)
#
#     # 创建画布并显示背景图片（使用传入的background_img，shape为[H, W, C]）
#     plt.figure(figsize=(10, 10))
#     plt.imshow(background_img)  # 直接显示背景图
#
#     # 绘制椭圆及相关元素
#     plt.plot(x_ellipse, y_ellipse, 'b-', linewidth=2)
#     plt.plot(x, y, 'ko', markersize=8)
#     plt.plot(P1[0], P1[1], 'ro', markersize=10)
#     plt.plot(P2[0], P2[1], 'go', markersize=10)

    # 转换为TensorBoard所需的张量格式 [C, H, W]
    # buf = io.BytesIO()
    # plt.savefig(buf, format='png', bbox_inches='tight')
    # buf.seek(0)
    # result_img = Image.open(buf).convert('RGB')
    # img_tensor = torch.from_numpy(np.array(result_img)).permute(2, 0, 1)
    # plt.close()
    #
    # return img_tensor
# 由低到高蓝黄红
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.data_type = kwargs['io']['data_type']
            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']
            self.augmentation= kwargs['train_params']["augmentation"]

    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")


        if  'points' in result:
            keypoint_img = draw_keypoints(boxed_image, result['points'], colors='red', width=3)

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


        if 'lines' in result:
            # 用自己写的函数画线段
            # line_image = draw_lines(boxed_image, result['lines'], color='red', width=3)
            print(f"shape of linescore:{result['lines_scores'].shape}")
            scores = result['lines_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")

        if 'arcs' in result:
            arcs = result['arcs'][0]
            print(f'arcs in dra w:{arcs}')

            ellipse_img = draw_el(arcs, background_img=im)

            # img_rgb = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)
            #
            # img_tensor =torch.tensor(img_rgb)
            # img_tensor = np.transpose(img_tensor)

            self.writer.add_image('z-out-arc', ellipse_img, global_step=epoch)

        if 'ins_masks' in result:
            # points=result['circles']
            # points=points.squeeze(1)
            ppp=result['ins_masks']
            bbb=result['boxes']
            print(f'boxes shape:{bbb.shape}')
            print(f'ppp:{ppp.shape}')
            ins_masks = result['ins_masks']
            ins_masks = ins_masks.squeeze(1)
            print(f'ins_masks shape:{ins_masks.shape}')
            features = result['features']

            circle_image = img.cpu().numpy().transpose((1, 2, 0))  # CHW -> HWC
            circle_image = (circle_image * 255).clip(0, 255).astype(np.uint8)

            sum_mask = ins_masks.sum(dim=0, keepdim=True)
            sum_mask = sum_mask / (sum_mask.max() + 1e-8)

            # keypoint_img = draw_keypoints((img * 255).to(torch.uint8), points, colors='red', width=3)
            self.writer.add_image('z-ins-masks', sum_mask.squeeze(0), global_step=epoch)


            result_imgs = draw_ellipses_on_image(img, ins_masks, threshold=0.5)
            self.writer.add_image('z-out-ellipses', result_imgs, dataformats='HWC', global_step= epoch)

            features=self.apply_gaussian_blur_to_tensor(features,sigma=3)
            self.writer.add_image('z-feature', features, global_step=epoch)

            # cv2.imshow('arc', img_rgb)
            # cv2.waitKey(1000000)



    def normalize_tensor(self,tensor):
        """Normalize tensor to [0, 1]"""
        min_val = tensor.min()
        max_val = tensor.max()
        return (tensor - min_val) / (max_val - min_val)

    def apply_gaussian_blur_to_tensor(self,feature_map, sigma=3):
        """
        Apply Gaussian blur to a feature map and convert it into an RGB heatmap.

        :param feature_map: Tensor of shape (H, W) or (1, H, W)
        :param sigma: Standard deviation for Gaussian kernel
        :return: Tensor of shape (3, H, W) representing the RGB heatmap
        """
        if feature_map.dim() == 3:
            if feature_map.shape[0] != 1:
                raise ValueError("Only single-channel feature map supported.")
            feature_map = feature_map.squeeze(0)

        # Normalize to [0, 1]
        normalized_feat = self.normalize_tensor(feature_map).cpu().numpy()

        # Apply Gaussian blur
        blurred_feat = gaussian_filter(normalized_feat, sigma=sigma)

        # Convert to colormap (e.g., 'jet')
        colormap = plt.get_cmap('jet')
        colored = colormap(blurred_feat)  # shape: (H, W, 4) RGBA

        # Convert to (3, H, W), drop alpha channel
        colored_rgb = colored[:, :, :3]  # (H, W, 3)
        colored_tensor = torch.from_numpy(colored_rgb).permute(2, 0, 1)  # (3, H, W)

        return colored_tensor.float()
    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,augmentation=self.augmentation, data_type=self.data_type, dataset_type='train')
        dataset_val = LineDataset(dataset_path=self.dataset_path,augmentation=self.augmentation, data_type=self.data_type, 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'],
            weight_decay=kwargs['train_params']['optim']['weight_decay'],

        )

        model, optimizer = self.load_best_model(model, optimizer,
                                                r"\\192.168.50.222\share\rlq\weights\250725_arc_res152_best_val.pth",
                                                device)
        # scheduler = StepLR(optimizer, step_size=10, gamma=0.1)
        scheduler = ReduceLROnPlateau(optimizer, 'min', patience=30)

        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)
            scheduler.step(epoch_train_loss)

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

            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())

                print(f'val losses:{losses}')
                # print(f'val result:{result}')
            else:
                loss_dict = model(imgs, targets)
                losses = sum(loss_dict.values())
                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}')
                from utils.data_process.show_prams import print_params
                print_params(imgs[0], result[0], epoch)
                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('')