1.数据集处理

将标签处理为 x中心点偏移量，y中心点偏移量，宽度，高度 （都进行归一化）

  1. with open(opt.train_annotation_path) as f:
  2.     train_lines = f.readlines()
  3.  with open(opt.val_annotation_path) as f:
  4.         val_lines   = f.readlines()
  5. num_train   = len(train_lines)
  6. num_val     = len(val_lines)
  8. train_dataset     = YoloDataset(train_lines, opt.input_shape, train = True)
  9. train_dataloader  = DataLoader(train_dataset, shuffle = True, batch_size = opt.batch_size, num_workers = 4, pin_memory=True,
  10.                                        drop_last=True)
  1. import cv2
  2. import numpy as np
  3. import torch
  4. from PIL import Image
  5. from torch.utils.data.dataset import Dataset
  7. class YoloDataset(Dataset):
  8.     def __init__(self, annotation_lines, input_shape, train):
  9.         super(YoloDataset, self).__init__()
  10.         self.annotation_lines   = annotation_lines
  11.         self.input_shape        = input_shape
  12.         self.length             = len(self.annotation_lines)
  13.         self.train              = train
  14.         self.build_target       = build_target()
  16.     def __len__(self):
  17.         return self.length
  19.     def __getitem__(self, index):
  20.         index       = index % self.length
  21.         #---------------------------------------------------#
  22.         #   数据增强
  23.         #---------------------------------------------------#
  24.         image, boxes  = self.get_random_data(self.annotation_lines[index], self.input_shape[0:2], random = self.train )
  25.         image = torch.from_numpy(image.astype(np.float32)).permute(2, 0, 1) / 255.
  26.         boxes         = np.array(boxes, dtype=np.float32)
  27.         if len(boxes) != 0:
  28.             boxes[:, [0, 2]] = boxes[:, [0, 2]] / self.input_shape[1]
  29.             boxes[:, [1, 3]] = boxes[:, [1, 3]] / self.input_shape[0]
  30.             boxes[:, 2:4] = boxes[:, 2:4] - boxes[:, 0:2]
  31.             boxes[:, 0:2] = boxes[:, 0:2] + boxes[:, 2:4] / 2
  34.         image, targets = self.build_target(image, boxes[:,:4], boxes[:,4])
  36.         return image, targets
  38.     def get_random_data(self, annotation_line, input_shape, jitter=.3, hue=.1, sat=0.7, val=0.4, random=True):
  39.         line    = annotation_line.split()
  41.         image   = Image.open(line[0])
  42.         image = image.convert('RGB')
  44.         iw, ih  = image.size
  45.         h, w    = input_shape
  46.         #------------------------------#
  47.         #   获取边框信息
  48.         #------------------------------#
  49.         box     = np.array([np.array(list(map(int,box.split(',')))) for box in line[1:]])
  51.         if not random:
  52.             scale = min(w/iw, h/ih)
  53.             nw = int(iw*scale)
  54.             nh = int(ih*scale)
  55.             dx = (w-nw)//2
  56.             dy = (h-nh)//2
  58.             #---------------------------------#
  59.             #   添加灰度条
  60.             #---------------------------------#
  61.             image       = image.resize((nw,nh), Image.BICUBIC)
  62.             new_image   = Image.new('RGB', (w,h), (128,128,128))
  63.             new_image.paste(image, (dx, dy))
  64.             image_data  = np.array(new_image, np.float32)
  66.             #---------------------------------#
  67.             #   调整边框
  68.             #---------------------------------#
  69.             if len(box)>0:
  70.                 np.random.shuffle(box)
  71.                 box[:, [0,2]] = box[:, [0,2]]*nw/iw + dx
  72.                 box[:, [1,3]] = box[:, [1,3]]*nh/ih + dy
  73.                 box[:, 0:2][box[:, 0:2]<0] = 0
  74.                 box[:, 2][box[:, 2]>w] = w
  75.                 box[:, 3][box[:, 3]>h] = h
  76.                 box_w = box[:, 2] - box[:, 0]
  77.                 box_h = box[:, 3] - box[:, 1]
  78.                 box = box[np.logical_and(box_w>1, box_h>1)] # discard invalid box
  79.             if box.size == 0:
  80.                 print(line[0])
  81.             if len(box)==0:
  82.                 print(line[0])
  84.             return image_data, box
  86.         #------------------------------------------#
  87.         #   对图像进行缩放，并进行长和宽的扭曲
  88.         #------------------------------------------#
  89.         new_ar = iw/ih * self.rand(1-jitter,1+jitter) / self.rand(1-jitter,1+jitter)
  90.         scale = self.rand(.25, 2)
  91.         if new_ar < 1:
  92.             nh = int(scale*h)
  93.             nw = int(nh*new_ar)
  94.         else:
  95.             nw = int(scale*w)
  96.             nh = int(nw/new_ar)
  97.         image = image.resize((nw,nh), Image.BICUBIC)
  99.         #------------------------------------------#
  100.         #   添加灰度条
  101.         #------------------------------------------#
  102.         dx = int(self.rand(0, w-nw))
  103.         dy = int(self.rand(0, h-nh))
  104.         new_image = Image.new('RGB', (w,h), (128,128,128))
  105.         new_image.paste(image, (dx, dy))
  106.         image = new_image
  108.         #------------------------------------------#
  109.         #   翻转图像
  110.         #------------------------------------------#
  111.         flip = self.rand()<.5
  112.         if flip: image = image.transpose(Image.FLIP_LEFT_RIGHT)
  114.         image_data      = np.array(image, np.uint8)
  115.         #---------------------------------#
  116.         #   对图像进行色域变换
  117.         #---------------------------------#
  118.         r               = np.random.uniform(-1, 1, 3) * [hue, sat, val] + 1
  120.         hue, sat, val   = cv2.split(cv2.cvtColor(image_data, cv2.COLOR_RGB2HSV))
  121.         dtype           = image_data.dtype
  123.         x       = np.arange(0, 256, dtype=r.dtype)
  124.         lut_hue = ((x * r[0]) % 180).astype(dtype)
  125.         lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
  126.         lut_val = np.clip(x * r[2], 0, 255).astype(dtype)
  128.         image_data = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val)))
  129.         image_data = cv2.cvtColor(image_data, cv2.COLOR_HSV2RGB)
  131.         #---------------------------------#
  132.         #   调整边框
  133.         #---------------------------------#
  134.         if len(box)>0:
  135.             np.random.shuffle(box)
  136.             box[:, [0,2]] = box[:, [0,2]]*nw/iw + dx
  137.             box[:, [1,3]] = box[:, [1,3]]*nh/ih + dy
  138.             if flip: box[:, [0,2]] = w - box[:, [2,0]]
  139.             box[:, 0:2][box[:, 0:2]<0] = 0
  140.             box[:, 2][box[:, 2]>w] = w
  141.             box[:, 3][box[:, 3]>h] = h
  142.             box_w = box[:, 2] - box[:, 0]
  143.             box_h = box[:, 3] - box[:, 1]
  144.             box = box[np.logical_and(box_w>1, box_h>1)] 
  146.         return image_data, box
  147.     def rand(self, a=0, b=1):
  148.         return np.random.rand()*(b-a) + a
  150. def build_target():
  151.         return TargetTransoform()
  153. class TargetTransoform(object):
  154.     def __init__(self, target_shape=(7, 7, 30), class_nums=20, cell_nums=7):
  155.         self.target_shape = target_shape
  156.         self.class_nums = class_nums
  157.         self.cell_nums = cell_nums
  159.     def __call__(self, image, boxes, labels):
  160.         """
  161.             labels = [1,2,3,4] 
  162.             boxes = [0.2 0.3 0.4 0.8]
  163.             return [self.S,self.S,self.B*5+self.C]
  164.             """
  166.         labels = np.array(labels,"int8")
  167.         np_target = np.zeros(self.target_shape)
  168.         np_class = np.zeros((len(boxes), self.class_nums))
  169.         for i in range(len(labels)):
  170.             np_class[i][labels[i]] = 1
  171.         step = 1.0 / self.cell_nums
  172.         for i in range(len(boxes)):
  173.             box = boxes[i]
  174.             label = np_class[i]
  175.             cx, cy, w, h = box
  176.             #获取属于哪个网格
  177.             bx = int(cx // (step + 1e-5))
  178.             by = int(cy // (step + 1e-5))
  179.             #每个网格左上角的坐标
  180.             cx = (cx % step) / step
  181.             cy = (cy % step) / step
  182.             box = [cx, cy, w, h]
  183.             np_target[by][bx][:4] = box
  184.             np_target[by][bx][4] = 1
  185.             np_target[by][bx][5:9] = box
  186.             np_target[by][bx][9] = 1
  187.             np_target[by][bx][10:] = label
  188.         return image, np_target
  190. if __name__ == "__main__":
  191.     tain_path = "/home/users/user1/Documents/AI_Files/gitfies/pytorch-YOLO-v1/2007_train.txt"
  192.     with open(tain_path) as f:
  193.         train_lines = f.readlines()
  194.     train_dataset     = YoloDataset(train_lines, [448,448],  train = True)
  195.     image,Target = train_dataset[3054]

2. 损失函数

训练使得网络输出 x中心点偏移量，y中心点偏移量，宽度开根号，高度开根号 （都进行归一化）

  1. import numpy as np
  2. import torch
  3. import torch.nn.functional as F
  4. from torch.autograd import Variable
  5. from torch.nn import *
  8. class yoloLoss(Module):
  9.     def __init__(self, num_class=20):
  10.         super(yoloLoss, self).__init__()
  11.         self.lambda_coord = 5
  12.         self.lambda_noobj = 0.5
  13.         self.S = 7
  14.         self.B = 2
  15.         self.C = num_class
  16.         self.step = 1.0 / 7
  18.     def compute_iou(self, box1, box2, index):
  19.         box1 = torch.clone(box1)
  20.         box2 = torch.clone(box2)
  21.         box1 = self.conver_box(box1, index)#[x_c,y_c,w,h]
  22.         box2 = self.conver_box(box2, index)
  23.         box1[:, 2] = torch.pow(box1[:, 2],2)
  24.         box1[:, 3] = torch.pow(box1[:, 3],2)
  25.         x1, y1, w1, h1 = box1[:, 0]- box1[:, 2] / 2, box1[:, 1]- box1[:, 3] / 2, box1[:, 2], box1[:, 3]
  26.         x2, y2, w2, h2 = box2[:, 0]- box2[:, 2] / 2, box2[:, 1]- box2[:, 3] / 2, box2[:, 2], box2[:, 3]
  28.         inter_w = (w1 + w2) - (torch.max(x1 + w1, x2 + w2) - torch.min(x1, x2))
  29.         inter_h = (h1 + h2) - (torch.max(y1 + h1, y2 + h2) - torch.min(y1, y2))
  30.         inter_h = torch.clamp(inter_h, 0)
  31.         inter_w = torch.clamp(inter_w, 0)
  32.         inter = inter_w * inter_h
  33.         union = w1 * h1 + w2 * h2 - inter
  34.         return inter / union
  36.     def conver_box(self, box, index):
  37.         i, j = index
  39.         box[:, 0], box[:, 1] = [(box[:, 0] + i) * self.step, (box[:, 1] + j) * self.step]
  40.         box = torch.clamp(box, 0)
  41.         return box
  43.     def forward(self, pred, target):
  44.         batch_size = pred.size(0)
  45.         target_boxes = target[:, :, :, :10].contiguous().reshape(
  46.             (-1, 7, 7, 2, 5))
  47.         pred_boxes = pred[:, :, :, :10].contiguous().reshape((-1, 7, 7, 2, 5))
  48.         target_cls = target[:, :, :, 10:]
  49.         pred_cls = pred[:, :, :, 10:]
  50.         #h获取含有目标的坐标
  51.         obj_mask = (target_boxes[..., 4] > 0).byte()
  52.         sig_mask = obj_mask[..., 1].bool()
  53.         index = torch.where(sig_mask == True)
  54.         #img_i代表第几张图片，y属于网格y坐标，x属于网格x坐标
  55.         for img_i, y, x in zip(*index):
  56.             img_i, y, x = img_i.item(), y.item(), x.item()
  57.             pbox = pred_boxes[img_i, y, x]
  58.             target_box = target_boxes[img_i, y, x]
  59.             ious = self.compute_iou(pbox[:, :4], target_box[:, :4], [x, y])
  60.             iou, max_i = ious.max(0)
  61.             #将有目标的置信度标签且IOU较大的框框的置信度设为IOU
  62.             target_boxes[img_i, y, x, max_i, 4] = iou.item()
  63.             #另一个框框设为0
  64.             target_boxes[img_i, y, x, 1 - max_i, 4] = 0
  65.             obj_mask[img_i, y, x, 1 - max_i] = 0
  66.         obj_mask = obj_mask.bool()
  67.         noobj_mask = ~obj_mask
  69.         noobj_loss = F.mse_loss(pred_boxes[noobj_mask][:, 4],
  70.                                 target_boxes[noobj_mask][:, 4],
  71.                                 reduction="sum")
  72.         obj_loss = F.mse_loss(pred_boxes[obj_mask][:, 4],
  73.                               target_boxes[obj_mask][:, 4],
  74.                               reduction="sum")
  75.         xy_loss = F.mse_loss(pred_boxes[obj_mask][:, :2],
  76.                              target_boxes[obj_mask][:, :2],
  77.                              reduction="sum")
  78.         wh_loss = F.mse_loss(pred_boxes[obj_mask][:, 2:4],
  79.                             torch.sqrt(target_boxes[obj_mask][:, 2:4]),
  80.                              reduction="sum")
  81.         class_loss = F.mse_loss(pred_cls[sig_mask],
  82.                                 target_cls[sig_mask],
  83.                                 reduction="sum")
  85.         loss = dict(conf_loss=(obj_loss + self.lambda_noobj * noobj_loss) /
  86.                     batch_size,
  87.                     reg_loss=(self.lambda_coord * xy_loss +
  88.                               self.lambda_coord * wh_loss) / batch_size,
  89.                     cls_loss=class_loss / batch_size)
  90.         return loss