Add files via upload

VOCdevkit/VOC2007/Annotations/README.md

+存放标签文件
\ No newline at end of file

VOCdevkit/VOC2007/ImageSets/Main/README.md

+存放训练索引文件
\ No newline at end of file

VOCdevkit/VOC2007/JPEGImages/README.md

+存放图片文件
\ No newline at end of file

VOCdevkit/VOC2007/voc2yolo4.py

+import os
+import random 
+ 
+xmlfilepath=r'./VOCdevkit/VOC2007/Annotations'
+saveBasePath=r"./VOCdevkit/VOC2007/ImageSets/Main/"
+ 
+trainval_percent=0
+train_percent=1
+
+temp_xml = os.listdir(xmlfilepath)
+total_xml = []
+for xml in temp_xml:
+    if xml.endswith(".xml"):
+        total_xml.append(xml)
+
+num=len(total_xml)  
+list=range(num)  
+tv=int(num*trainval_percent)  
+tr=int(tv*train_percent)  
+trainval= random.sample(list,tv)  
+train=random.sample(trainval,tr)  
+ 
+print("train and val size",tv)
+print("traub suze",tr)
+ftrainval = open(os.path.join(saveBasePath,'trainval.txt'), 'w')  
+ftest = open(os.path.join(saveBasePath,'test.txt'), 'w')  
+ftrain = open(os.path.join(saveBasePath,'train.txt'), 'w')  
+fval = open(os.path.join(saveBasePath,'val.txt'), 'w')  
+ 
+for i  in list:  
+    name=total_xml[i][:-4]+'\n'  
+    if i in trainval:  
+        ftrainval.write(name)  
+        if i in train:  
+            ftrain.write(name)  
+        else:  
+            fval.write(name)  
+    else:  
+        ftest.write(name)  
+  
+ftrainval.close()  
+ftrain.close()  
+fval.close()  
+ftest .close()

ciou_test.py

+import torch
+import math
+import numpy as np
+def box_ciou(b1, b2):
+    """
+    输入为：
+    ----------
+    b1: tensor, shape=(batch, feat_w, feat_h, anchor_num, 4), xywh
+    b2: tensor, shape=(batch, feat_w, feat_h, anchor_num, 4), xywh
+
+    返回为：
+    -------
+    ciou: tensor, shape=(batch, feat_w, feat_h, anchor_num, 1)
+    """
+    # 求出预测框左上角右下角
+    b1_xy = b1[..., :2]
+    b1_wh = b1[..., 2:4]
+    b1_wh_half = b1_wh/2.
+    b1_mins = b1_xy - b1_wh_half
+    b1_maxes = b1_xy + b1_wh_half
+    # 求出真实框左上角右下角
+    b2_xy = b2[..., :2]
+    b2_wh = b2[..., 2:4]
+    b2_wh_half = b2_wh/2.
+    b2_mins = b2_xy - b2_wh_half
+    b2_maxes = b2_xy + b2_wh_half
+
+    # 求真实框和预测框所有的iou
+    intersect_mins = torch.max(b1_mins, b2_mins)
+    intersect_maxes = torch.min(b1_maxes, b2_maxes)
+    intersect_wh = torch.max(intersect_maxes - intersect_mins, torch.zeros_like(intersect_maxes))
+    intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]
+    b1_area = b1_wh[..., 0] * b1_wh[..., 1]
+    b2_area = b2_wh[..., 0] * b2_wh[..., 1]
+    union_area = b1_area + b2_area - intersect_area
+    iou = intersect_area / (union_area + 1e-7)
+
+    # 计算中心的差距
+    center_distance = torch.sum(torch.pow((b1_xy - b2_xy), 2), axis=-1)
+    # 找到包裹两个框的最小框的左上角和右下角
+    enclose_mins = torch.min(b1_mins, b2_mins)
+    enclose_maxes = torch.max(b1_maxes, b2_maxes)
+    enclose_wh = torch.max(enclose_maxes - enclose_mins, torch.zeros_like(intersect_maxes))
+    # 计算对角线距离
+    enclose_diagonal = torch.sum(torch.pow(enclose_wh,2), axis=-1)
+    ciou = iou - 1.0 * (center_distance) / (enclose_diagonal + 1e-7)
+
+    v = (4 / (math.pi ** 2)) * torch.pow((torch.atan(b1_wh[..., 0]/b1_wh[..., 1]) - torch.atan(b2_wh[..., 0]/b2_wh[..., 1])), 2)
+    alpha = v / (1.0 - iou + v)
+    ciou = ciou - alpha * v
+    return ciou
+
+box1 = torch.from_numpy(np.array([[25,25,40,40]])).type(torch.FloatTensor)
+box2 = torch.from_numpy(np.array([[25,25,30,40]])).type(torch.FloatTensor)
+
+print(box_ciou(box1,box2))
\ No newline at end of file

get_dr_txt.py

+#-------------------------------------#
+#       mAP所需文件计算代码
+#       具体教程请查看Bilibili
+#       Bubbliiiing
+#-------------------------------------#
+import cv2
+import keras
+import numpy as np
+import colorsys
+import os
+import torch
+import torch.nn as nn
+import torch.backends.cudnn as cudnn
+from torch.autograd import Variable
+from yolo import YOLO
+from nets.yolo4 import YoloBody
+from PIL import Image,ImageFont, ImageDraw
+from utils.utils import non_max_suppression, bbox_iou, DecodeBox,letterbox_image,yolo_correct_boxes
+
+class mAP_Yolo(YOLO):
+    #---------------------------------------------------#
+    #   检测图片
+    #---------------------------------------------------#
+    def detect_image(self,image_id,image):
+        self.confidence = 0.05
+        f = open("./input/detection-results/"+image_id+".txt","w") 
+        image_shape = np.array(np.shape(image)[0:2])
+
+        crop_img = np.array(letterbox_image(image, (self.model_image_size[0],self.model_image_size[1])))
+        photo = np.array(crop_img,dtype = np.float32)
+        photo /= 255.0
+        photo = np.transpose(photo, (2, 0, 1))
+        photo = photo.astype(np.float32)
+        images = []
+        images.append(photo)
+        images = np.asarray(images)
+
+        with torch.no_grad():
+            images = torch.from_numpy(images)
+            if self.cuda:
+                images = images.cuda()
+                
+            
+        outputs = self.net(images)
+        output_list = []
+        for i in range(3):
+            output_list.append(self.yolo_decodes[i](outputs[i]))
+        output = torch.cat(output_list, 1)
+        batch_detections = non_max_suppression(output, len(self.class_names),
+                                                conf_thres=self.confidence,
+                                                nms_thres=0.3)
+
+        try:
+            batch_detections = batch_detections[0].cpu().numpy()
+        except:
+            return image
+            
+        top_index = batch_detections[:,4]*batch_detections[:,5] > self.confidence
+        top_conf = batch_detections[top_index,4]*batch_detections[top_index,5]
+        top_label = np.array(batch_detections[top_index,-1],np.int32)
+        top_bboxes = np.array(batch_detections[top_index,:4])
+        top_xmin, top_ymin, top_xmax, top_ymax = np.expand_dims(top_bboxes[:,0],-1),np.expand_dims(top_bboxes[:,1],-1),np.expand_dims(top_bboxes[:,2],-1),np.expand_dims(top_bboxes[:,3],-1)
+
+        # 去掉灰条
+        boxes = yolo_correct_boxes(top_ymin,top_xmin,top_ymax,top_xmax,np.array([self.model_image_size[0],self.model_image_size[1]]),image_shape)
+
+        for i, c in enumerate(top_label):
+            predicted_class = self.class_names[c]
+            score = str(top_conf[i])
+
+            top, left, bottom, right = boxes[i]
+            f.write("%s %s %s %s %s %s\n" % (predicted_class, score[:6], str(int(left)), str(int(top)), str(int(right)),str(int(bottom))))
+
+        f.close()
+        return 
+
+yolo = mAP_Yolo()
+image_ids = open('VOCdevkit/VOC2007/ImageSets/Main/test.txt').read().strip().split()
+
+if not os.path.exists("./input"):
+    os.makedirs("./input")
+if not os.path.exists("./input/detection-results"):
+    os.makedirs("./input/detection-results")
+if not os.path.exists("./input/images-optional"):
+    os.makedirs("./input/images-optional")
+
+
+for image_id in image_ids:
+    image_path = "./VOCdevkit/VOC2007/JPEGImages/"+image_id+".jpg"
+    image = Image.open(image_path)
+    # 开启后在之后计算mAP可以可视化
+    # image.save("./input/images-optional/"+image_id+".jpg")
+    yolo.detect_image(image_id,image)
+    print(image_id," done!")
+    
+
+print("Conversion completed!")
\ No newline at end of file

get_gt_txt.py

+#-------------------------------------#
+#       mAP所需文件计算代码
+#       具体教程请查看Bilibili
+#       Bubbliiiing
+#-------------------------------------#
+import sys
+import os
+import glob
+import xml.etree.ElementTree as ET
+
+image_ids = open('VOCdevkit/VOC2007/ImageSets/Main/test.txt').read().strip().split()
+
+if not os.path.exists("./input"):
+    os.makedirs("./input")
+if not os.path.exists("./input/ground-truth"):
+    os.makedirs("./input/ground-truth")
+
+for image_id in image_ids:
+    with open("./input/ground-truth/"+image_id+".txt", "w") as new_f:
+        root = ET.parse("VOCdevkit/VOC2007/Annotations/"+image_id+".xml").getroot()
+        for obj in root.findall('object'):
+            if obj.find('difficult')!=None:
+                difficult = obj.find('difficult').text
+                if int(difficult)==1:
+                    continue
+            obj_name = obj.find('name').text
+            bndbox = obj.find('bndbox')
+            left = bndbox.find('xmin').text
+            top = bndbox.find('ymin').text
+            right = bndbox.find('xmax').text
+            bottom = bndbox.find('ymax').text
+            new_f.write("%s %s %s %s %s\n" % (obj_name, left, top, right, bottom))
+print("Conversion completed!")
\ No newline at end of file

logs/README.md

+用于存放训练好的文件
\ No newline at end of file

model_data/coco_classes.txt

+person
+bicycle
+car
+motorbike
+aeroplane
+bus
+train
+truck
+boat
+traffic light
+fire hydrant
+stop sign
+parking meter
+bench
+bird
+cat
+dog
+horse
+sheep
+cow
+elephant
+bear
+zebra
+giraffe
+backpack
+umbrella
+handbag
+tie
+suitcase
+frisbee
+skis
+snowboard
+sports ball
+kite
+baseball bat
+baseball glove
+skateboard
+surfboard
+tennis racket
+bottle
+wine glass
+cup
+fork
+knife
+spoon
+bowl
+banana
+apple
+sandwich
+orange
+broccoli
+carrot
+hot dog
+pizza
+donut
+cake
+chair
+sofa
+pottedplant
+bed
+diningtable
+toilet
+tvmonitor
+laptop
+mouse
+remote
+keyboard
+cell phone
+microwave
+oven
+toaster
+sink
+refrigerator
+book
+clock
+vase
+scissors
+teddy bear
+hair drier
+toothbrush

model_data/voc_classes.txt

+aeroplane
+bicycle
+bird
+boat
+bottle
+bus
+car
+cat
+chair
+cow
+diningtable
+dog
+horse
+motorbike
+person
+pottedplant
+sheep
+sofa
+train
+tvmonitor
\ No newline at end of file

model_data/yolo_anchors.txt

+12, 16,  19, 36,  40, 28,  36, 75,  76, 55,  72, 146,  142, 110,  192, 243,  459, 401
\ No newline at end of file

nets/CSPdarknet.py

+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+import math
+from collections import OrderedDict
+
+#-------------------------------------------------#
+#   MISH激活函数
+#-------------------------------------------------#
+class Mish(nn.Module):
+    def __init__(self):
+        super(Mish, self).__init__()
+
+    def forward(self, x):
+        return x * torch.tanh(F.softplus(x))
+
+#-------------------------------------------------#
+#   卷积块
+#   CONV+BATCHNORM+MISH
+#-------------------------------------------------#
+class BasicConv(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1):
+        super(BasicConv, self).__init__()
+
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, kernel_size//2, bias=False)
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.activation = Mish()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.activation(x)
+        return x
+
+#---------------------------------------------------#
+#   CSPdarknet的结构块的组成部分
+#   内部堆叠的残差块
+#---------------------------------------------------#
+class Resblock(nn.Module):
+    def __init__(self, channels, hidden_channels=None, residual_activation=nn.Identity()):
+        super(Resblock, self).__init__()
+
+        if hidden_channels is None:
+            hidden_channels = channels
+
+        self.block = nn.Sequential(
+            BasicConv(channels, hidden_channels, 1),
+            BasicConv(hidden_channels, channels, 3)
+        )
+
+    def forward(self, x):
+        return x + self.block(x)
+
+#---------------------------------------------------#
+#   CSPdarknet的结构块
+#   存在一个大残差边
+#   这个大残差边绕过了很多的残差结构
+#---------------------------------------------------#
+class Resblock_body(nn.Module):
+    def __init__(self, in_channels, out_channels, num_blocks, first):
+        super(Resblock_body, self).__init__()
+
+        self.downsample_conv = BasicConv(in_channels, out_channels, 3, stride=2)
+
+        if first:
+            self.split_conv0 = BasicConv(out_channels, out_channels, 1)
+            self.split_conv1 = BasicConv(out_channels, out_channels, 1)  
+            self.blocks_conv = nn.Sequential(
+                Resblock(channels=out_channels, hidden_channels=out_channels//2),
+                BasicConv(out_channels, out_channels, 1)
+            )
+            self.concat_conv = BasicConv(out_channels*2, out_channels, 1)
+        else:
+            self.split_conv0 = BasicConv(out_channels, out_channels//2, 1)
+            self.split_conv1 = BasicConv(out_channels, out_channels//2, 1)
+
+            self.blocks_conv = nn.Sequential(
+                *[Resblock(out_channels//2) for _ in range(num_blocks)],
+                BasicConv(out_channels//2, out_channels//2, 1)
+            )
+            self.concat_conv = BasicConv(out_channels, out_channels, 1)
+
+    def forward(self, x):
+        x = self.downsample_conv(x)
+
+        x0 = self.split_conv0(x)
+
+        x1 = self.split_conv1(x)
+        x1 = self.blocks_conv(x1)
+
+        x = torch.cat([x1, x0], dim=1)
+        x = self.concat_conv(x)
+
+        return x
+
+class CSPDarkNet(nn.Module):
+    def __init__(self, layers):
+        super(CSPDarkNet, self).__init__()
+        self.inplanes = 32
+        self.conv1 = BasicConv(3, self.inplanes, kernel_size=3, stride=1)
+        self.feature_channels = [64, 128, 256, 512, 1024]
+
+        self.stages = nn.ModuleList([
+            Resblock_body(self.inplanes, self.feature_channels[0], layers[0], first=True),
+            Resblock_body(self.feature_channels[0], self.feature_channels[1], layers[1], first=False),
+            Resblock_body(self.feature_channels[1], self.feature_channels[2], layers[2], first=False),
+            Resblock_body(self.feature_channels[2], self.feature_channels[3], layers[3], first=False),
+            Resblock_body(self.feature_channels[3], self.feature_channels[4], layers[4], first=False)
+        ])
+
+        self.num_features = 1
+        # 进行权值初始化
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+
+    def forward(self, x):
+        x = self.conv1(x)
+
+        x = self.stages[0](x)
+        x = self.stages[1](x)
+        out3 = self.stages[2](x)
+        out4 = self.stages[3](out3)
+        out5 = self.stages[4](out4)
+
+        return out3, out4, out5
+
+def darknet53(pretrained, **kwargs):
+    model = CSPDarkNet([1, 2, 8, 8, 4])
+    if pretrained:
+        if isinstance(pretrained, str):
+            model.load_state_dict(torch.load(pretrained))
+        else:
+            raise Exception("darknet request a pretrained path. got [{}]".format(pretrained))
+    return model

nets/yolo4.py

+import torch
+import torch.nn as nn
+from collections import OrderedDict
+from nets.CSPdarknet import darknet53
+
+def conv2d(filter_in, filter_out, kernel_size, stride=1):
+    pad = (kernel_size - 1) // 2 if kernel_size else 0
+    return nn.Sequential(OrderedDict([
+        ("conv", nn.Conv2d(filter_in, filter_out, kernel_size=kernel_size, stride=stride, padding=pad, bias=False)),
+        ("bn", nn.BatchNorm2d(filter_out)),
+        ("relu", nn.LeakyReLU(0.1)),
+    ]))
+
+#---------------------------------------------------#
+#   SPP结构，利用不同大小的池化核进行池化
+#   池化后堆叠
+#---------------------------------------------------#
+class SpatialPyramidPooling(nn.Module):
+    def __init__(self, pool_sizes=[5, 9, 13]):
+        super(SpatialPyramidPooling, self).__init__()
+
+        self.maxpools = nn.ModuleList([nn.MaxPool2d(pool_size, 1, pool_size//2) for pool_size in pool_sizes])
+
+    def forward(self, x):
+        features = [maxpool(x) for maxpool in self.maxpools[::-1]]
+        features = torch.cat(features + [x], dim=1)
+
+        return features
+
+#---------------------------------------------------#
+#   卷积 + 上采样
+#---------------------------------------------------#
+class Upsample(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(Upsample, self).__init__()
+
+        self.upsample = nn.Sequential(
+            conv2d(in_channels, out_channels, 1),
+            nn.Upsample(scale_factor=2, mode='nearest')
+        )
+
+    def forward(self, x,):
+        x = self.upsample(x)
+        return x
+
+#---------------------------------------------------#
+#   三次卷积块
+#---------------------------------------------------#
+def make_three_conv(filters_list, in_filters):
+    m = nn.Sequential(
+        conv2d(in_filters, filters_list[0], 1),
+        conv2d(filters_list[0], filters_list[1], 3),
+        conv2d(filters_list[1], filters_list[0], 1),
+    )
+    return m
+
+#---------------------------------------------------#
+#   五次卷积块
+#---------------------------------------------------#
+def make_five_conv(filters_list, in_filters):
+    m = nn.Sequential(
+        conv2d(in_filters, filters_list[0], 1),
+        conv2d(filters_list[0], filters_list[1], 3),
+        conv2d(filters_list[1], filters_list[0], 1),
+        conv2d(filters_list[0], filters_list[1], 3),
+        conv2d(filters_list[1], filters_list[0], 1),
+    )
+    return m
+
+#---------------------------------------------------#
+#   最后获得yolov4的输出
+#---------------------------------------------------#
+def yolo_head(filters_list, in_filters):
+    m = nn.Sequential(
+        conv2d(in_filters, filters_list[0], 3),
+        nn.Conv2d(filters_list[0], filters_list[1], 1),
+    )
+    return m
+
+#---------------------------------------------------#
+#   yolo_body
+#---------------------------------------------------#
+class YoloBody(nn.Module):
+    def __init__(self, num_anchors, num_classes):
+        super(YoloBody, self).__init__()
+        #  backbone
+        self.backbone = darknet53(None)
+
+        self.conv1 = make_three_conv([512,1024],1024)
+        self.SPP = SpatialPyramidPooling()
+        self.conv2 = make_three_conv([512,1024],2048)
+
+        self.upsample1 = Upsample(512,256)
+        self.conv_for_P4 = conv2d(512,256,1)
+        self.make_five_conv1 = make_five_conv([256, 512],512)
+
+        self.upsample2 = Upsample(256,128)
+        self.conv_for_P3 = conv2d(256,128,1)
+        self.make_five_conv2 = make_five_conv([128, 256],256)
+        # 3*(5+num_classes)=3*(5+20)=3*(4+1+20)=75
+        # 4+1+num_classes
+        final_out_filter2 = num_anchors * (5 + num_classes)
+        self.yolo_head3 = yolo_head([256, final_out_filter2],128)
+
+        self.down_sample1 = conv2d(128,256,3,stride=2)
+        self.make_five_conv3 = make_five_conv([256, 512],512)
+        # 3*(5+num_classes)=3*(5+20)=3*(4+1+20)=75
+        final_out_filter1 =  num_anchors * (5 + num_classes)
+        self.yolo_head2 = yolo_head([512, final_out_filter1],256)
+
+
+        self.down_sample2 = conv2d(256,512,3,stride=2)
+        self.make_five_conv4 = make_five_conv([512, 1024],1024)
+        # 3*(5+num_classes)=3*(5+20)=3*(4+1+20)=75
+        final_out_filter0 =  num_anchors * (5 + num_classes)
+        self.yolo_head1 = yolo_head([1024, final_out_filter0],512)
+
+
+    def forward(self, x):
+        #  backbone
+        x2, x1, x0 = self.backbone(x)
+
+        P5 = self.conv1(x0)
+        P5 = self.SPP(P5)
+        P5 = self.conv2(P5)
+
+        P5_upsample = self.upsample1(P5)
+        P4 = self.conv_for_P4(x1)
+        P4 = torch.cat([P4,P5_upsample],axis=1)
+        P4 = self.make_five_conv1(P4)
+
+        P4_upsample = self.upsample2(P4)
+        P3 = self.conv_for_P3(x2)
+        P3 = torch.cat([P3,P4_upsample],axis=1)
+        P3 = self.make_five_conv2(P3)
+
+        P3_downsample = self.down_sample1(P3)
+        P4 = torch.cat([P3_downsample,P4],axis=1)
+        P4 = self.make_five_conv3(P4)
+
+        P4_downsample = self.down_sample2(P4)
+        P5 = torch.cat([P4_downsample,P5],axis=1)
+        P5 = self.make_five_conv4(P5)
+
+        out2 = self.yolo_head3(P3)
+        out1 = self.yolo_head2(P4)
+        out0 = self.yolo_head1(P5)
+
+        return out0, out1, out2
+

predict.py

+#-------------------------------------#
+#       对单张图片进行预测
+#-------------------------------------#
+from yolo import YOLO
+from PIL import Image
+
+yolo = YOLO()
+
+while True:
+    img = input('Input image filename:')
+    try:
+        image = Image.open(img)
+    except:
+        print('Open Error! Try again!')
+        continue
+    else:
+        r_image = yolo.detect_image(image)
+        r_image.show()

test.py

+import torch
+from torchsummary import summary
+from nets.CSPdarknet import darknet53
+from nets.yolo4 import YoloBody
+
+if __name__ == "__main__":
+    # 需要使用device来指定网络在GPU还是CPU运行
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = YoloBody(3,20).to(device)
+    summary(model, input_size=(3, 416, 416))

train.py

+#-------------------------------------#
+#       对数据集进行训练
+#-------------------------------------#
+import os
+import numpy as np
+import time
+import torch
+from torch.autograd import Variable
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+import torch.backends.cudnn as cudnn
+from nets.yolo_training import YOLOLoss,Generator
+from nets.yolo4 import YoloBody
+
+
+#---------------------------------------------------#
+#   获得类和先验框
+#---------------------------------------------------#
+def get_classes(classes_path):
+    '''loads the classes'''
+    with open(classes_path) as f:
+        class_names = f.readlines()
+    class_names = [c.strip() for c in class_names]
+    return class_names
+
+def get_anchors(anchors_path):
+    '''loads the anchors from a file'''
+    with open(anchors_path) as f:
+        anchors = f.readline()
+    anchors = [float(x) for x in anchors.split(',')]
+    return np.array(anchors).reshape([-1,3,2])[::-1,:,:]
+
+def fit_ont_epoch(net,yolo_losses,epoch,epoch_size,epoch_size_val,gen,genval,Epoch,cuda):
+    total_loss = 0
+    val_loss = 0
+    for iteration in range(epoch_size):
+        start_time = time.time()
+        images, targets = next(gen)
+        with torch.no_grad():
+            if cuda:
+                images = Variable(torch.from_numpy(images).type(torch.FloatTensor)).cuda()
+                targets = [Variable(torch.from_numpy(ann).type(torch.FloatTensor)) for ann in targets]
+            else:
+                images = Variable(torch.from_numpy(images).type(torch.FloatTensor))
+                targets = [Variable(torch.from_numpy(ann).type(torch.FloatTensor)) for ann in targets]
+        # print(images)
+        optimizer.zero_grad()
+        outputs = net(images)
+        losses = []
+        for i in range(3):
+            loss_item = yolo_losses[i](outputs[i], targets)
+            losses.append(loss_item[0])
+        loss = sum(losses)
+        loss.backward()
+        optimizer.step()
+
+        total_loss += loss
+        waste_time = time.time() - start_time
+        print('\nEpoch:'+ str(epoch+1) + '/' + str(Epoch))
+        print('iter:' + str(iteration) + '/' + str(epoch_size) + ' || Total Loss: %.4f || %.4fs/step' % (total_loss/(iteration+1),waste_time))
+
+    print('Start Validation')
+    for iteration in range(epoch_size_val):
+        images_val, targets_val = next(genval)
+
+        with torch.no_grad():
+            if cuda:
+                images = Variable(torch.from_numpy(images).cuda().type(torch.FloatTensor))
+                targets = [Variable(torch.from_numpy(ann).type(torch.FloatTensor)) for ann in targets]
+            else:
+                images = Variable(torch.from_numpy(images).type(torch.FloatTensor))
+                targets = [Variable(torch.from_numpy(ann).type(torch.FloatTensor)) for ann in targets]
+            optimizer.zero_grad()
+            outputs = net(images_val)
+            losses = []
+            for i in range(3):
+                loss_item = yolo_losses[i](outputs[i], targets_val)
+                losses.append(loss_item[0])
+            loss = sum(losses)
+            val_loss += loss
+    print('Finish Validation')
+    print('\nEpoch:'+ str(epoch+1) + '/' + str(Epoch))
+    print('Total Loss: %.4f || Val Loss: %.4f ' % (total_loss/(epoch_size+1),val_loss/(epoch_size_val+1)))
+
+    print('Saving state, iter:', str(epoch+1))
+    torch.save(model.state_dict(), 'logs/Epoch%d-Total_Loss%.4f-Val_Loss%.4f.pth'%((epoch+1),total_loss/(epoch_size+1),val_loss/(epoch_size_val+1)))
+
+
+if __name__ == "__main__":
+    #-------------------------------#
+    #   输入的shape大小
+    #   显存比较小可以使用416x416
+    #   显存比较大可以使用608x608
+    #-------------------------------#
+    input_shape = (416,416)
+    #-------------------------------#
+    #   tricks的使用设置
+    #-------------------------------#
+    Cosine_lr = False
+    mosaic = True
+    # 用于设定是否使用cuda
+    Cuda = True
+    smoooth_label = 0
+
+    annotation_path = '2007_train.txt'
+    #-------------------------------#
+    #   获得先验框和类
+    #-------------------------------#
+    anchors_path = 'model_data/yolo_anchors.txt'
+    classes_path = 'model_data/voc_classes.txt'   
+    class_names = get_classes(classes_path)
+    anchors = get_anchors(anchors_path)
+    num_classes = len(class_names)
+    
+    # 创建模型
+    model = YoloBody(len(anchors[0]),num_classes)
+    model_path = "model_data/yolo4_weights.pth"
+    # 加快模型训练的效率
+    print('Loading weights into state dict...')
+    model_dict = model.state_dict()
+    pretrained_dict = torch.load(model_path)
+    pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) ==  np.shape(v)}
+    model_dict.update(pretrained_dict)
+    model.load_state_dict(model_dict)
+    print('Finished!')
+
+    net = model.train()
+
+    if Cuda:
+        net = torch.nn.DataParallel(model)
+        cudnn.benchmark = True
+        net = net.cuda()
+
+    # 建立loss函数
+    yolo_losses = []
+    for i in range(3):
+        yolo_losses.append(YOLOLoss(np.reshape(anchors,[-1,2]),num_classes, \
+                                (input_shape[1], input_shape[0]), smoooth_label, Cuda))
+
+    # 0.1用于验证，0.9用于训练
+    val_split = 0.1
+    with open(annotation_path) as f:
+        lines = f.readlines()
+    np.random.seed(10101)
+    np.random.shuffle(lines)
+    np.random.seed(None)
+    num_val = int(len(lines)*val_split)
+    num_train = len(lines) - num_val
+    
+    if True:
+        lr = 1e-3
+        Batch_size = 4
+        Init_Epoch = 0
+        Freeze_Epoch = 25
+        
+        optimizer = optim.Adam(net.parameters(),lr)
+        if Cosine_lr:
+            lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=5, eta_min=1e-5)
+        else:
+            lr_scheduler = optim.lr_scheduler.StepLR(optimizer,step_size=1,gamma=0.95)
+
+        gen = Generator(Batch_size, lines[:num_train],
+                        (input_shape[0], input_shape[1])).generate(mosaic = mosaic)
+        gen_val = Generator(Batch_size, lines[num_train:],
+                        (input_shape[0], input_shape[1])).generate(mosaic = False)
+                        
+        epoch_size = int(max(1, num_train//Batch_size//2.5)) if mosaic else max(1, num_train//Batch_size)
+        epoch_size_val = num_val//Batch_size
+        #------------------------------------#
+        #   冻结一定部分训练
+        #------------------------------------#
+        for param in model.backbone.parameters():
+            param.requires_grad = False
+
+        for epoch in range(Init_Epoch,Freeze_Epoch):
+            fit_ont_epoch(net,yolo_losses,epoch,epoch_size,epoch_size_val,gen,gen_val,Freeze_Epoch,Cuda)
+            lr_scheduler.step()
+
+    if True:
+        lr = 1e-4
+        Batch_size = 2
+        Freeze_Epoch = 25
+        Unfreeze_Epoch = 50
+
+        optimizer = optim.Adam(net.parameters(),lr)
+        if Cosine_lr:
+            lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=5, eta_min=1e-5)
+        else:
+            lr_scheduler = optim.lr_scheduler.StepLR(optimizer,step_size=1,gamma=0.95)
+
+        gen = Generator(Batch_size, lines[:num_train],
+                        (input_shape[0], input_shape[1])).generate(mosaic = mosaic)
+        gen_val = Generator(Batch_size, lines[num_train:],
+                        (input_shape[0], input_shape[1])).generate(mosaic = False)
+                        
+        epoch_size = int(max(1, num_train//Batch_size//2.5)) if mosaic else max(1, num_train//Batch_size)
+        epoch_size_val = num_val//Batch_size
+        #------------------------------------#
+        #   解冻后训练
+        #------------------------------------#
+        for param in model.backbone.parameters():
+            param.requires_grad = True
+
+        for epoch in range(Freeze_Epoch,Unfreeze_Epoch):
+            fit_ont_epoch(net,yolo_losses,epoch,epoch_size,epoch_size_val,gen,gen_val,Unfreeze_Epoch,Cuda)
+            lr_scheduler.step()
\ No newline at end of file

utils/utils.py

+from __future__ import division
+import os
+import math
+import time
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+import matplotlib.pyplot as plt
+
+class DecodeBox(nn.Module):
+    def __init__(self, anchors, num_classes, img_size):
+        super(DecodeBox, self).__init__()
+        self.anchors = anchors
+        self.num_anchors = len(anchors)
+        self.num_classes = num_classes
+        self.bbox_attrs = 5 + num_classes
+        self.img_size = img_size
+
+    def forward(self, input):
+        # input为bs,3*(1+4+num_classes),13,13
+
+        # 一共多少张图片
+        batch_size = input.size(0)
+        # 13，13
+        input_height = input.size(2)
+        input_width = input.size(3)
+
+        # 计算步长
+        # 每一个特征点对应原来的图片上多少个像素点
+        # 如果特征层为13x13的话，一个特征点就对应原来的图片上的32个像素点
+        # 416/13 = 32
+        stride_h = self.img_size[1] / input_height
+        stride_w = self.img_size[0] / input_width
+
+        # 把先验框的尺寸调整成特征层大小的形式
+        # 计算出先验框在特征层上对应的宽高
+        scaled_anchors = [(anchor_width / stride_w, anchor_height / stride_h) for anchor_width, anchor_height in self.anchors]
+
+        # bs,3*(5+num_classes),13,13 -> bs,3,13,13,(5+num_classes)
+        prediction = input.view(batch_size, self.num_anchors,
+                                self.bbox_attrs, input_height, input_width).permute(0, 1, 3, 4, 2).contiguous()
+
+        # 先验框的中心位置的调整参数
+        x = torch.sigmoid(prediction[..., 0])  
+        y = torch.sigmoid(prediction[..., 1])
+        # 先验框的宽高调整参数
+        w = prediction[..., 2]  # Width
+        h = prediction[..., 3]  # Height
+
+        # 获得置信度，是否有物体
+        conf = torch.sigmoid(prediction[..., 4])
+        # 种类置信度
+        pred_cls = torch.sigmoid(prediction[..., 5:])  # Cls pred.
+
+        FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
+        LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor
+
+        # 生成网格，先验框中心，网格左上角 batch_size,3,13,13
+        grid_x = torch.linspace(0, input_width - 1, input_width).repeat(input_width, 1).repeat(
+            batch_size * self.num_anchors, 1, 1).view(x.shape).type(FloatTensor)
+        grid_y = torch.linspace(0, input_height - 1, input_height).repeat(input_height, 1).t().repeat(
+            batch_size * self.num_anchors, 1, 1).view(y.shape).type(FloatTensor)
+
+        # 生成先验框的宽高
+        anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0]))
+        anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1]))
+        anchor_w = anchor_w.repeat(batch_size, 1).repeat(1, 1, input_height * input_width).view(w.shape)
+        anchor_h = anchor_h.repeat(batch_size, 1).repeat(1, 1, input_height * input_width).view(h.shape)
+        
+        # 计算调整后的先验框中心与宽高
+        pred_boxes = FloatTensor(prediction[..., :4].shape)
+        pred_boxes[..., 0] = x.data + grid_x
+        pred_boxes[..., 1] = y.data + grid_y
+        pred_boxes[..., 2] = torch.exp(w.data) * anchor_w
+        pred_boxes[..., 3] = torch.exp(h.data) * anchor_h
+
+        # fig = plt.figure()
+        # ax = fig.add_subplot(121)
+        # if input_height==13:
+        #     plt.ylim(0,13)
+        #     plt.xlim(0,13)
+        # elif input_height==26:
+        #     plt.ylim(0,26)
+        #     plt.xlim(0,26)
+        # elif input_height==52:
+        #     plt.ylim(0,52)
+        #     plt.xlim(0,52)
+        # plt.scatter(grid_x.cpu(),grid_y.cpu())
+
+        # anchor_left = grid_x - anchor_w/2 
+        # anchor_top = grid_y - anchor_h/2 
+
+        # rect1 = plt.Rectangle([anchor_left[0,0,5,5],anchor_top[0,0,5,5]],anchor_w[0,0,5,5],anchor_h[0,0,5,5],color="r",fill=False)
+        # rect2 = plt.Rectangle([anchor_left[0,1,5,5],anchor_top[0,1,5,5]],anchor_w[0,1,5,5],anchor_h[0,1,5,5],color="r",fill=False)
+        # rect3 = plt.Rectangle([anchor_left[0,2,5,5],anchor_top[0,2,5,5]],anchor_w[0,2,5,5],anchor_h[0,2,5,5],color="r",fill=False)
+
+        # ax.add_patch(rect1)
+        # ax.add_patch(rect2)
+        # ax.add_patch(rect3)
+
+        # ax = fig.add_subplot(122)
+        # if input_height==13:
+        #     plt.ylim(0,13)
+        #     plt.xlim(0,13)
+        # elif input_height==26:
+        #     plt.ylim(0,26)
+        #     plt.xlim(0,26)
+        # elif input_height==52:
+        #     plt.ylim(0,52)
+        #     plt.xlim(0,52)
+        # plt.scatter(grid_x.cpu(),grid_y.cpu())
+        # plt.scatter(pred_boxes[0,:,5,5,0].cpu(),pred_boxes[0,:,5,5,1].cpu(),c='r')
+
+        # pre_left = pred_boxes[...,0] - pred_boxes[...,2]/2 
+        # pre_top = pred_boxes[...,1] - pred_boxes[...,3]/2 
+
+        # rect1 = plt.Rectangle([pre_left[0,0,5,5],pre_top[0,0,5,5]],pred_boxes[0,0,5,5,2],pred_boxes[0,0,5,5,3],color="r",fill=False)
+        # rect2 = plt.Rectangle([pre_left[0,1,5,5],pre_top[0,1,5,5]],pred_boxes[0,1,5,5,2],pred_boxes[0,1,5,5,3],color="r",fill=False)
+        # rect3 = plt.Rectangle([pre_left[0,2,5,5],pre_top[0,2,5,5]],pred_boxes[0,2,5,5,2],pred_boxes[0,2,5,5,3],color="r",fill=False)
+
+        # ax.add_patch(rect1)
+        # ax.add_patch(rect2)
+        # ax.add_patch(rect3)
+
+        # plt.show()
+        # 用于将输出调整为相对于416x416的大小
+        _scale = torch.Tensor([stride_w, stride_h] * 2).type(FloatTensor)
+        output = torch.cat((pred_boxes.view(batch_size, -1, 4) * _scale,
+                            conf.view(batch_size, -1, 1), pred_cls.view(batch_size, -1, self.num_classes)), -1)
+        return output.data
+        
+def letterbox_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
+
+def yolo_correct_boxes(top, left, bottom, right, input_shape, image_shape):
+    new_shape = image_shape*np.min(input_shape/image_shape)
+
+    offset = (input_shape-new_shape)/2./input_shape
+    scale = input_shape/new_shape
+
+    box_yx = np.concatenate(((top+bottom)/2,(left+right)/2),axis=-1)/input_shape
+    box_hw = np.concatenate((bottom-top,right-left),axis=-1)/input_shape
+
+    box_yx = (box_yx - offset) * scale
+    box_hw *= scale
+
+    box_mins = box_yx - (box_hw / 2.)
+    box_maxes = box_yx + (box_hw / 2.)
+    boxes =  np.concatenate([
+        box_mins[:, 0:1],
+        box_mins[:, 1:2],
+        box_maxes[:, 0:1],
+        box_maxes[:, 1:2]
+    ],axis=-1)
+    print(np.shape(boxes))
+    boxes *= np.concatenate([image_shape, image_shape],axis=-1)
+    return boxes
+
+def bbox_iou(box1, box2, x1y1x2y2=True):
+    """
+        计算IOU
+    """
+    if not x1y1x2y2:
+        b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2
+        b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2
+        b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2
+        b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2
+    else:
+        b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
+        b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]
+
+    inter_rect_x1 = torch.max(b1_x1, b2_x1)
+    inter_rect_y1 = torch.max(b1_y1, b2_y1)
+    inter_rect_x2 = torch.min(b1_x2, b2_x2)
+    inter_rect_y2 = torch.min(b1_y2, b2_y2)
+
+    inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1 + 1, min=0) * \
+                 torch.clamp(inter_rect_y2 - inter_rect_y1 + 1, min=0)
+                 
+    b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1)
+    b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1)
+
+    iou = inter_area / (b1_area + b2_area - inter_area + 1e-16)
+
+    return iou
+
+
+def non_max_suppression(prediction, num_classes, conf_thres=0.5, nms_thres=0.4):
+    # 求左上角和右下角
+    box_corner = prediction.new(prediction.shape)
+    box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2
+    box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2
+    box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2
+    box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2
+    prediction[:, :, :4] = box_corner[:, :, :4]
+
+    output = [None for _ in range(len(prediction))]
+    for image_i, image_pred in enumerate(prediction):
+        # 利用置信度进行第一轮筛选
+        conf_mask = (image_pred[:, 4] >= conf_thres).squeeze()
+        image_pred = image_pred[conf_mask]
+
+        if not image_pred.size(0):
+            continue
+
+        # 获得种类及其置信度
+        class_conf, class_pred = torch.max(image_pred[:, 5:5 + num_classes], 1, keepdim=True)
+
+        # 获得的内容为(x1, y1, x2, y2, obj_conf, class_conf, class_pred)
+        detections = torch.cat((image_pred[:, :5], class_conf.float(), class_pred.float()), 1)
+
+        # 获得种类
+        unique_labels = detections[:, -1].cpu().unique()
+
+        if prediction.is_cuda:
+            unique_labels = unique_labels.cuda()
+
+        for c in unique_labels:
+            # 获得某一类初步筛选后全部的预测结果
+            detections_class = detections[detections[:, -1] == c]
+            # 按照存在物体的置信度排序
+            _, conf_sort_index = torch.sort(detections_class[:, 4], descending=True)
+            detections_class = detections_class[conf_sort_index]
+            # 进行非极大抑制
+            max_detections = []
+            while detections_class.size(0):
+                # 取出这一类置信度最高的，一步一步往下判断，判断重合程度是否大于nms_thres，如果是则去除掉
+                max_detections.append(detections_class[0].unsqueeze(0))
+                if len(detections_class) == 1:
+                    break
+                ious = bbox_iou(max_detections[-1], detections_class[1:])
+                detections_class = detections_class[1:][ious < nms_thres]
+            # 堆叠
+            max_detections = torch.cat(max_detections).data
+            # Add max detections to outputs
+            output[image_i] = max_detections if output[image_i] is None else torch.cat(
+                (output[image_i], max_detections))
+
+    return output
+
+def merge_bboxes(bboxes, cutx, cuty):
+    merge_bbox = []
+    for i in range(len(bboxes)):
+        for box in bboxes[i]:
+            tmp_box = []
+            x1,y1,x2,y2 = box[0], box[1], box[2], box[3]
+
+            if i == 0:
+                if y1 > cuty or x1 > cutx:
+                    continue
+                if y2 >= cuty and y1 <= cuty:
+                    y2 = cuty
+                    if y2-y1 < 5:
+                        continue
+                if x2 >= cutx and x1 <= cutx:
+                    x2 = cutx
+                    if x2-x1 < 5:
+                        continue
+                
+            if i == 1:
+                if y2 < cuty or x1 > cutx:
+                    continue
+
+                if y2 >= cuty and y1 <= cuty:
+                    y1 = cuty
+                    if y2-y1 < 5:
+                        continue
+                
+                if x2 >= cutx and x1 <= cutx:
+                    x2 = cutx
+                    if x2-x1 < 5:
+                        continue
+
+            if i == 2:
+                if y2 < cuty or x2 < cutx:
+                    continue
+
+                if y2 >= cuty and y1 <= cuty:
+                    y1 = cuty
+                    if y2-y1 < 5:
+                        continue
+
+                if x2 >= cutx and x1 <= cutx:
+                    x1 = cutx
+                    if x2-x1 < 5:
+                        continue
+
+            if i == 3:
+                if y1 > cuty or x2 < cutx:
+                    continue
+
+                if y2 >= cuty and y1 <= cuty:
+                    y2 = cuty
+                    if y2-y1 < 5:
+                        continue
+
+                if x2 >= cutx and x1 <= cutx:
+                    x1 = cutx
+                    if x2-x1 < 5:
+                        continue
+
+            tmp_box.append(x1)
+            tmp_box.append(y1)
+            tmp_box.append(x2)
+            tmp_box.append(y2)
+            tmp_box.append(box[-1])
+            merge_bbox.append(tmp_box)
+    return merge_bbox
\ No newline at end of file

video.py

+#-------------------------------------#
+#       调用摄像头检测
+#-------------------------------------#
+from yolo import YOLO
+from PIL import Image
+import numpy as np
+import cv2
+import time
+yolo = YOLO()
+# 调用摄像头
+capture=cv2.VideoCapture(0) # capture=cv2.VideoCapture("1.mp4")
+
+fps = 0.0
+while(True):
+    t1 = time.time()
+    # 读取某一帧
+    ref,frame=capture.read()
+    # 格式转变，BGRtoRGB
+    frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
+    # 转变成Image
+    frame = Image.fromarray(np.uint8(frame))
+
+    # 进行检测
+    frame = np.array(yolo.detect_image(frame))
+
+    # RGBtoBGR满足opencv显示格式
+    frame = cv2.cvtColor(frame,cv2.COLOR_RGB2BGR)
+
+    fps  = ( fps + (1./(time.time()-t1)) ) / 2
+    print("fps= %.2f"%(fps))
+    frame = cv2.putText(frame, "fps= %.2f"%(fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+
+    cv2.imshow("video",frame)
+
+
+    c= cv2.waitKey(30) & 0xff 
+    if c==27:
+        capture.release()
+        break

voc_annotation.py

+import xml.etree.ElementTree as ET
+from os import getcwd
+
+sets=[('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
+
+wd = getcwd()
+classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
+
+def convert_annotation(year, image_id, list_file):
+    in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml'%(year, image_id))
+    tree=ET.parse(in_file)
+    root = tree.getroot()
+    if root.find('object')==None:
+        return
+    list_file.write('%s/VOCdevkit/VOC%s/JPEGImages/%s.jpg'%(wd, year, image_id))
+    for obj in root.iter('object'):
+        difficult = obj.find('difficult').text
+        cls = obj.find('name').text
+        if cls not in classes or int(difficult)==1:
+            continue
+        cls_id = classes.index(cls)
+        xmlbox = obj.find('bndbox')
+        b = (int(xmlbox.find('xmin').text), int(xmlbox.find('ymin').text), int(xmlbox.find('xmax').text), int(xmlbox.find('ymax').text))
+        list_file.write(" " + ",".join([str(a) for a in b]) + ',' + str(cls_id))
+
+    list_file.write('\n')
+
+for year, image_set in sets:
+    image_ids = open('VOCdevkit/VOC%s/ImageSets/Main/%s.txt'%(year, image_set)).read().strip().split()
+    list_file = open('%s_%s.txt'%(year, image_set), 'w')
+    for image_id in image_ids:
+        convert_annotation(year, image_id, list_file)
+    list_file.close()

yolo.py

+#-------------------------------------#
+#       创建YOLO类
+#-------------------------------------#
+import cv2
+import numpy as np
+import colorsys
+import os
+import torch
+import torch.nn as nn
+from nets.yolo4 import YoloBody
+import torch.backends.cudnn as cudnn
+from PIL import Image,ImageFont, ImageDraw
+from torch.autograd import Variable
+from utils.utils import non_max_suppression, bbox_iou, DecodeBox,letterbox_image,yolo_correct_boxes
+
+class YOLO(object):
+    _defaults = {
+        "model_path": 'model_data/yolo4_voc_weights.pth',
+        "anchors_path": 'model_data/yolo_anchors.txt',
+        "classes_path": 'model_data/voc_classes.txt',
+        "model_image_size" : (416, 416, 3),
+        "confidence": 0.5,
+        "cuda": True
+    }
+
+    @classmethod
+    def get_defaults(cls, n):
+        if n in cls._defaults:
+            return cls._defaults[n]
+        else:
+            return "Unrecognized attribute name '" + n + "'"
+
+    #---------------------------------------------------#
+    #   初始化YOLO
+    #---------------------------------------------------#
+    def __init__(self, **kwargs):
+        self.__dict__.update(self._defaults)
+        self.class_names = self._get_class()
+        self.anchors = self._get_anchors()
+        self.generate()
+    #---------------------------------------------------#
+    #   获得所有的分类
+    #---------------------------------------------------#
+    def _get_class(self):
+        classes_path = os.path.expanduser(self.classes_path)
+        with open(classes_path) as f:
+            class_names = f.readlines()
+        class_names = [c.strip() for c in class_names]
+        return class_names
+    
+    #---------------------------------------------------#
+    #   获得所有的先验框
+    #---------------------------------------------------#
+    def _get_anchors(self):
+        anchors_path = os.path.expanduser(self.anchors_path)
+        with open(anchors_path) as f:
+            anchors = f.readline()
+        anchors = [float(x) for x in anchors.split(',')]
+        return np.array(anchors).reshape([-1, 3, 2])[::-1,:,:]
+
+    #---------------------------------------------------#
+    #   获得所有的分类
+    #---------------------------------------------------#
+    def generate(self):
+        os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+        self.net = YoloBody(len(self.anchors[0]),len(self.class_names)).eval()
+
+        # 加快模型训练的效率
+        print('Loading weights into state dict...')
+        state_dict = torch.load(self.model_path)
+        self.net.load_state_dict(state_dict)
+        self.net = nn.DataParallel(self.net)
+        if self.cuda:
+            self.net = self.net.cuda()
+    
+        print('Finished!')
+
+        self.yolo_decodes = []
+        for i in range(3):
+            self.yolo_decodes.append(DecodeBox(self.anchors[i], len(self.class_names),  (self.model_image_size[1], self.model_image_size[0])))
+
+
+        print('{} model, anchors, and classes loaded.'.format(self.model_path))
+        # 画框设置不同的颜色
+        hsv_tuples = [(x / len(self.class_names), 1., 1.)
+                      for x in range(len(self.class_names))]
+        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
+        self.colors = list(
+            map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
+                self.colors))
+
+    #---------------------------------------------------#
+    #   检测图片
+    #---------------------------------------------------#
+    def detect_image(self, image):
+        image_shape = np.array(np.shape(image)[0:2])
+
+        crop_img = np.array(letterbox_image(image, (self.model_image_size[0],self.model_image_size[1])))
+        photo = np.array(crop_img,dtype = np.float32)
+        photo /= 255.0
+        photo = np.transpose(photo, (2, 0, 1))
+        photo = photo.astype(np.float32)
+        images = []
+        images.append(photo)
+        images = np.asarray(images)
+
+        with torch.no_grad():
+            images = torch.from_numpy(images)
+            if self.cuda:
+                images = images.cuda()
+            outputs = self.net(images)
+            
+        output_list = []
+        for i in range(3):
+            output_list.append(self.yolo_decodes[i](outputs[i]))
+        output = torch.cat(output_list, 1)
+        batch_detections = non_max_suppression(output, len(self.class_names),
+                                                conf_thres=self.confidence,
+                                                nms_thres=0.3)
+        try:
+            batch_detections = batch_detections[0].cpu().numpy()
+        except:
+            return image
+            
+        top_index = batch_detections[:,4]*batch_detections[:,5] > self.confidence
+        top_conf = batch_detections[top_index,4]*batch_detections[top_index,5]
+        top_label = np.array(batch_detections[top_index,-1],np.int32)
+        top_bboxes = np.array(batch_detections[top_index,:4])
+        top_xmin, top_ymin, top_xmax, top_ymax = np.expand_dims(top_bboxes[:,0],-1),np.expand_dims(top_bboxes[:,1],-1),np.expand_dims(top_bboxes[:,2],-1),np.expand_dims(top_bboxes[:,3],-1)
+
+        # 去掉灰条
+        boxes = yolo_correct_boxes(top_ymin,top_xmin,top_ymax,top_xmax,np.array([self.model_image_size[0],self.model_image_size[1]]),image_shape)
+
+        font = ImageFont.truetype(font='model_data/simhei.ttf',size=np.floor(3e-2 * np.shape(image)[1] + 0.5).astype('int32'))
+
+        thickness = (np.shape(image)[0] + np.shape(image)[1]) // self.model_image_size[0]
+
+        for i, c in enumerate(top_label):
+            predicted_class = self.class_names[c]
+            score = top_conf[i]
+
+            top, left, bottom, right = boxes[i]
+            top = top - 5
+            left = left - 5
+            bottom = bottom + 5
+            right = right + 5
+
+            top = max(0, np.floor(top + 0.5).astype('int32'))
+            left = max(0, np.floor(left + 0.5).astype('int32'))
+            bottom = min(np.shape(image)[0], np.floor(bottom + 0.5).astype('int32'))
+            right = min(np.shape(image)[1], np.floor(right + 0.5).astype('int32'))
+
+            # 画框框
+            label = '{} {:.2f}'.format(predicted_class, score)
+            draw = ImageDraw.Draw(image)
+            label_size = draw.textsize(label, font)
+            label = label.encode('utf-8')
+            print(label)
+            
+            if top - label_size[1] >= 0:
+                text_origin = np.array([left, top - label_size[1]])
+            else:
+                text_origin = np.array([left, top + 1])
+
+            for i in range(thickness):
+                draw.rectangle(
+                    [left + i, top + i, right - i, bottom - i],
+                    outline=self.colors[self.class_names.index(predicted_class)])
+            draw.rectangle(
+                [tuple(text_origin), tuple(text_origin + label_size)],
+                fill=self.colors[self.class_names.index(predicted_class)])
+            draw.text(text_origin, str(label,'UTF-8'), fill=(0, 0, 0), font=font)
+            del draw
+        return image
+