PeterYoung777/EfficientNetV2-For-Flower-Detection

1import os2import sys3import json4import pickle5import random6 7import torch8from tqdm import tqdm9 10import matplotlib.pyplot as plt11 12 13def read_split_data(root: str, val_rate: float = 0.2):14    random.seed(0)  # 保证随机结果可复现15    assert os.path.exists(root), "dataset root: {} does not exist.".format(root)16 17    # 遍历文件夹，一个文件夹对应一个类别18    flower_class = [cla for cla in os.listdir(root) if os.path.isdir(os.path.join(root, cla))]19    # 排序，保证顺序一致20    flower_class.sort()21    # 生成类别名称以及对应的数字索引22    class_indices = dict((k, v) for v, k in enumerate(flower_class))23    json_str = json.dumps(dict((val, key) for key, val in class_indices.items()), indent=4)24    with open('class_indices.json', 'w') as json_file:25        json_file.write(json_str)26 27    train_images_path = []  # 存储训练集的所有图片路径28    train_images_label = []  # 存储训练集图片对应索引信息29    val_images_path = []  # 存储验证集的所有图片路径30    val_images_label = []  # 存储验证集图片对应索引信息31    every_class_num = []  # 存储每个类别的样本总数32    supported = [".jpg", ".JPG", ".png", ".PNG"]  # 支持的文件后缀类型33    # 遍历每个文件夹下的文件34    for cla in flower_class:35        cla_path = os.path.join(root, cla)36        # 遍历获取supported支持的所有文件路径37        images = [os.path.join(root, cla, i) for i in os.listdir(cla_path)38                  if os.path.splitext(i)[-1] in supported]39        # 获取该类别对应的索引40        image_class = class_indices[cla]41        # 记录该类别的样本数量42        every_class_num.append(len(images))43        # 按比例随机采样验证样本44        val_path = random.sample(images, k=int(len(images) * val_rate))45 46        for img_path in images:47            if img_path in val_path:  # 如果该路径在采样的验证集样本中则存入验证集48                val_images_path.append(img_path)49                val_images_label.append(image_class)50            else:  # 否则存入训练集51                train_images_path.append(img_path)52                train_images_label.append(image_class)53 54    print("{} images were found in the dataset.".format(sum(every_class_num)))55    print("{} images for training.".format(len(train_images_path)))56    print("{} images for validation.".format(len(val_images_path)))57 58    plot_image = False59    if plot_image:60        # 绘制每种类别个数柱状图61        plt.bar(range(len(flower_class)), every_class_num, align='center')62        # 将横坐标0,1,2,3,4替换为相应的类别名称63        plt.xticks(range(len(flower_class)), flower_class)64        # 在柱状图上添加数值标签65        for i, v in enumerate(every_class_num):66            plt.text(x=i, y=v + 5, s=str(v), ha='center')67        # 设置x坐标68        plt.xlabel('image class')69        # 设置y坐标70        plt.ylabel('number of images')71        # 设置柱状图的标题72        plt.title('flower class distribution')73        plt.show()74 75    return train_images_path, train_images_label, val_images_path, val_images_label76 77 78def plot_data_loader_image(data_loader):79    batch_size = data_loader.batch_size80    plot_num = min(batch_size, 4)81 82    json_path = './class_indices.json'83    assert os.path.exists(json_path), json_path + " does not exist."84    json_file = open(json_path, 'r')85    class_indices = json.load(json_file)86 87    for data in data_loader:88        images, labels = data89        for i in range(plot_num):90            # [C, H, W] -> [H, W, C]91            img = images[i].numpy().transpose(1, 2, 0)92            # 反Normalize操作93            img = (img * [0.229, 0.224, 0.225] + [0.485, 0.456, 0.406]) * 25594            label = labels[i].item()95            plt.subplot(1, plot_num, i+1)96            plt.xlabel(class_indices[str(label)])97            plt.xticks([])  # 去掉x轴的刻度98            plt.yticks([])  # 去掉y轴的刻度99            plt.imshow(img.astype('uint8'))100        plt.show()101 102 103def write_pickle(list_info: list, file_name: str):104    with open(file_name, 'wb') as f:105        pickle.dump(list_info, f)106 107 108def read_pickle(file_name: str) -> list:109    with open(file_name, 'rb') as f:110        info_list = pickle.load(f)111        return info_list112 113 114def train_one_epoch(model, optimizer, data_loader, device, epoch):115    model.train()116    loss_function = torch.nn.CrossEntropyLoss()117    accu_loss = torch.zeros(1).to(device)  # 累计损失118    accu_num = torch.zeros(1).to(device)   # 累计预测正确的样本数119    optimizer.zero_grad()120 121    sample_num = 0122    data_loader = tqdm(data_loader)123    for step, data in enumerate(data_loader):124        images, labels = data125        sample_num += images.shape[0]126 127        pred = model(images.to(device))128        pred_classes = torch.max(pred, dim=1)[1]129        accu_num += torch.eq(pred_classes, labels.to(device)).sum()130 131        loss = loss_function(pred, labels.to(device))132        loss.backward()133        accu_loss += loss.detach()134 135        data_loader.desc = "[train epoch {}] loss: {:.3f}, acc: {:.3f}".format(epoch,136                                                                               accu_loss.item() / (step + 1),137                                                                               accu_num.item() / sample_num)138 139        if not torch.isfinite(loss):140            print('WARNING: non-finite loss, ending training ', loss)141            sys.exit(1)142 143        optimizer.step()144        optimizer.zero_grad()145 146    return accu_loss.item() / (step + 1), accu_num.item() / sample_num147 148 149@torch.no_grad()150def evaluate(model, data_loader, device, epoch):151    loss_function = torch.nn.CrossEntropyLoss()152 153    model.eval()154 155    accu_num = torch.zeros(1).to(device)   # 累计预测正确的样本数156    accu_loss = torch.zeros(1).to(device)  # 累计损失157 158    sample_num = 0159    data_loader = tqdm(data_loader)160    for step, data in enumerate(data_loader):161        images, labels = data162        sample_num += images.shape[0]163 164        pred = model(images.to(device))165        pred_classes = torch.max(pred, dim=1)[1]166        accu_num += torch.eq(pred_classes, labels.to(device)).sum()167 168        loss = loss_function(pred, labels.to(device))169        accu_loss += loss170 171        data_loader.desc = "[valid epoch {}] loss: {:.3f}, acc: {:.3f}".format(epoch,172                                                                               accu_loss.item() / (step + 1),173                                                                               accu_num.item() / sample_num)174 175    return accu_loss.item() / (step + 1), accu_num.item() / sample_num176