首页 > 分享 > 【深度学习】快速制作图像标签数据集以及训练

【深度学习】快速制作图像标签数据集以及训练

花匠小妙招
2024-12-17 22:34

快速制作图像标签数据集以及训练

制作DataSet

先从网络收集十张图片每种十张
在这里插入图片描述

定义dataSet和dataloader

import glob import torch from torch.utils import data from PIL import Image import numpy as np from torchvision import transforms import matplotlib.pyplot as plt # 通过创建data.Dataset子类Mydataset来创建输入 class Mydataset(data.Dataset): # init() 初始化方法，传入数据文件夹路径 def __init__(self, root): self.imgs_path = root # getitem() 切片方法，根据索引下标，获得相应的图片 def __getitem__(self, index): img_path = self.imgs_path[index] # len() 计算长度方法，返回整个数据文件夹下所有文件的个数 def __len__(self): return len(self.imgs_path) # 使用glob方法来获取数据图片的所有路径 all_imgs_path = glob.glob(r"./Data/*/*.jpg") # 数据文件夹路径 # 利用自定义类Mydataset创建对象brake_dataset # 将所有的路径塞进dataset 使用每张图片的路径进行索引图片 brake_dataset = Mydataset(all_imgs_path) # print("图片总数:{}".format(len(brake_dataset))) # 返回文件夹中图片总个数 # 制作dataloader brake_dataloader = torch.utils.data.DataLoader(brake_dataset, batch_size=2) # 每次迭代时返回4个数据 # print(next(iter(break_dataloader))) 1234567891011121314151617181920212223242526272829303132333435363738

制作标签

# 为每张图片制作对应标签 species = ['sun', 'rain', 'cloud'] species_to_id = dict((c, i) for i, c in enumerate(species)) # print(species_to_id) id_to_species = dict((v, k) for k, v in species_to_id.items()) # print(id_to_species) # 对所有图片路径进行迭代 all_labels = [] for img in all_imgs_path:# 区分出每个img，应该属于什么类别for i, c in enumerate(species):if c in img:all_labels.append(i) # print(all_labels) 1234567891011121314151617

制作数据和标签一起的dataset和dataloader

上面的dataset不够完善

# 将数据转换为张量数据 # 对数据进行转换处理 transform = transforms.Compose([transforms.Resize((256, 256)), # 做的第一步转换transforms.ToTensor() # 第二步转换，作用：第一转换成Tensor，第二将图片取值范围转换成0-1之间，第三会将channel置前 ]) class Mydatasetpro(data.Dataset):def __init__(self, img_paths, labels, transform):self.imgs = img_pathsself.labels = labelsself.transforms = transform# 进行切片def __getitem__(self, index):img = self.imgs[index]label = self.labels[index]pil_img = Image.open(img) # pip install pillowpil_img = pil_img.convert('RGB')data = self.transforms(pil_img)return data, label# 返回长度def __len__(self):return len(self.imgs) BATCH_SIZE = 4 brake_dataset = Mydatasetpro(all_imgs_path, all_labels, transform) brake_dataloader = data.DataLoader(brake_dataset,batch_size=BATCH_SIZE,shuffle=True ) 123456789101112131415161718192021222324252627282930313233343536

imgs_batch, labels_batch = next(iter(brake_dataloader)) # 4 X 3 X 256 X 256 print(imgs_batch.shape) plt.figure(figsize=(12, 8)) for i, (img, label) in enumerate(zip(imgs_batch[:10], labels_batch[:10])):img = img.permute(1, 2, 0).numpy()plt.subplot(2, 3, i + 1)plt.title(id_to_species.get(label.item()))plt.imshow(img) plt.show() # 展示图片 1234567891011121314

制作训练集和测试集

# 划分数据集和测试集 index = np.random.permutation(len(all_imgs_path)) # 打乱所有图片的索引 print(index) # 根据索引获取所有图片的路径 all_imgs_path = np.array(all_imgs_path)[index] all_labels = np.array(all_labels)[index] print("打乱顺序之后的所有图片路径{}".format(all_imgs_path)) print("打乱顺序之后的所有图片索引{}".format(all_labels)) # 80%做训练集 s = int(len(all_imgs_path) * 0.8) # print(s) train_imgs = all_imgs_path[:s] # print(train_imgs) train_labels = all_labels[:s] test_imgs = all_imgs_path[s:] test_labels = all_labels[s:] # 将训练集和标签制作dataset 需要转换为张量 train_ds = Mydatasetpro(train_imgs, train_labels, transform) # TrainSet TensorData test_ds = Mydatasetpro(test_imgs, test_labels, transform) # TestSet TensorData # print(train_ds) # print(test_ds) print("**********") # 制作trainLoader train_dl = data.DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True) # TrainSet Labels test_dl = data.DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True) # TestSet Labels 12345678910111213141516171819202122232425262728293031323334353637

训练代码

import torch import torchvision.models as models from torch import nn from torch import optim from DataSetMake import brake_dataloader from DataSetMake import train_dl, test_dl # 判断是否使用GPU DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 使用resnet 训练 model_ft = models.resnet50(pretrained=True) # 使用迁移学习，加载预训练权 in_features = model_ft.fc.in_features model_ft.fc = nn.Sequential(nn.Linear(in_features, 256),nn.ReLU(),# nn.Dropout(0, 4),nn.Linear(256, 4),nn.LogSoftmax(dim=1)) model_ft = model_ft.to(DEVICE) # 将模型迁移到gpu # 优化器 loss_fn = nn.CrossEntropyLoss() loss_fn = loss_fn.to(DEVICE) # 将loss_fn迁移到GPU # Adam损失函数 optimizer = optim.Adam(model_ft.fc.parameters(), lr=0.003) epochs = 50 # 迭代次数 steps = 0 running_loss = 0 print_every = 10 train_losses, test_losses = [], [] for epoch in range(epochs):model_ft.train()# 遍历训练集数据for imgs, labels in brake_dataloader:steps += 1# 标签转换为 tensorlabels = torch.tensor(labels, dtype=torch.long)# 将图片和标签放到设备上imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)optimizer.zero_grad() # 梯度归零# 前向推理outputs = model_ft(imgs)# 计算lossloss = loss_fn(outputs, labels)loss.backward() # 反向传播计算梯度optimizer.step() # 梯度优化# 累加lossrunning_loss += loss.item()if steps % print_every == 0:test_loss = 0accuracy = 0# 验证模式model_ft.eval()# 测试集不需要计算梯度with torch.no_grad():# 遍历测试集数据for imgs, labels in test_dl:# 转换为tensorlabels = torch.tensor(labels, dtype=torch.long)# 数据标签部署到gpuimgs, labels = imgs.to(DEVICE), labels.to(DEVICE)# 前向推理outputs = model_ft(imgs)# 计算损失loss = loss_fn(outputs, labels)# 累加测试机的损失test_loss += loss.item()ps = torch.exp(outputs)top_p, top_class = ps.topk(1, dim=1)equals = top_class == labels.view(*top_class.shape)accuracy += torch.mean(equals.type(torch.FloatTensor)).item()train_losses.append(running_loss / len(train_dl))test_losses.append(test_loss / len(test_dl))print(f"Epoch {epoch + 1}/{epochs}.. " f"Train loss: {running_loss / print_every:.3f}.. " f"Test loss: {test_loss / len(test_dl):.3f}.. " f"Test accuracy: {accuracy / len(test_dl):.3f}")# 回到训练模式训练误差清0running_loss = 0model_ft.train() torch.save(model_ft, "aerialmodel.pth") 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111

在这里插入图片描述

预测代码

import os import torch from PIL import Image from torch import nn from torchvision import transforms, models i = 0 # 识别图片计数 # 这里最好新建一个test_data文件随机放一些上面整理好的图片进去 root_path = r"D:CODEImageClassifyTest" # 待测试文件夹 names = os.listdir(root_path) for name in names:print(name)i = i + 1data_class = ['sun', 'rain', 'cloud'] # 按文件索引顺序排列# 找出文件夹中的所有图片image_path = os.path.join(root_path, name)image = Image.open(image_path)print(image)# 张量定义格式transform = transforms.Compose([transforms.Resize((256, 256)),transforms.ToTensor()])# 图片转换为张量image = transform(image)print(image.shape)# 定义resnet模型model_ft = models.resnet50()# 模型结构in_features = model_ft.fc.in_featuresmodel_ft.fc = nn.Sequential(nn.Linear(in_features, 256),nn.ReLU(),# nn.Dropout(0, 4),nn.Linear(256, 4),nn.LogSoftmax(dim=1))# 加载已经训练好的模型参数model = torch.load("aerialmodel.pth", map_location=torch.device("cpu"))# 将每张图片调整维度image = torch.reshape(image, (1, 3, 256, 256)) # 修改待预测图片尺寸，需要与训练时一致model.eval()# 速出预测结果with torch.no_grad():output = model(image)print(output) # 输出预测结果# print(int(output.argmax(1)))# 对结果进行处理，使直接显示出预测的种类根据索引判别是哪一类print("第{}张图片预测为：{}".format(i, data_class[int(output.argmax(1))])) 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657