Upload 23 files

utils/__init__.py

@@ -0,0 +1 @@
+#

utils/create_exp_folder.py

@@ -0,0 +1,59 @@
+import os
+
+def create_exp_folder():
+    # Step 1: 创建run文件夹（如果不存在）
+    if not os.path.exists("run"):
+        os.mkdir("run")
+
+    # Step 2: 创建train文件夹（如果不存在）
+    train_folder = os.path.join("run", "train")
+    if not os.path.exists(train_folder):
+        os.mkdir(train_folder)
+
+    # Step 3: 创建exp文件夹（检查是否存在）
+    exp_folder = os.path.join(train_folder, "exp")
+    if not os.path.exists(exp_folder):
+        os.mkdir(exp_folder)
+        os.mkdir(os.path.join(exp_folder, "weights"))  # 创建weights文件夹
+        return exp_folder, os.path.join(exp_folder, "weights")  # 返回exp和weights文件夹路径
+
+    # 如果exp文件夹已存在，则查找exp1, exp2, 等
+    exp_num = 1
+    while True:
+        # 动态命名exp1, exp2, ...
+        exp_folder_name = f"exp{exp_num}"
+        exp_folder = os.path.join(train_folder, exp_folder_name)
+        if not os.path.exists(exp_folder):
+            os.mkdir(exp_folder)  # 创建新的exp文件夹
+            os.mkdir(os.path.join(exp_folder, "weights"))  # 创建weights文件夹
+            return exp_folder, os.path.join(exp_folder, "weights")  # 返回exp和weights文件夹路径
+        exp_num += 1  # 如果文件夹已存在，增加数字，继续查找下一个文件夹
+
+
+def create_val_exp_folder():
+    # Step 1: 创建run文件夹（如果不存在）
+    if not os.path.exists("run"):
+        os.mkdir("run")
+
+    # Step 2: 创建train文件夹（如果不存在）
+    train_folder = os.path.join("run", "predict")
+    if not os.path.exists(train_folder):
+        os.mkdir(train_folder)
+
+    # Step 3: 创建exp文件夹（检查是否存在）
+    exp_folder = os.path.join(train_folder, "exp")
+    if not os.path.exists(exp_folder):
+        os.mkdir(exp_folder)
+
+    # 如果exp文件夹已存在，则查找exp1, exp2, 等
+    exp_num = 1
+    while True:
+        # 动态命名exp1, exp2, ...
+        exp_folder_name = f"exp{exp_num}"
+        exp_folder = os.path.join(train_folder, exp_folder_name)
+        if not os.path.exists(exp_folder):
+            os.mkdir(exp_folder)  # 创建新的exp文件夹
+            return exp_folder  # 返回新创建的文件夹路径
+        exp_num += 1  # 如果文件夹已存在，增加数字，继续查找下一个文件夹
+
+

utils/dataloader.py

@@ -0,0 +1,157 @@
+import torch
+from utils.utils import cvtColor, preprocess_input
+import os
+from PIL import Image
+import numpy as np
+from torch.utils.data import Dataset, DataLoader
+import cv2
+
+class UnetDataset(Dataset):
+    def __init__(self, data_path, input_shape, num_classes, augmentation=True ,txt_name: str = "train.txt"):
+        # 读取train.txt和val.txt,test.txt文件，获取训练和验证集的图像ID
+        with open(os.path.join(data_path, "VOC2012/ImageSets/Segmentation", txt_name), "r") as f:
+            self.annotation_lines = f.readlines()
+
+        # 初始化其他参数
+        self.length = len(self.annotation_lines)  # 数据集的长度
+        self.input_shape = input_shape  # 输入图像的形状（宽和高）
+        self.num_classes = num_classes  # 类别数目
+        self.augmentation = augmentation  # 是否在训练阶段，用来控制是否使用数据增强
+        self.data_path = data_path  # 数据集路径
+
+    def __len__(self):
+        # 返回数据集的大小
+        return self.length
+
+    def __getitem__(self, index):
+        # 读取单个样本
+        annotation_line = self.annotation_lines[index]  # 获取对应的annotation
+        name = annotation_line.split()[0]  # 获取文件名，通常是图像文件的名称
+
+        # 读取JPEG图像
+        jpg = Image.open(os.path.join(self.data_path, "VOC2012/JPEGImages", name + ".png"))
+        # 读取PNG标签图像
+        png = Image.open(os.path.join(self.data_path, "VOC2012/SegmentationClass", name + ".png"))
+
+        # 如果是训练阶段，进行随机数据增强
+        jpg, png = self.get_random_data(jpg, png, self.input_shape, random=self.augmentation)
+
+        # 图像预处理
+        jpg = np.transpose(preprocess_input(np.array(jpg, np.float64)), [2, 0, 1])
+        # 标签转换成numpy数组
+        png = np.array(png)
+
+        # 将标签值大于类别数的部分设置为类别数（忽略这些区域）
+        png[png >= self.num_classes] = self.num_classes
+
+        # 将标签转换为one-hot编码
+        seg_labels = np.eye(self.num_classes + 1)[png.reshape([-1])]
+
+        # 重塑标签为目标形状
+        seg_labels = seg_labels.reshape((int(self.input_shape[0]), int(self.input_shape[1]), self.num_classes + 1))
+
+        # 返回图像、标签以及one-hot编码的标签
+        return jpg, png, seg_labels
+
+    # 生成随机数的函数
+    def rand(self, a=0, b=1):
+        return np.random.rand() * (b - a) + a
+
+    # 对图像和标签进行随机数据增强的函数
+    def get_random_data(self, image, label, input_shape, jitter=.3, hue=.1, sat=0.7, val=0.3, random=True):
+        # 将图像转为RGB格式
+        image = cvtColor(image)
+        label = Image.fromarray(np.array(label))
+
+        iw, ih = image.size  # 获取图像的宽和高
+        h, w = input_shape  # 获取目标图像的高和宽
+
+        if not random:
+            # 如果不进行随机增强（例如在验证阶段）
+            iw, ih = image.size
+            scale = min(w / iw, h / ih)  # 计算缩放比例
+            nw = int(iw * scale)  # 根据比例计算缩放后的宽
+            nh = int(ih * scale)  # 根据比例计算缩放后的高
+
+            # 缩放图像并进行中心裁剪
+            image = image.resize((nw, nh), Image.BICUBIC)
+            new_image = Image.new('RGB', [w, h], (128, 128, 128))  # 创建一个灰色背景的图像
+            new_image.paste(image, ((w - nw) // 2, (h - nh) // 2))  # 将缩放后的图像粘贴到目标图像中
+
+            # 缩放标签图像并进行中心裁剪
+            label = label.resize((nw, nh), Image.NEAREST)   # 标签使用最近邻插值
+            new_label = Image.new('L', [w, h], (0))  # 创建一个空白标签图像
+            new_label.paste(label, ((w - nw) // 2, (h - nh) // 2))  # 将缩放后的标签图像粘贴到目标图像中
+            return new_image, new_label
+
+        # 获取一个新的宽高比（通过调整宽和高的比例）
+        new_ar = iw / ih * self.rand(1 - jitter, 1 + jitter) / self.rand(1 - jitter, 1 + jitter)
+        scale = self.rand(0.25, 2)  # 随机缩放比例
+        if new_ar < 1:
+            nh = int(scale * h)
+            nw = int(nh * new_ar)
+        else:
+            nw = int(scale * w)
+            nh = int(nw / new_ar)
+        image = image.resize((nw, nh), Image.BICUBIC)
+        label = label.resize((nw, nh), Image.NEAREST)
+
+        # 随机翻转图像
+        flip = self.rand() < .5
+        if flip:
+            image = image.transpose(Image.FLIP_LEFT_RIGHT)
+            label = label.transpose(Image.FLIP_LEFT_RIGHT)
+
+        # 在图像周围随机添加灰色边框
+        dx = int(self.rand(0, w - nw))
+        dy = int(self.rand(0, h - nh))
+        new_image = Image.new('RGB', (w, h), (128, 128, 128))
+        new_label = Image.new('L', (w, h), (0))
+        new_image.paste(image, (dx, dy))   # 将图像粘贴到新图像中
+        new_label.paste(label, (dx, dy))   # 将标签粘贴到新标签中
+        image = new_image
+        label = new_label
+
+        # 转换图像为数组
+        image_data = np.array(image, np.uint8)
+
+
+        r = np.random.uniform(-1, 1, 3) * [hue, sat, val] + 1  # 随机调整色调、饱和度和亮度
+        hue, sat, val = cv2.split(cv2.cvtColor(image_data, cv2.COLOR_RGB2HSV))  # 转为HSV色域
+        dtype = image_data.dtype  # 获取数据类型
+        x = np.arange(0, 256, dtype=r.dtype)  # 获取颜色值范围
+        lut_hue = ((x * r[0]) % 180).astype(dtype)  # 应用色调变换
+        lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)  # 应用饱和度变换
+        lut_val = np.clip(x * r[2], 0, 255).astype(dtype)  # 应用亮度变换
+
+        # 使用查找表（LUT）应用变换
+        image_data = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val)))
+        image_data = cv2.cvtColor(image_data, cv2.COLOR_HSV2RGB)  # 转换回RGB色域
+
+        return image_data, label  # 返回经过增强的图像和标签
+
+
+# DataLoader中collate_fn使用
+def unet_dataset_collate(batch):
+    # 初始化三个列表，用于存储每个批次中的图像、标签和one-hot编码标签
+    images = []        # 用来存储图像数据
+    pngs = []          # 用来存储原始标签（通常是类别标签）
+    seg_labels = []    # 用来存储one-hot编码的标签
+
+    # 遍历当前批次中的每个样本（img, png, labels）
+    for img, png, labels in batch:
+        images.append(img)            # 将图像添加到images列表中
+        pngs.append(png)              # 将原始标签添加到pngs列表中
+        seg_labels.append(labels)     # 将one-hot标签添加到seg_labels列表中
+
+    # 将列表转换为NumPy数组，然后转换为torch张量
+    # images的张量需要是float类型，通常用于输入图像
+    images = torch.from_numpy(np.array(images)).type(torch.FloatTensor)
+    # pngs的张量需要是long类型，通常用于标签索引
+    pngs = torch.from_numpy(np.array(pngs)).long()
+    # seg_labels的张量需要是float类型，通常用于标签的one-hot编码
+    seg_labels = torch.from_numpy(np.array(seg_labels)).type(torch.FloatTensor)
+
+    # 返回三个张量，分别对应图像、原始标签和one-hot标签
+    return images, pngs, seg_labels
+

utils/plot_results.py

@@ -0,0 +1,46 @@
+import matplotlib.pyplot as plt
+import os
+
+def plot_training_curves(train_losses, val_losses, val_metrics_history, weights_folder):
+    # 准备数据
+    epochs = range(1, len(train_losses) + 1)
+    pixel_acc_list = [m["Pixel Accuracy"] for m in val_metrics_history]
+    mean_acc_list = [m["Mean Accuracy"] for m in val_metrics_history]
+    mean_iou_list = [m["Mean IoU"] for m in val_metrics_history]
+    fw_iou_list = [m["Frequency Weighted IoU"] for m in val_metrics_history]
+
+    # ========================
+    # 📈 绘制 Loss 曲线
+    # ========================
+    plt.figure(figsize=(8,6))
+    plt.plot(epochs, train_losses, label="Train Loss", linewidth=2)
+    plt.plot(epochs, val_losses, label="Val Loss", linewidth=2)
+
+    plt.xlabel("Epoch", fontsize=14, fontname='Times New Roman')
+    plt.ylabel("Loss", fontsize=14, fontname='Times New Roman')
+    plt.xticks(fontsize=12, fontname='Times New Roman')
+    plt.yticks(fontsize=12, fontname='Times New Roman')
+    plt.grid(True, which='both', linestyle='--', alpha=0.5)
+    plt.legend(prop={'family':'Times New Roman', 'size':12})
+    plt.tight_layout()
+    plt.savefig(os.path.join(weights_folder, "loss_curve.png"), dpi=300)
+    plt.close()
+
+    # =========================
+    # 📈 绘制指标曲线
+    # =========================
+    plt.figure(figsize=(8,6))
+    plt.plot(epochs, pixel_acc_list, label="Pixel Accuracy", linewidth=2)
+    plt.plot(epochs, mean_acc_list, label="Mean Accuracy", linewidth=2)
+    plt.plot(epochs, mean_iou_list, label="Mean IoU", linewidth=2)
+    plt.plot(epochs, fw_iou_list, label="FWIoU", linewidth=2)
+
+    plt.xlabel("Epoch", fontsize=14, fontname='Times New Roman')
+    plt.ylabel("Score", fontsize=14, fontname='Times New Roman')
+    plt.xticks(fontsize=12, fontname='Times New Roman')
+    plt.yticks(fontsize=12, fontname='Times New Roman')
+    plt.grid(True, which='both', linestyle='--', alpha=0.5)
+    plt.legend(prop={'family':'Times New Roman', 'size':12})
+    plt.tight_layout()
+    plt.savefig(os.path.join(weights_folder, "metrics_curve.png"), dpi=300)
+    plt.close()

utils/train_and_eval.py

@@ -0,0 +1,312 @@
+import os
+
+import torch
+import numpy as np
+from model.unet_training import CE_Loss, Dice_loss, Focal_Loss
+
+from utils.utils import get_lr
+from torch.cuda.amp import autocast, GradScaler
+import time
+
+
+class LogColor:
+    GREEN = "\033[1;32m"
+    YELLOW = "\033[1;33m"
+    RED = "\033[1;31m"
+    RESET = "\033[0m"
+    BLUE = "\033[1;34m"
+
+
+def pixel_accuracy(output, target):
+    with torch.no_grad():
+        _, predicted = torch.max(output, 1)
+        correct = (predicted == target).float()
+        correct_pixels = correct.sum().item()
+        total_pixels = target.numel()
+        return correct_pixels / total_pixels
+
+def mean_accuracy(output, target, num_classes):
+    """
+    计算 Mean Pixel Accuracy (MPA).
+    :param output: torch.Tensor, shape [N, C, H, W]
+    :param target: torch.Tensor, shape [N, H, W]
+    :param num_classes: int
+    :return: float, mean pixel accuracy over valid classes
+    """
+    with torch.no_grad():
+        # 取出每个像素的预测类别索引
+        _, predicted = torch.max(output, dim=1)  # shape [N, H, W]
+
+        accuracies = []
+        for i in range(num_classes):
+            # 找到该类别在标签和预测中的位置
+            target_mask = (target == i)
+            predicted_mask = (predicted == i)
+
+            # 交集：预测正确的像素数（即 TP）
+            intersection = torch.logical_and(target_mask, predicted_mask).sum().item()
+            total = target_mask.sum().item()  # 标签中该类的总像素数
+
+            if total > 0:
+                acc = intersection / total
+                accuracies.append(acc)
+            # 如果该类别在 GT 中没有出现，则跳过，不计入平均
+
+        # 防止所有类别都未出现
+        if len(accuracies) == 0:
+            return 0.0
+        else:
+            return sum(accuracies) / len(accuracies)
+
+
+# 计算Mean IoU
+def mean_iou(output, target, num_classes):
+    """
+    计算 mean IoU，只在 target 出现的类别中取平均
+    """
+    with torch.no_grad():
+        _, predicted = torch.max(output, dim=1)  # (N, H, W)
+        ious = []
+        for i in range(num_classes):
+            target_mask = (target == i)
+            pred_mask = (predicted == i)
+
+            intersection = torch.logical_and(target_mask, pred_mask).sum().item()
+            union = torch.logical_or(target_mask, pred_mask).sum().item()
+
+            if target_mask.sum().item() > 0:  # 只对 target 中存在的类求 IoU
+                ious.append(intersection / union if union > 0 else 0.0)
+        if len(ious) == 0:
+            return 0.0
+        return sum(ious) / len(ious)
+
+
+# 计算Frequency Weighted IoU
+def frequency_weighted_iou(output, target, num_classes):
+    with torch.no_grad():
+        _, predicted = torch.max(output, 1)
+        ious = []
+        frequencies = []
+        for i in range(num_classes):
+            target_mask = (target == i)
+            pred_mask = (predicted == i)
+            intersection = torch.logical_and(target_mask, pred_mask).sum().item()
+            union = torch.logical_or(target_mask, pred_mask).sum().item()
+            freq = target_mask.sum().item()
+            frequencies.append(freq)
+            ious.append((intersection / union) if union > 0 else 0.0)
+
+        total = sum(frequencies)
+        if total == 0:
+            return 0.0
+        fw_iou = sum(f * iou for f, iou in zip(frequencies, ious)) / total
+        return fw_iou
+
+
+def train_one_epoch(model, optimizer, train_loader, device, dice_loss, focal_loss,
+                    gpu_used, num_classes, scaler, epoch, train_epoch):
+
+    # 设置类别权重参数。它是用来处理类别不平衡的问题的
+    cls_weights = np.ones([num_classes], np.float32)
+    epoch_loss = 0.0 # 总的训练损失
+
+    # 设置模型为训练模式
+    model_train = model.train()
+    model_train = model_train.cuda()
+
+    # 遍历训练数据
+    for iteration, batch in enumerate(train_loader):
+        imgs, pngs, labels = batch  # 获取输入图像、标签和分割目标
+        print(f"图像像素值范围: {pngs.min().item()} ~ {pngs.max().item()}")  # 仍为0/1/2/3
+
+        # 数据准备阶段：使用 `.to(device)` 自动将数据移到设备上
+        weights = torch.tensor(cls_weights).to(device)  # 转换类别权重并移动到GPU
+        imgs = imgs.to(device)
+        pngs = pngs.to(device)
+        labels = labels.to(device)
+
+        optimizer.zero_grad()  # 清除之前的梯度
+
+        # 混合精度训练
+        if scaler is None:
+            # 前向传播
+            outputs = model_train(imgs)
+
+            # 损失计算
+            if focal_loss:
+                loss = Focal_Loss(outputs, pngs, weights, num_classes=num_classes)
+            else:
+                loss = CE_Loss(outputs, pngs, weights, num_classes=num_classes)
+
+            if dice_loss:
+                # 如果使用Dice Loss，则加上Dice损失
+                main_dice = Dice_loss(outputs, labels)
+                # main_dice = Dice_loss(outputs, pngs)
+                loss = loss + main_dice
+
+            # 反向传播
+            loss.backward()
+            optimizer.step()  # 更新模型参数
+        else:
+            with autocast():
+                outputs = model_train(imgs)  # 通过模型获取预测结果
+
+                # 损失计算
+                if focal_loss:
+                    loss = Focal_Loss(outputs, pngs, weights, num_classes=num_classes)
+                else:
+                    loss = CE_Loss(outputs, pngs, weights, num_classes=num_classes)
+
+                if dice_loss:
+                    main_dice = Dice_loss(outputs, labels)
+                    loss = loss + main_dice
+
+            # 反向传播
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)  # 使用混合精度更新梯度
+            scaler.update()  # 更新scaler
+
+        # 累加训练损失和F-score
+        epoch_loss += loss.item()
+
+        # 打印标题（每个epoch开始时打印一次）
+        if iteration == 0:  # 只在第一个 batch 打印标题
+            print(f"{LogColor.GREEN}Epoch{LogColor.RESET}{' ' * 12}"
+                  f"{LogColor.YELLOW}data_num{LogColor.RESET}{' ' * 12}"
+                  f"{LogColor.YELLOW}GPU Mem{LogColor.RESET}{' ' * 12}"
+                  f"{LogColor.YELLOW}Loss{LogColor.RESET}{' ' * 12}"
+                  f"{LogColor.YELLOW}LR{LogColor.RESET}{' ' * 12}"
+                  f"{LogColor.YELLOW}Image_size{LogColor.RESET}{' ' * 12}"
+                  )
+
+        # 每10个batch打印一次信息
+        if iteration % 1 == 0:
+            if len(train_loader) < 1:
+                a = len(train_loader)
+            else:
+                a = 1
+
+        Epoch_len = len("Epoch") + 12 - len(str(f"{epoch + 1}/{train_epoch}"))
+        batch_len = len("data_num") + 12 - len(str(f"{iteration + a}/{len(train_loader)}"))
+        GPU_len = len("GPU Mem") + 12 - len(str(f"{gpu_used:.2f} MB"))
+        Loss_len = len("Loss") + 12 - len(str(f"{loss.item():.8f}"))
+        LR_len = len("LR") + 12 - len(str(f"{get_lr(optimizer):.8f}"))
+
+        # 使用 \r 在同一行更新输出
+        print(f"\r{epoch + 1}/{train_epoch}{' ' * Epoch_len}"
+              f"{iteration + a}/{len(train_loader)}{' ' * batch_len}" 
+              f"{gpu_used:.2f} MB{' ' * GPU_len}"
+              f"{loss.item():.8f}{' ' * Loss_len}"
+              f"{get_lr(optimizer):.8f}{' ' * LR_len}"
+              f"{imgs.shape[2]}", end='', flush=True)
+
+    # 每个epoch结束后打印一次
+    print(f"{LogColor.GREEN}")
+    time.sleep(1)  # 加一点延迟，防止输出闪烁过快
+
+    # ➕ 返回平均loss
+    return epoch_loss / len(train_loader)
+
+def evaluate(model, val_loader, device, dice_loss, focal_loss, num_classes):
+
+    cls_weights = np.ones([num_classes], np.float32)
+    val_loss = 0  # 记录总的验证损失
+
+    # 设置模型为验证模式
+    model_eval = model.eval()
+    model_eval = model_eval.cuda()
+
+    # 初始化累积变量
+    total_pixel_acc = 0
+    total_mean_acc = 0
+    total_mean_iou = 0
+    total_fw_iou = 0
+    num_batches = len(val_loader)
+
+    # 遍历验证数据，前向传播
+    with torch.no_grad():
+        for iteration, batch in enumerate(val_loader):
+            imgs, pngs, labels = batch  # 获取验证数据
+
+            # 数据准备阶段
+            weights = torch.tensor(cls_weights).to(device)  # 转换类别权重并移动到GPU
+            imgs = imgs.to(device)
+            pngs = pngs.to(device)
+            labels = labels.to(device)
+            outputs = model_eval(imgs)
+
+            # print("outputs", outputs.shape)
+            # print("pngs", pngs.shape)
+            # 损失计算
+            if focal_loss:
+                loss = Focal_Loss(outputs, pngs, weights, num_classes=num_classes)
+            else:
+                loss = CE_Loss(outputs, pngs, weights, num_classes=num_classes)
+
+            if dice_loss:
+                main_dice = Dice_loss(outputs, labels)
+                # main_dice = Dice_loss(outputs, pngs)
+                loss = loss + main_dice
+
+            # 计算各个指标
+            pixel_acc = pixel_accuracy(outputs, pngs)
+            mean_acc = mean_accuracy(outputs, pngs, num_classes)
+            mean_iou_value = mean_iou(outputs, pngs, num_classes)
+            fw_iou = frequency_weighted_iou(outputs, pngs, num_classes)
+
+            # 累加到总结果
+            total_pixel_acc += pixel_acc
+            total_mean_acc += mean_acc
+            total_mean_iou += mean_iou_value
+            total_fw_iou += fw_iou
+            val_loss += loss.item()
+
+            # 打印标题（每个epoch开始时打印一次）
+            if iteration == 0:  # 只在第一个 batch 打印标题
+                epoch_len = len("Epoch") + 12
+                data_num_len = len("data_num") - len("data_num") + 12
+                Pixelacc_len = len("GPU Mem") - len("Pixelacc") + 12
+                Meanacc_len = len("Loss") - len("Meanacc") + 12
+                Meaniou_len = len("LR") - len("Meaniou") + 12
+
+                print(f"{' ' * epoch_len}"
+                      f"{LogColor.RED}data_num{LogColor.RESET}{' ' * data_num_len}"
+                      f"{LogColor.RED}Pixelacc{LogColor.RESET}{' ' * Pixelacc_len}"
+                      f"{LogColor.RED}Meanacc{LogColor.RESET}{' ' * Meanacc_len}"
+                      f"{LogColor.RED}Meaniou{LogColor.RESET}{' ' * Meaniou_len}"
+                      f"{LogColor.RED}Fwiou{LogColor.RESET}")
+
+    # 计算平均值
+    avg_pixel_acc = total_pixel_acc / num_batches
+    avg_mean_acc = total_mean_acc / num_batches
+    avg_mean_iou = total_mean_iou / num_batches
+    avg_fw_iou = total_fw_iou / num_batches
+    avg_loss = val_loss / num_batches  # ➕ 平均 loss
+
+
+    # 将结果保存到字典中
+    metrics = {
+        'Pixel Accuracy': avg_pixel_acc,
+        'Mean Accuracy': avg_mean_acc,
+        'Mean IoU': avg_mean_iou,
+        'Frequency Weighted IoU': avg_fw_iou,
+        'Loss': avg_loss  # ➕ 加入字典
+    }
+
+    epoch_len = len("Epoch") + 12
+    batch_len = data_num_len + len("data_num") - len(str(f"{iteration + 1}/{len(val_loader)}"))
+    avg_pixel_acc_len = Pixelacc_len + len("Pixelacc") - len(str(f"{avg_pixel_acc:.2f}"))
+    avg_mean_acc_len = Meanacc_len + len("Meanacc") - len(str(f"{avg_mean_acc:.2f}"))
+    avg_Mean_iou_len = Meaniou_len + len("Meaniou") - len(str(f"{avg_mean_iou:.2f}"))
+
+    # 使用 \r 在同一行更新输出
+    print(f"{' ' * (epoch_len)}"
+          f"{iteration + 1}/{len(val_loader)}{' ' * batch_len}"
+          f"{avg_pixel_acc:.2f}{' ' * avg_pixel_acc_len}"
+          f"{avg_mean_acc:.2f}{' ' * avg_mean_acc_len}"
+          f"{avg_mean_iou:.2f}{' ' * avg_Mean_iou_len}"
+          f"{avg_fw_iou:.2f}", end='', flush=True)
+    print(f"\n{LogColor.GREEN}")
+    time.sleep(1)  # 加一点延迟，防止输出闪烁过快
+
+    return metrics

utils/utils.py

@@ -0,0 +1,72 @@
+import random
+
+import numpy as np
+import torch
+from PIL import Image
+
+
+# ---------------------------------------------------------#
+#   将图像转换成RGB图像，防止灰度图在预测时报错。
+#   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
+# ---------------------------------------------------------#
+def cvtColor(image):
+    if len(np.shape(image)) == 3 and np.shape(image)[2] == 3:
+        return image
+    else:
+        image = image.convert('RGB')
+        return image
+
+    # ---------------------------------------------------#
+
+
+#   对输入图像进行resize
+# ---------------------------------------------------#
+def resize_image(image, size):
+    iw, ih = image.size
+    w, h = size
+
+    scale = min(w / iw, h / ih)
+    nw = int(iw * scale)
+    nh = int(ih * scale)
+
+    image = image.resize((nw, nh), Image.BICUBIC)
+    new_image = Image.new('RGB', size, (128, 128, 128))
+    new_image.paste(image, ((w - nw) // 2, (h - nh) // 2))
+
+    return new_image, nw, nh
+
+
+# ---------------------------------------------------#
+#   获得学习率
+# ---------------------------------------------------#
+def get_lr(optimizer):
+    for param_group in optimizer.param_groups:
+        return param_group['lr']
+
+
+# ---------------------------------------------------#
+#   设置种子
+# ---------------------------------------------------#
+def seed_everything(seed=11):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+# ---------------------------------------------------#
+#   设置Dataloader的种子
+# ---------------------------------------------------#
+def worker_init_fn(rank, seed):
+    worker_seed = rank + seed
+    random.seed(worker_seed)
+    np.random.seed(worker_seed)
+    torch.manual_seed(worker_seed)
+
+
+def preprocess_input(image):
+    image /= 255.0
+    return image