first commit

.idea/.gitignore

@@ -0,0 +1,3 @@
+# Default ignored files
+/shelf/
+/workspace.xml

.idea/MPKFNet.iml

@@ -0,0 +1,12 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="jdk" jdkName="Python 3.6 (py36pt18)" jdkType="Python SDK" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+  <component name="PyDocumentationSettings">
+    <option name="format" value="GOOGLE" />
+    <option name="myDocStringFormat" value="Google" />
+  </component>
+</module>

.idea/inspectionProfiles/Project_Default.xml

@@ -0,0 +1,14 @@
+<component name="InspectionProjectProfileManager">
+  <profile version="1.0">
+    <option name="myName" value="Project Default" />
+    <inspection_tool class="PyPackageRequirementsInspection" enabled="true" level="WARNING" enabled_by_default="true">
+      <option name="ignoredPackages">
+        <value>
+          <list size="1">
+            <item index="0" class="java.lang.String" itemvalue="glob3" />
+          </list>
+        </value>
+      </option>
+    </inspection_tool>
+  </profile>
+</component>

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/misc.xml

@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.6 (py36pt18)" project-jdk-type="Python SDK" />
+</project>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/MPKFNet.iml" filepath="$PROJECT_DIR$/.idea/MPKFNet.iml" />
+    </modules>
+  </component>
+</project>

.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="$PROJECT_DIR$" vcs="Git" />
+  </component>
+</project>

Check_VOC_Hatcher.py

@@ -0,0 +1,56 @@
+'''
+author ：  Hatcher W
+Email： hcwang@semi.ac.cn
+use: 用于检测voc数据集是否正确，
+input:（图像的文件夹路径，图像标签文件夹路径）
+output：check路径，所有的可视化结果都在里面
+'''
+import os
+import cv2
+import  xml.dom.minidom
+from tqdm import tqdm
+
+image_path="./JPEGImages/"
+annotation_path="./Annatations/"
+save_path ="./check/"
+num = 0
+files_name = os.listdir(image_path)
+for filename_ in tqdm(files_name):
+    filename, extension= os.path.splitext(filename_)
+    #filename=filename.split('__')[0]
+    img_path =image_path+filename+'.jpg'
+    xml_path =annotation_path+filename+'.xml'
+    img = cv2.imread(img_path)
+    if img is None:
+        pass
+    try:
+        dom = xml.dom.minidom.parse(xml_path)
+    except:
+        #os.remove(img_path)
+        continue
+    root = dom.documentElement
+    objects=dom.getElementsByTagName("object")
+    for object in objects:
+        bndbox = object.getElementsByTagName('bndbox')[0]
+        xmin = bndbox.getElementsByTagName('xmin')[0]
+        ymin = bndbox.getElementsByTagName('ymin')[0]
+        xmax = bndbox.getElementsByTagName('xmax')[0]
+        ymax = bndbox.getElementsByTagName('ymax')[0]
+        #print(xmin.childNodes[0].data)
+        #print(type(xmin.childNodes[0].data))
+        xmin_data=int(xmin.childNodes[0].data)
+        ymin_data=int(ymin.childNodes[0].data)
+        xmax_data=int(xmax.childNodes[0].data)
+        ymax_data=int(ymax.childNodes[0].data)
+        label_name=object.getElementsByTagName('name')[0].childNodes[0].data
+        cv2.rectangle(img,(xmin_data,ymin_data),(xmax_data,ymax_data),(55,255,155),1)
+        num +=1
+        cv2.putText(img,label_name,(xmin_data,ymin_data-3),cv2.FONT_HERSHEY_SIMPLEX,0.5,(0,255,0),2)
+    flag=0
+    flag=cv2.imwrite(save_path+"{}.jpg".format(filename),img)
+    if not (flag):
+        print(filename,"error")
+print("all done ====================================")
+print("number = ",num)
+
+#(int((xmin_data+xmax_data/2)),int((ymin_data+ymax_data)/2))

Dcl/DclConfig.py

@@ -0,0 +1,137 @@
+import os
+import pandas as pd
+import torch
+
+from Dcl.transforms import transforms
+
+
+
+# pretrained model checkpoints
+pretrained_model = {'resnet50' : '/data1/Hatcher/model/data_MPKFNet/weigth/FinedGrainedRecognition/pretrained/vgg16-397923af.pth',}
+
+# transforms dict
+def load_data_transformers(resize_reso=512, crop_reso=448, swap_num=[7, 7]):
+    center_resize = 600
+    Normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    data_transforms = {
+       	'swap': transforms.Compose([
+            transforms.Randomswap((swap_num[0], swap_num[1])),
+        ]),
+        'common_aug': transforms.Compose([
+            transforms.Resize((resize_reso, resize_reso)),
+            transforms.RandomRotation(degrees=15), # 将图片进行随机旋转
+            transforms.RandomCrop((crop_reso,crop_reso)), # 图片进行随机裁剪
+            transforms.RandomHorizontalFlip(), # 随机水平翻转
+        ]),
+        'train_totensor': transforms.Compose([
+            transforms.Resize((crop_reso, crop_reso)),
+            # ImageNetPolicy(),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]),
+        'val_totensor': transforms.Compose([
+            transforms.Resize((crop_reso, crop_reso)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]),
+        'test_totensor': transforms.Compose([
+            transforms.Resize((resize_reso, resize_reso)),
+            transforms.CenterCrop((crop_reso, crop_reso)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]),
+        'None': None,
+    }
+    return data_transforms
+
+''' change three paths dataset save '''
+
+class LoadConfig(object):
+    def __init__(self, args, version):
+        get_list = []
+        if version == 'train':
+            get_list = ['train', 'val']
+        elif version == 'val':
+            get_list = ['val']
+        elif version == 'test':
+            get_list = ['test']
+
+
+        ###############################
+        #### add dataset info here ####
+        ###############################
+
+        # put image data in $PATH/data
+        # put annotation txt file in $PATH/anno
+
+        if args.dataset == 'product':
+            self.dataset = args.dataset
+            self.rawdata_root = './../FGVC_product/data'
+            self.anno_root = './../FGVC_product/anno'
+            self.numcls = 2019
+        elif args.dataset == 'CUB':
+            self.dataset = args.dataset
+            self.rawdata_root = './datasets/CUB_200_2011/data'
+            self.anno_root = './datasets/CUB_200_2011/anno'
+            self.numcls = 200
+        elif args.dataset == 'STCAR':
+            self.dataset = args.dataset
+            self.rawdata_root = './dataset/st_car/data'
+            self.anno_root = './dataset/st_car/anno'
+            self.numcls = 196
+        elif args.dataset == 'AIR':
+            self.dataset = args.dataset
+            self.rawdata_root = './dataset/aircraft/data'
+            self.anno_root = './dataset/aircraft/anno'
+            self.numcls = 200
+
+        elif args.dataset == 'HCData':
+            self.dataset = args.dataset
+            self.rawdata_root = '../../data_MPKFNet/data/DCL/datasets/HCData/data'
+            self.anno_root = '../../data_MPKFNet/data/DCL/datasets/HCData/anno'
+            self.numcls = 50
+        else:
+            raise Exception('dataset not defined ???')
+
+        # annotation file organized as :
+        # path/image_name cls_num\n
+
+        if 'train' in get_list:
+             self.train_anno = pd.read_csv(os.path.join(self.anno_root, 'train.txt'),
+                                           sep=" ",
+                                           header=None,
+                                           names=['ImageName', 'label'])
+
+        if 'val' in get_list:
+            self.val_anno = pd.read_csv(os.path.join(self.anno_root, 'test.txt'),
+                                           sep=" ",
+                                           header=None,
+                                           names=['ImageName', 'label'])
+
+        if 'test' in get_list:
+            self.test_anno = pd.read_csv(os.path.join(self.anno_root, 'test.txt'),
+                                           sep=" ",
+                                           header=None,
+                                           names=['ImageName', 'label'])
+
+        self.swap_num = args.swap_num
+        self.backbone = args.backbone
+
+        self.use_dcl = True
+        self.use_backbone = False if self.use_dcl else True
+        self.use_Asoftmax = False
+        self.use_focal_loss = False
+        self.use_fpn = False
+        self.use_hier = False
+
+        self.weighted_sample = False
+        self.cls_2 = True
+        self.cls_2xmul = False
+
+        self.log_folder = './logs'
+        if not os.path.exists(self.log_folder):
+            os.mkdir(self.log_folder)
+
+
+
+

Dcl/config.py

@@ -0,0 +1,140 @@
+import os
+import pandas as pd
+import torch
+
+from transforms import transforms
+from util.autoaugment import ImageNetPolicy
+
+# pretrained model checkpoints
+pretrained_model = {'resnet50' : '/data1/Hatcher/model/data_MPKFNet/weigth/FinedGrainedRecognition/pretrained/vgg16-397923af.pth',}
+
+# transforms dict
+def load_data_transformers(resize_reso=512, crop_reso=448, swap_num=[7, 7]):
+    center_resize = 600
+    Normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    data_transforms = {
+       	'swap': transforms.Compose([
+            transforms.Randomswap((swap_num[0], swap_num[1])),
+        ]),
+        'common_aug': transforms.Compose([
+            transforms.Resize((resize_reso, resize_reso)),
+            transforms.RandomRotation(degrees=15), # 将图片进行随机旋转
+            transforms.RandomCrop((crop_reso,crop_reso)), # 图片进行随机裁剪
+            transforms.RandomHorizontalFlip(), # 随机水平翻转
+        ]),
+        'train_totensor': transforms.Compose([
+            transforms.Resize((crop_reso, crop_reso)),
+            # ImageNetPolicy(),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]),
+        'val_totensor': transforms.Compose([
+            transforms.Resize((crop_reso, crop_reso)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]),
+        'test_totensor': transforms.Compose([
+            transforms.Resize((resize_reso, resize_reso)),
+            transforms.CenterCrop((crop_reso, crop_reso)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]),
+        'None': None,
+    }
+    return data_transforms
+
+''' change three paths dataset save '''
+
+class LoadConfig(object):
+    def __init__(self, args, version):
+        if version == 'train':
+            get_list = ['train', 'val']
+        elif version == 'val':
+            get_list = ['val']
+        elif version == 'test':
+            get_list = ['test']
+        else:
+            raise Exception("train/val/test ???\n")
+
+        ###############################
+        #### add dataset info here ####
+        ###############################
+
+        # put image data in $PATH/data
+        # put annotation txt file in $PATH/anno
+
+        if args.dataset == 'product':
+            self.dataset = args.dataset
+            self.rawdata_root = './../FGVC_product/data'
+            self.anno_root = './../FGVC_product/anno'
+            self.numcls = 2019
+        elif args.dataset == 'CUB':
+            self.dataset = args.dataset
+            self.rawdata_root = './datasets/CUB_200_2011/data'
+            self.anno_root = './datasets/CUB_200_2011/anno'
+            self.numcls = 200
+        elif args.dataset == 'STCAR':
+            self.dataset = args.dataset
+            self.rawdata_root = './dataset/st_car/data'
+            self.anno_root = './dataset/st_car/anno'
+            self.numcls = 196
+        elif args.dataset == 'AIR':
+            self.dataset = args.dataset
+            self.rawdata_root = './dataset/aircraft/data'
+            self.anno_root = './dataset/aircraft/anno'
+            self.numcls = 200
+
+        elif args.dataset == 'HCData':
+            self.dataset = args.dataset
+            self.rawdata_root = '../../data_MPKFNet/data/DCL/datasets/HCData/data'
+            self.anno_root = '../../data_MPKFNet/data/DCL/datasets/HCData/anno'
+            self.numcls = 50
+        else:
+            raise Exception('dataset not defined ???')
+
+        # annotation file organized as :
+        # path/image_name cls_num\n
+
+        if 'train' in get_list:
+             self.train_anno = pd.read_csv(os.path.join(self.anno_root, 'train.txt'),
+                                           sep=" ",
+                                           header=None,
+                                           names=['ImageName', 'label'])
+
+        if 'val' in get_list:
+            self.val_anno = pd.read_csv(os.path.join(self.anno_root, 'test.txt'),
+                                           sep=" ",
+                                           header=None,
+                                           names=['ImageName', 'label'])
+
+        if 'test' in get_list:
+            self.test_anno = pd.read_csv(os.path.join(self.anno_root, 'test.txt'),
+                                           sep=" ",
+                                           header=None,
+                                           names=['ImageName', 'label'])
+
+        self.swap_num = args.swap_num
+
+        self.save_dir = '../../data_MPKFNet/weigth/FinedGrainedRecognition/Train'
+        if not os.path.exists(self.save_dir):
+            os.mkdir(self.save_dir)
+        self.backbone = args.backbone
+
+        self.use_dcl = True
+        self.use_backbone = False if self.use_dcl else True
+        self.use_Asoftmax = False
+        self.use_focal_loss = False
+        self.use_fpn = False
+        self.use_hier = False
+
+        self.weighted_sample = False
+        self.cls_2 = True
+        self.cls_2xmul = False
+
+        self.log_folder = './logs'
+        if not os.path.exists(self.log_folder):
+            os.mkdir(self.log_folder)
+
+
+
+

Dcl/models/LoadModel.py

@@ -0,0 +1,80 @@
+import numpy as np
+from torch import nn
+import torch
+from torchvision import models, transforms, datasets
+import torch.nn.functional as F
+# import pretrainedmodels
+
+pretrained_model = {'resnet50' : '../data_MPKFNet/weigth/FinedGrainedRecognition/pretrained/vgg16-397923af.pth',}
+
+
+import pdb
+
+class MainModel(nn.Module):
+    def __init__(self, config):
+        super(MainModel, self).__init__()
+        self.use_dcl = config.use_dcl
+        self.num_classes = config.numcls
+        self.backbone_arch = config.backbone
+        self.use_Asoftmax = config.use_Asoftmax
+        # 返回属性值
+        if self.backbone_arch in dir(models):
+            self.model = getattr(models, 'vgg16')() # getattr models.vgg16
+            if self.backbone_arch in pretrained_model:
+                self.model.load_state_dict(torch.load(pretrained_model[self.backbone_arch])) # 加载权重
+        else:
+            # 在 Python 类的内部，无论是类属性还是实例属性，都是以字典的形式进行存储的，其中属性名作为键，而值作为该键对应的值
+            # 将键对应的键值提取出来
+            if self.backbone_arch in pretrained_model:
+                self.model = pretrainedmodels.__dict__[self.backbone_arch](num_classes=1000, pretrained=None)
+            else:
+                self.model = pretrainedmodels.__dict__[self.backbone_arch](num_classes=1000)
+        # 去掉全连接层 模型结构转换
+        if self.backbone_arch == 'resnet50' or self.backbone_arch == 'se_resnet50':
+            self.model = nn.Sequential(*list(self.model.children())[:-2])
+        if self.backbone_arch == 'senet154':
+            self.model = nn.Sequential(*list(self.model.children())[:-3])
+        if self.backbone_arch == 'se_resnext101_32x4d':
+            self.model = nn.Sequential(*list(self.model.children())[:-2])
+        if self.backbone_arch == 'se_resnet101':
+            self.model = nn.Sequential(*list(self.model.children())[:-2])
+        self.avgpool = nn.AdaptiveAvgPool2d(output_size=1)
+        self.classifier = nn.Linear(512, self.num_classes, bias=False) # 输入通道进行了改变 原来是2048
+        if self.use_dcl:
+            # if config.cls_2:
+            self.classifier_swap = nn.Linear(512, 2, bias=False)
+            if config.cls_2xmul:
+                self.classifier_swap = nn.Linear(512, 2*self.num_classes, bias=False)
+            self.Convmask = nn.Conv2d(512, 1, 1, stride=1, padding=0, bias=True) # 输入 输出 卷积核大小 步长
+            self.avgpool2 = nn.AvgPool2d(2, stride=2)
+
+        # if self.use_Asoftmax:
+        #     self.Aclassifier = AngleLinear(2048, self.num_classes, bias=False)
+# 分类网络：output[0] 对抗损失网络：output[1] 区域对齐网络：output[2]
+    def forward(self, x, last_cont=None):
+        x = self.model(x)
+        if self.use_dcl:
+            mask = self.Convmask(x) # 加一个全卷积层
+            mask = self.avgpool2(mask)# 平均池化
+            mask = torch.tanh(mask)
+            mask = mask.view(mask.size(0), -1)
+        x = self.avgpool(x) # 平均值池化
+        x = x.view(x.size(0), -1)
+        out = []
+        out.append(self.classifier(x))
+
+        if self.use_dcl:
+            out.append(self.classifier_swap(x))
+            out.append(mask)
+
+        # if self.use_Asoftmax:
+        #     if last_cont is None:
+        #         x_size = x.size(0)
+        #         out.append(self.Aclassifier(x[0:x_size:2]))
+        #     else:
+        #         last_x = self.model(last_cont)
+        #         last_x = self.avgpool(last_x)
+        #         last_x = last_x.view(last_x.size(0), -1)
+        #         out.append(self.Aclassifier(last_x))
+
+        return out

Dcl/transforms/__init__.py

@@ -0,0 +1 @@
+from .transforms import *

Dcl/transforms/functional.py

@@ -0,0 +1,750 @@
+from __future__ import division
+import torch
+import math
+import random
+from PIL import Image, ImageOps, ImageEnhance
+try:
+    import accimage
+except ImportError:
+    accimage = None
+import numpy as np
+import numbers
+import types
+import collections
+import warnings
+
+
+def _is_pil_image(img):
+    if accimage is not None:
+        return isinstance(img, (Image.Image, accimage.Image))
+    else:
+        return isinstance(img, Image.Image)
+
+
+def _is_tensor_image(img):
+    return torch.is_tensor(img) and img.ndimension() == 3
+
+
+def _is_numpy_image(img):
+    return isinstance(img, np.ndarray) and (img.ndim in {2, 3})
+
+
+def to_tensor(pic):
+    """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor.
+
+    See ``ToTensor`` for more details.
+
+    Args:
+        pic (PIL Image or numpy.ndarray): Image to be converted to tensor.
+
+    Returns:
+        Tensor: Converted image.
+    """
+    if not(_is_pil_image(pic) or _is_numpy_image(pic)):
+        raise TypeError('pic should be PIL Image or ndarray. Got {}'.format(type(pic)))
+
+    if isinstance(pic, np.ndarray):
+        # handle numpy array
+        img = torch.from_numpy(pic.transpose((2, 0, 1)))
+        # backward compatibility
+        if isinstance(img, torch.ByteTensor):
+            return img.float().div(255)
+        else:
+            return img
+
+    if accimage is not None and isinstance(pic, accimage.Image):
+        nppic = np.zeros([pic.channels, pic.height, pic.width], dtype=np.float32)
+        pic.copyto(nppic)
+        return torch.from_numpy(nppic)
+
+    # handle PIL Image
+    if pic.mode == 'I':
+        img = torch.from_numpy(np.array(pic, np.int32, copy=False))
+    elif pic.mode == 'I;16':
+        img = torch.from_numpy(np.array(pic, np.int16, copy=False))
+    elif pic.mode == 'F':
+        img = torch.from_numpy(np.array(pic, np.float32, copy=False))
+    elif pic.mode == '1':
+        img = 255 * torch.from_numpy(np.array(pic, np.uint8, copy=False))
+    else:
+        img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes()))
+    # PIL image mode: L, P, I, F, RGB, YCbCr, RGBA, CMYK
+    if pic.mode == 'YCbCr':
+        nchannel = 3
+    elif pic.mode == 'I;16':
+        nchannel = 1
+    else:
+        nchannel = len(pic.mode)
+    img = img.view(pic.size[1], pic.size[0], nchannel)
+    # put it from HWC to CHW format
+    # yikes, this transpose takes 80% of the loading time/CPU
+    img = img.transpose(0, 1).transpose(0, 2).contiguous()
+    if isinstance(img, torch.ByteTensor):
+        return img.float().div(255)
+    else:
+        return img
+
+
+def to_pil_image(pic, mode=None):
+    """Convert a tensor or an ndarray to PIL Image.
+
+    See :class:`~torchvision.transforms.ToPIlImage` for more details.
+
+    Args:
+        pic (Tensor or numpy.ndarray): Image to be converted to PIL Image.
+        mode (`PIL.Image mode`_): color space and pixel depth of input data (optional).
+
+    .. _PIL.Image mode: http://pillow.readthedocs.io/en/3.4.x/handbook/concepts.html#modes
+
+    Returns:
+        PIL Image: Image converted to PIL Image.
+    """
+    if not(_is_numpy_image(pic) or _is_tensor_image(pic)):
+        raise TypeError('pic should be Tensor or ndarray. Got {}.'.format(type(pic)))
+
+    npimg = pic
+    if isinstance(pic, torch.FloatTensor):
+        pic = pic.mul(255).byte()
+    if torch.is_tensor(pic):
+        npimg = np.transpose(pic.numpy(), (1, 2, 0))
+
+    if not isinstance(npimg, np.ndarray):
+        raise TypeError('Input pic must be a torch.Tensor or NumPy ndarray, ' +
+                        'not {}'.format(type(npimg)))
+
+    if npimg.shape[2] == 1:
+        expected_mode = None
+        npimg = npimg[:, :, 0]
+        if npimg.dtype == np.uint8:
+            expected_mode = 'L'
+        elif npimg.dtype == np.int16:
+            expected_mode = 'I;16'
+        elif npimg.dtype == np.int32:
+            expected_mode = 'I'
+        elif npimg.dtype == np.float32:
+            expected_mode = 'F'
+        if mode is not None and mode != expected_mode:
+            raise ValueError("Incorrect mode ({}) supplied for input type {}. Should be {}"
+                             .format(mode, np.dtype, expected_mode))
+        mode = expected_mode
+
+    elif npimg.shape[2] == 4:
+        permitted_4_channel_modes = ['RGBA', 'CMYK']
+        if mode is not None and mode not in permitted_4_channel_modes:
+            raise ValueError("Only modes {} are supported for 4D inputs".format(permitted_4_channel_modes))
+
+        if mode is None and npimg.dtype == np.uint8:
+            mode = 'RGBA'
+    else:
+        permitted_3_channel_modes = ['RGB', 'YCbCr', 'HSV']
+        if mode is not None and mode not in permitted_3_channel_modes:
+            raise ValueError("Only modes {} are supported for 3D inputs".format(permitted_3_channel_modes))
+        if mode is None and npimg.dtype == np.uint8:
+            mode = 'RGB'
+
+    if mode is None:
+        raise TypeError('Input type {} is not supported'.format(npimg.dtype))
+
+    return Image.fromarray(npimg, mode=mode)
+
+
+def normalize(tensor, mean, std):
+    """Normalize a tensor image with mean and standard deviation.
+
+    See ``Normalize`` for more details.
+
+    Args:
+        tensor (Tensor): Tensor image of size (C, H, W) to be normalized.
+        mean (sequence): Sequence of means for each channel.
+        std (sequence): Sequence of standard deviations for each channely.
+
+    Returns:
+        Tensor: Normalized Tensor image.
+    """
+    if not _is_tensor_image(tensor):
+        raise TypeError('tensor is not a torch image.')
+    # TODO: make efficient
+    for t, m, s in zip(tensor, mean, std):
+        t.sub_(m).div_(s)
+    return tensor
+
+
+def resize(img, size, interpolation=Image.BILINEAR):
+    """Resize the input PIL Image to the given size.
+
+    Args:
+        img (PIL Image): Image to be resized.
+        size (sequence or int): Desired output size. If size is a sequence like
+            (h, w), the output size will be matched to this. If size is an int,
+            the smaller edge of the image will be matched to this number maintaing
+            the aspect ratio. i.e, if height > width, then image will be rescaled to
+            (size * height / width, size)
+        interpolation (int, optional): Desired interpolation. Default is
+            ``PIL.Image.BILINEAR``
+
+    Returns:
+        PIL Image: Resized image.
+    """
+    if not _is_pil_image(img):
+        raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
+    if not (isinstance(size, int) or (isinstance(size, collections.Iterable) and len(size) == 2)):
+        raise TypeError('Got inappropriate size arg: {}'.format(size))
+
+    if isinstance(size, int):
+        w, h = img.size
+        if (w <= h and w == size) or (h <= w and h == size):
+            return img
+        if w < h:
+            ow = size
+            oh = int(size * h / w)
+            return img.resize((ow, oh), interpolation)
+        else:
+            oh = size
+            ow = int(size * w / h)
+            return img.resize((ow, oh), interpolation)
+    else:
+        return img.resize(size[::-1], interpolation)
+
+
+def scale(*args, **kwargs):
+    warnings.warn("The use of the transforms.Scale transform is deprecated, " +
+                  "please use transforms.Resize instead.")
+    return resize(*args, **kwargs)
+
+
+def pad(img, padding, fill=0, padding_mode='constant'):
+    """Pad the given PIL Image on all sides with speficified padding mode and fill value.
+
+    Args:
+        img (PIL Image): Image to be padded.
+        padding (int or tuple): Padding on each border. If a single int is provided this
+            is used to pad all borders. If tuple of length 2 is provided this is the padding
+            on left/right and top/bottom respectively. If a tuple of length 4 is provided
+            this is the padding for the left, top, right and bottom borders
+            respectively.
+        fill: Pixel fill value for constant fill. Default is 0. If a tuple of
+            length 3, it is used to fill R, G, B channels respectively.
+            This value is only used when the padding_mode is constant
+        padding_mode: Type of padding. Should be: constant, edge, reflect or symmetric. Default is constant.
+            constant: pads with a constant value, this value is specified with fill
+            edge: pads with the last value on the edge of the image
+            reflect: pads with reflection of image (without repeating the last value on the edge)
+                padding [1, 2, 3, 4] with 2 elements on both sides in reflect mode
+                will result in [3, 2, 1, 2, 3, 4, 3, 2]
+            symmetric: pads with reflection of image (repeating the last value on the edge)
+                padding [1, 2, 3, 4] with 2 elements on both sides in symmetric mode
+                will result in [2, 1, 1, 2, 3, 4, 4, 3]
+
+    Returns:
+        PIL Image: Padded image.
+    """
+    if not _is_pil_image(img):
+        raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
+
+    if not isinstance(padding, (numbers.Number, tuple)):
+        raise TypeError('Got inappropriate padding arg')
+    if not isinstance(fill, (numbers.Number, str, tuple)):
+        raise TypeError('Got inappropriate fill arg')
+    if not isinstance(padding_mode, str):
+        raise TypeError('Got inappropriate padding_mode arg')
+
+    if isinstance(padding, collections.Sequence) and len(padding) not in [2, 4]:
+        raise ValueError("Padding must be an int or a 2, or 4 element tuple, not a " +
+                         "{} element tuple".format(len(padding)))
+
+    assert padding_mode in ['constant', 'edge', 'reflect', 'symmetric'], \
+        'Padding mode should be either constant, edge, reflect or symmetric'
+
+    if padding_mode == 'constant':
+        return ImageOps.expand(img, border=padding, fill=fill)
+    else:
+        if isinstance(padding, int):
+            pad_left = pad_right = pad_top = pad_bottom = padding
+        if isinstance(padding, collections.Sequence) and len(padding) == 2:
+            pad_left = pad_right = padding[0]
+            pad_top = pad_bottom = padding[1]
+        if isinstance(padding, collections.Sequence) and len(padding) == 4:
+            pad_left = padding[0]
+            pad_top = padding[1]
+            pad_right = padding[2]
+            pad_bottom = padding[3]
+
+        img = np.asarray(img)
+        # RGB image
+        if len(img.shape) == 3:
+            img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)), padding_mode)
+        # Grayscale image
+        if len(img.shape) == 2:
+            img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right)), padding_mode)
+
+        return Image.fromarray(img)
+
+
+def crop(img, i, j, h, w):
+    """Crop the given PIL Image.
+
+    Args:
+        img (PIL Image): Image to be cropped.
+        i: Upper pixel coordinate.
+        j: Left pixel coordinate.
+        h: Height of the cropped image.
+        w: Width of the cropped image.
+
+    Returns:
+        PIL Image: Cropped image.
+    """
+    if not _is_pil_image(img):
+        raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
+
+    return img.crop((j, i, j + w, i + h))
+
+
+def center_crop(img, output_size):
+    if isinstance(output_size, numbers.Number):
+        output_size = (int(output_size), int(output_size))
+    w, h = img.size
+    th, tw = output_size
+    i = int(round((h - th) / 2.))
+    j = int(round((w - tw) / 2.))
+    return crop(img, i, j, th, tw)
+
+
+def resized_crop(img, i, j, h, w, size, interpolation=Image.BILINEAR):
+    """Crop the given PIL Image and resize it to desired size.
+
+    Notably used in RandomResizedCrop.
+
+    Args:
+        img (PIL Image): Image to be cropped.
+        i: Upper pixel coordinate.
+        j: Left pixel coordinate.
+        h: Height of the cropped image.
+        w: Width of the cropped image.
+        size (sequence or int): Desired output size. Same semantics as ``scale``.
+        interpolation (int, optional): Desired interpolation. Default is
+            ``PIL.Image.BILINEAR``.
+    Returns:
+        PIL Image: Cropped image.
+    """
+    assert _is_pil_image(img), 'img should be PIL Image'
+    img = crop(img, i, j, h, w)
+    img = resize(img, size, interpolation)
+    return img
+
+
+def hflip(img):
+    """Horizontally flip the given PIL Image.
+
+    Args:
+        img (PIL Image): Image to be flipped.
+
+    Returns:
+        PIL Image:  Horizontall flipped image.
+    """
+    if not _is_pil_image(img):
+        raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
+
+    return img.transpose(Image.FLIP_LEFT_RIGHT)
+
+
+def vflip(img):
+    """Vertically flip the given PIL Image.
+
+    Args:
+        img (PIL Image): Image to be flipped.
+
+    Returns:
+        PIL Image:  Vertically flipped image.
+    """
+    if not _is_pil_image(img):
+        raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
+
+    return img.transpose(Image.FLIP_TOP_BOTTOM)
+
+
+def swap(img, crop):
+    def crop_image(image, cropnum):
+        width, high = image.size
+        crop_x = [int((width / cropnum[0]) * i) for i in range(cropnum[0] + 1)]
+        crop_y = [int((high / cropnum[1]) * i) for i in range(cropnum[1] + 1)]
+        im_list = []
+        for j in range(len(crop_y) - 1):
+            for i in range(len(crop_x) - 1):
+                im_list.append(image.crop((crop_x[i], crop_y[j], min(crop_x[i + 1], width), min(crop_y[j + 1], high))))
+        return im_list
+
+    widthcut, highcut = img.size
+    img = img.crop((10, 10, widthcut-10, highcut-10))
+    images = crop_image(img, crop)
+    pro = 5
+    if pro >= 5:          
+        tmpx = []
+        tmpy = []
+        count_x = 0
+        count_y = 0
+        k = 1
+        RAN = 2
+        for i in range(crop[1] * crop[0]):
+            tmpx.append(images[i])
+            count_x += 1
+            if len(tmpx) >= k:
+                tmp = tmpx[count_x - RAN:count_x]
+                random.shuffle(tmp)
+                tmpx[count_x - RAN:count_x] = tmp
+            if count_x == crop[0]:
+                tmpy.append(tmpx)
+                count_x = 0
+                count_y += 1
+                tmpx = []
+            if len(tmpy) >= k:
+                tmp2 = tmpy[count_y - RAN:count_y]
+                random.shuffle(tmp2)
+                tmpy[count_y - RAN:count_y] = tmp2
+        random_im = []
+        for line in tmpy:
+            random_im.extend(line)
+        
+        # random.shuffle(images)
+        width, high = img.size
+        iw = int(width / crop[0])
+        ih = int(high / crop[1])
+        toImage = Image.new('RGB', (iw * crop[0], ih * crop[1]))
+        x = 0
+        y = 0
+        for i in random_im:
+            i = i.resize((iw, ih), Image.ANTIALIAS)
+            toImage.paste(i, (x * iw, y * ih))
+            x += 1
+            if x == crop[0]:
+                x = 0
+                y += 1
+    else:
+        toImage = img
+    toImage = toImage.resize((widthcut, highcut))
+    return toImage
+
+
+
+def five_crop(img, size):
+    """Crop the given PIL Image into four corners and the central crop.
+
+    .. Note::
+        This transform returns a tuple of images and there may be a
+        mismatch in the number of inputs and targets your ``Dataset`` returns.
+
+    Args:
+       size (sequence or int): Desired output size of the crop. If size is an
+           int instead of sequence like (h, w), a square crop (size, size) is
+           made.
+    Returns:
+        tuple: tuple (tl, tr, bl, br, center) corresponding top left,
+            top right, bottom left, bottom right and center crop.
+    """
+    if isinstance(size, numbers.Number):
+        size = (int(size), int(size))
+    else:
+        assert len(size) == 2, "Please provide only two dimensions (h, w) for size."
+
+    w, h = img.size
+    crop_h, crop_w = size
+    if crop_w > w or crop_h > h:
+        raise ValueError("Requested crop size {} is bigger than input size {}".format(size,
+                                                                                      (h, w)))
+    tl = img.crop((0, 0, crop_w, crop_h))
+    tr = img.crop((w - crop_w, 0, w, crop_h))
+    bl = img.crop((0, h - crop_h, crop_w, h))
+    br = img.crop((w - crop_w, h - crop_h, w, h))
+    center = center_crop(img, (crop_h, crop_w))
+    return (tl, tr, bl, br, center)
+
+
+def ten_crop(img, size, vertical_flip=False):
+    """Crop the given PIL Image into four corners and the central crop plus the
+       flipped version of these (horizontal flipping is used by default).
+
+    .. Note::
+        This transform returns a tuple of images and there may be a
+        mismatch in the number of inputs and targets your ``Dataset`` returns.
+
+       Args:
+           size (sequence or int): Desired output size of the crop. If size is an
+               int instead of sequence like (h, w), a square crop (size, size) is
+               made.
+           vertical_flip (bool): Use vertical flipping instead of horizontal
+
+        Returns:
+            tuple: tuple (tl, tr, bl, br, center, tl_flip, tr_flip, bl_flip,
+                br_flip, center_flip) corresponding top left, top right,
+                bottom left, bottom right and center crop and same for the
+                flipped image.
+    """
+    if isinstance(size, numbers.Number):
+        size = (int(size), int(size))
+    else:
+        assert len(size) == 2, "Please provide only two dimensions (h, w) for size."
+
+    first_five = five_crop(img, size)
+
+    if vertical_flip:
+        img = vflip(img)
+    else:
+        img = hflip(img)
+
+    second_five = five_crop(img, size)
+    return first_five + second_five
+
+
+def adjust_brightness(img, brightness_factor):
+    """Adjust brightness of an Image.
+
+    Args:
+        img (PIL Image): PIL Image to be adjusted.
+        brightness_factor (float):  How much to adjust the brightness. Can be
+            any non negative number. 0 gives a black image, 1 gives the
+            original image while 2 increases the brightness by a factor of 2.
+
+    Returns:
+        PIL Image: Brightness adjusted image.
+    """
+    if not _is_pil_image(img):
+        raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
+
+    enhancer = ImageEnhance.Brightness(img)
+    img = enhancer.enhance(brightness_factor)
+    return img
+
+
+def adjust_contrast(img, contrast_factor):
+    """Adjust contrast of an Image.
+
+    Args:
+        img (PIL Image): PIL Image to be adjusted.
+        contrast_factor (float): How much to adjust the contrast. Can be any
+            non negative number. 0 gives a solid gray image, 1 gives the
+            original image while 2 increases the contrast by a factor of 2.
+
+    Returns:
+        PIL Image: Contrast adjusted image.
+    """
+    if not _is_pil_image(img):
+        raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
+
+    enhancer = ImageEnhance.Contrast(img)
+    img = enhancer.enhance(contrast_factor)
+    return img
+
+
+def adjust_saturation(img, saturation_factor):
+    """Adjust color saturation of an image.
+
+    Args:
+        img (PIL Image): PIL Image to be adjusted.
+        saturation_factor (float):  How much to adjust the saturation. 0 will
+            give a black and white image, 1 will give the original image while
+            2 will enhance the saturation by a factor of 2.
+
+    Returns:
+        PIL Image: Saturation adjusted image.
+    """
+    if not _is_pil_image(img):
+        raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
+
+    enhancer = ImageEnhance.Color(img)
+    img = enhancer.enhance(saturation_factor)
+    return img
+
+
+def adjust_hue(img, hue_factor):
+    """Adjust hue of an image.
+
+    The image hue is adjusted by converting the image to HSV and
+    cyclically shifting the intensities in the hue channel (H).
+    The image is then converted back to original image mode.
+
+    `hue_factor` is the amount of shift in H channel and must be in the
+    interval `[-0.5, 0.5]`.
+
+    See https://en.wikipedia.org/wiki/Hue for more details on Hue.
+
+    Args:
+        img (PIL Image): PIL Image to be adjusted.
+        hue_factor (float):  How much to shift the hue channel. Should be in
+            [-0.5, 0.5]. 0.5 and -0.5 give complete reversal of hue channel in
+            HSV space in positive and negative direction respectively.
+            0 means no shift. Therefore, both -0.5 and 0.5 will give an image
+            with complementary colors while 0 gives the original image.
+
+    Returns:
+        PIL Image: Hue adjusted image.
+    """
+    if not(-0.5 <= hue_factor <= 0.5):
+        raise ValueError('hue_factor is not in [-0.5, 0.5].'.format(hue_factor))
+
+    if not _is_pil_image(img):
+        raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
+
+    input_mode = img.mode
+    if input_mode in {'L', '1', 'I', 'F'}:
+        return img
+
+    h, s, v = img.convert('HSV').split()
+
+    np_h = np.array(h, dtype=np.uint8)
+    # uint8 addition take cares of rotation across boundaries
+    with np.errstate(over='ignore'):
+        np_h += np.uint8(hue_factor * 255)
+    h = Image.fromarray(np_h, 'L')
+
+    img = Image.merge('HSV', (h, s, v)).convert(input_mode)
+    return img
+
+
+def adjust_gamma(img, gamma, gain=1):
+    """Perform gamma correction on an image.
+
+    Also known as Power Law Transform. Intensities in RGB mode are adjusted
+    based on the following equation:
+
+        I_out = 255 * gain * ((I_in / 255) ** gamma)
+
+    See https://en.wikipedia.org/wiki/Gamma_correction for more details.
+
+    Args:
+        img (PIL Image): PIL Image to be adjusted.
+        gamma (float): Non negative real number. gamma larger than 1 make the
+            shadows darker, while gamma smaller than 1 make dark regions
+            lighter.
+        gain (float): The constant multiplier.
+    """
+    if not _is_pil_image(img):
+        raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
+
+    if gamma < 0:
+        raise ValueError('Gamma should be a non-negative real number')
+
+    input_mode = img.mode
+    img = img.convert('RGB')
+
+    gamma_map = [255 * gain * pow(ele / 255., gamma) for ele in range(256)] * 3
+    img = img.point(gamma_map)  # use PIL's point-function to accelerate this part
+
+    img = img.convert(input_mode)
+    return img
+
+
+def rotate(img, angle, resample=False, expand=False, center=None):
+    """Rotate the image by angle.
+
+
+    Args:
+        img (PIL Image): PIL Image to be rotated.
+        angle ({float, int}): In degrees degrees counter clockwise order.
+        resample ({PIL.Image.NEAREST, PIL.Image.BILINEAR, PIL.Image.BICUBIC}, optional):
+            An optional resampling filter.
+            See http://pillow.readthedocs.io/en/3.4.x/handbook/concepts.html#filters
+            If omitted, or if the image has mode "1" or "P", it is set to PIL.Image.NEAREST.
+        expand (bool, optional): Optional expansion flag.
+            If true, expands the output image to make it large enough to hold the entire rotated image.
+            If false or omitted, make the output image the same size as the input image.
+            Note that the expand flag assumes rotation around the center and no translation.
+        center (2-tuple, optional): Optional center of rotation.
+            Origin is the upper left corner.
+            Default is the center of the image.
+    """
+
+    if not _is_pil_image(img):
+        raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
+
+    return img.rotate(angle, resample, expand, center)
+
+
+def _get_inverse_affine_matrix(center, angle, translate, scale, shear):
+    # Helper method to compute inverse matrix for affine transformation
+
+    # As it is explained in PIL.Image.rotate
+    # We need compute INVERSE of affine transformation matrix: M = T * C * RSS * C^-1
+    # where T is translation matrix: [1, 0, tx | 0, 1, ty | 0, 0, 1]
+    #       C is translation matrix to keep center: [1, 0, cx | 0, 1, cy | 0, 0, 1]
+    #       RSS is rotation with scale and shear matrix
+    #       RSS(a, scale, shear) = [ cos(a)*scale    -sin(a + shear)*scale     0]
+    #                              [ sin(a)*scale    cos(a + shear)*scale     0]
+    #                              [     0                  0          1]
+    # Thus, the inverse is M^-1 = C * RSS^-1 * C^-1 * T^-1
+
+    angle = math.radians(angle)
+    shear = math.radians(shear)
+    scale = 1.0 / scale
+
+    # Inverted rotation matrix with scale and shear
+    d = math.cos(angle + shear) * math.cos(angle) + math.sin(angle + shear) * math.sin(angle)
+    matrix = [
+        math.cos(angle + shear), math.sin(angle + shear), 0,
+        -math.sin(angle), math.cos(angle), 0
+    ]
+    matrix = [scale / d * m for m in matrix]
+
+    # Apply inverse of translation and of center translation: RSS^-1 * C^-1 * T^-1
+    matrix[2] += matrix[0] * (-center[0] - translate[0]) + matrix[1] * (-center[1] - translate[1])
+    matrix[5] += matrix[3] * (-center[0] - translate[0]) + matrix[4] * (-center[1] - translate[1])
+
+    # Apply center translation: C * RSS^-1 * C^-1 * T^-1
+    matrix[2] += center[0]
+    matrix[5] += center[1]
+    return matrix
+
+
+def affine(img, angle, translate, scale, shear, resample=0, fillcolor=None):
+    """Apply affine transformation on the image keeping image center invariant
+
+    Args:
+        img (PIL Image): PIL Image to be rotated.
+        angle ({float, int}): rotation angle in degrees between -180 and 180, clockwise direction.
+        translate (list or tuple of integers): horizontal and vertical translations (post-rotation translation)
+        scale (float): overall scale
+        shear (float): shear angle value in degrees between -180 to 180, clockwise direction.
+        resample ({PIL.Image.NEAREST, PIL.Image.BILINEAR, PIL.Image.BICUBIC}, optional):
+            An optional resampling filter.
+            See http://pillow.readthedocs.io/en/3.4.x/handbook/concepts.html#filters
+            If omitted, or if the image has mode "1" or "P", it is set to PIL.Image.NEAREST.
+        fillcolor (int): Optional fill color for the area outside the transform in the output image. (Pillow>=5.0.0)
+    """
+    if not _is_pil_image(img):
+        raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
+
+    assert isinstance(translate, (tuple, list)) and len(translate) == 2, \
+        "Argument translate should be a list or tuple of length 2"
+
+    assert scale > 0.0, "Argument scale should be positive"
+
+    output_size = img.size
+    center = (img.size[0] * 0.5 + 0.5, img.size[1] * 0.5 + 0.5)
+    matrix = _get_inverse_affine_matrix(center, angle, translate, scale, shear)
+    kwargs = {"fillcolor": fillcolor} if PILLOW_VERSION[0] == '5' else {}
+    return img.transform(output_size, Image.AFFINE, matrix, resample, **kwargs)
+
+
+def to_grayscale(img, num_output_channels=1):
+    """Convert image to grayscale version of image.
+
+    Args:
+        img (PIL Image): Image to be converted to grayscale.
+
+    Returns:
+        PIL Image:  Grayscale version of the image.
+                    if num_output_channels == 1 : returned image is single channel
+                    if num_output_channels == 3 : returned image is 3 channel with r == g == b
+    """
+    if not _is_pil_image(img):
+        raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
+
+    if num_output_channels == 1:
+        img = img.convert('L')
+    elif num_output_channels == 3:
+        img = img.convert('L')
+        np_img = np.array(img, dtype=np.uint8)
+        np_img = np.dstack([np_img, np_img, np_img])
+        img = Image.fromarray(np_img, 'RGB')
+    else:
+        raise ValueError('num_output_channels should be either 1 or 3')
+
+    return img

Dcl/util/Asoftmax_loss.py

@@ -0,0 +1,47 @@
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+import torch.nn.functional as F
+from torch.nn import Parameter
+import math
+
+import pdb
+
+class AngleLoss(nn.Module):
+    def __init__(self, gamma=0):
+        super(AngleLoss, self).__init__()
+        self.gamma   = gamma
+        self.it = 0
+        self.LambdaMin = 50.0
+        self.LambdaMax = 1500.0
+        self.lamb = 1500.0
+
+    def forward(self, input, target, decay=None):
+        self.it += 1
+        cos_theta,phi_theta = input
+        target = target.view(-1,1) #size=(B,1)
+
+        index = cos_theta.data * 0.0 #size=(B,Classnum)
+        index.scatter_(1,target.data.view(-1,1),1)
+        index = index.byte()
+        index = Variable(index)
+
+        if decay is None:
+            self.lamb = max(self.LambdaMin,self.LambdaMax/(1+0.1*self.it ))
+        else:
+            self.LambdaMax *= decay 
+            self.lamb = max(self.LambdaMin, self.LambdaMax) 
+        output = cos_theta * 1.0 #size=(B,Classnum)
+        output[index] -= cos_theta[index]*(1.0+0)/(1+self.lamb)
+        output[index] += phi_theta[index]*(1.0+0)/(1+self.lamb)
+
+        logpt = F.log_softmax(output, 1)
+        logpt = logpt.gather(1,target)
+        logpt = logpt.view(-1)
+        pt = Variable(logpt.data.exp())
+
+        loss = -1 * (1-pt)**self.gamma * logpt
+        loss = loss.mean()
+
+        return loss
+

Dcl/util/autoaugment.py

@@ -0,0 +1,232 @@
+from PIL import Image, ImageEnhance, ImageOps
+import numpy as np
+import random
+
+
+class ImageNetPolicy(object):
+    """ Randomly choose one of the best 24 Sub-policies on ImageNet.
+        Example:
+        >>> policy = ImageNetPolicy()
+        >>> transformed = policy(image)
+        Example as a PyTorch Transform:
+        >>> transform=transforms.Compose([
+        >>>     transforms.Resize(256),
+        >>>     ImageNetPolicy(),
+        >>>     transforms.ToTensor()])
+    """
+    def __init__(self, fillcolor=(128, 128, 128)):
+        self.policies = [
+            SubPolicy(0.4, "posterize", 8, 0.6, "rotate", 9, fillcolor),
+            SubPolicy(0.6, "solarize", 5, 0.6, "autocontrast", 5, fillcolor),
+            SubPolicy(0.8, "equalize", 8, 0.6, "equalize", 3, fillcolor),
+            SubPolicy(0.6, "posterize", 7, 0.6, "posterize", 6, fillcolor),
+            SubPolicy(0.4, "equalize", 7, 0.2, "solarize", 4, fillcolor),
+
+            SubPolicy(0.4, "equalize", 4, 0.8, "rotate", 8, fillcolor),
+            SubPolicy(0.6, "solarize", 3, 0.6, "equalize", 7, fillcolor),
+            SubPolicy(0.8, "posterize", 5, 1.0, "equalize", 2, fillcolor),
+            SubPolicy(0.2, "rotate", 3, 0.6, "solarize", 8, fillcolor),
+            SubPolicy(0.6, "equalize", 8, 0.4, "posterize", 6, fillcolor),
+
+            SubPolicy(0.8, "rotate", 8, 0.4, "color", 0, fillcolor),
+            SubPolicy(0.4, "rotate", 9, 0.6, "equalize", 2, fillcolor),
+            SubPolicy(0.0, "equalize", 7, 0.8, "equalize", 8, fillcolor),
+            SubPolicy(0.6, "invert", 4, 1.0, "equalize", 8, fillcolor),
+            SubPolicy(0.6, "color", 4, 1.0, "contrast", 8, fillcolor),
+
+            SubPolicy(0.8, "rotate", 8, 1.0, "color", 2, fillcolor),
+            SubPolicy(0.8, "color", 8, 0.8, "solarize", 7, fillcolor),
+            SubPolicy(0.4, "sharpness", 7, 0.6, "invert", 8, fillcolor),
+            SubPolicy(0.6, "shearX", 5, 1.0, "equalize", 9, fillcolor),
+            SubPolicy(0.4, "color", 0, 0.6, "equalize", 3, fillcolor),
+
+            SubPolicy(0.4, "equalize", 7, 0.2, "solarize", 4, fillcolor),
+            SubPolicy(0.6, "solarize", 5, 0.6, "autocontrast", 5, fillcolor),
+            SubPolicy(0.6, "invert", 4, 1.0, "equalize", 8, fillcolor),
+            SubPolicy(0.6, "color", 4, 1.0, "contrast", 8, fillcolor),
+            SubPolicy(0.8, "equalize", 8, 0.6, "equalize", 3, fillcolor)
+        ]
+
+
+    def __call__(self, img):
+        policy_idx = random.randint(0, len(self.policies) - 1)
+        return self.policies[policy_idx](img)
+
+    def __repr__(self):
+        return "AutoAugment ImageNet Policy"
+
+
+class CIFAR10Policy(object):
+    """ Randomly choose one of the best 25 Sub-policies on CIFAR10.
+        Example:
+        >>> policy = CIFAR10Policy()
+        >>> transformed = policy(image)
+        Example as a PyTorch Transform:
+        >>> transform=transforms.Compose([
+        >>>     transforms.Resize(256),
+        >>>     CIFAR10Policy(),
+        >>>     transforms.ToTensor()])
+    """
+    def __init__(self, fillcolor=(128, 128, 128)):
+        self.policies = [
+            SubPolicy(0.1, "invert", 7, 0.2, "contrast", 6, fillcolor),
+            SubPolicy(0.7, "rotate", 2, 0.3, "translateX", 9, fillcolor),
+            SubPolicy(0.8, "sharpness", 1, 0.9, "sharpness", 3, fillcolor),
+            SubPolicy(0.5, "shearY", 8, 0.7, "translateY", 9, fillcolor),
+            SubPolicy(0.5, "autocontrast", 8, 0.9, "equalize", 2, fillcolor),
+
+            SubPolicy(0.2, "shearY", 7, 0.3, "posterize", 7, fillcolor),
+            SubPolicy(0.4, "color", 3, 0.6, "brightness", 7, fillcolor),
+            SubPolicy(0.3, "sharpness", 9, 0.7, "brightness", 9, fillcolor),
+            SubPolicy(0.6, "equalize", 5, 0.5, "equalize", 1, fillcolor),
+            SubPolicy(0.6, "contrast", 7, 0.6, "sharpness", 5, fillcolor),
+
+            SubPolicy(0.7, "color", 7, 0.5, "translateX", 8, fillcolor),
+            SubPolicy(0.3, "equalize", 7, 0.4, "autocontrast", 8, fillcolor),
+            SubPolicy(0.4, "translateY", 3, 0.2, "sharpness", 6, fillcolor),
+            SubPolicy(0.9, "brightness", 6, 0.2, "color", 8, fillcolor),
+            SubPolicy(0.5, "solarize", 2, 0.0, "invert", 3, fillcolor),
+
+            SubPolicy(0.2, "equalize", 0, 0.6, "autocontrast", 0, fillcolor),
+            SubPolicy(0.2, "equalize", 8, 0.8, "equalize", 4, fillcolor),
+            SubPolicy(0.9, "color", 9, 0.6, "equalize", 6, fillcolor),
+            SubPolicy(0.8, "autocontrast", 4, 0.2, "solarize", 8, fillcolor),
+            SubPolicy(0.1, "brightness", 3, 0.7, "color", 0, fillcolor),
+
+            SubPolicy(0.4, "solarize", 5, 0.9, "autocontrast", 3, fillcolor),
+            SubPolicy(0.9, "translateY", 9, 0.7, "translateY", 9, fillcolor),
+            SubPolicy(0.9, "autocontrast", 2, 0.8, "solarize", 3, fillcolor),
+            SubPolicy(0.8, "equalize", 8, 0.1, "invert", 3, fillcolor),
+            SubPolicy(0.7, "translateY", 9, 0.9, "autocontrast", 1, fillcolor)
+        ]
+
+
+    def __call__(self, img):
+        policy_idx = random.randint(0, len(self.policies) - 1)
+        return self.policies[policy_idx](img)
+
+    def __repr__(self):
+        return "AutoAugment CIFAR10 Policy"
+
+
+class SVHNPolicy(object):
+    """ Randomly choose one of the best 25 Sub-policies on SVHN.
+        Example:
+        >>> policy = SVHNPolicy()
+        >>> transformed = policy(image)
+        Example as a PyTorch Transform:
+        >>> transform=transforms.Compose([
+        >>>     transforms.Resize(256),
+        >>>     SVHNPolicy(),
+        >>>     transforms.ToTensor()])
+    """
+    def __init__(self, fillcolor=(128, 128, 128)):
+        self.policies = [
+            SubPolicy(0.9, "shearX", 4, 0.2, "invert", 3, fillcolor),
+            SubPolicy(0.9, "shearY", 8, 0.7, "invert", 5, fillcolor),
+            SubPolicy(0.6, "equalize", 5, 0.6, "solarize", 6, fillcolor),
+            SubPolicy(0.9, "invert", 3, 0.6, "equalize", 3, fillcolor),
+            SubPolicy(0.6, "equalize", 1, 0.9, "rotate", 3, fillcolor),
+
+            SubPolicy(0.9, "shearX", 4, 0.8, "autocontrast", 3, fillcolor),
+            SubPolicy(0.9, "shearY", 8, 0.4, "invert", 5, fillcolor),
+            SubPolicy(0.9, "shearY", 5, 0.2, "solarize", 6, fillcolor),
+            SubPolicy(0.9, "invert", 6, 0.8, "autocontrast", 1, fillcolor),
+            SubPolicy(0.6, "equalize", 3, 0.9, "rotate", 3, fillcolor),
+
+            SubPolicy(0.9, "shearX", 4, 0.3, "solarize", 3, fillcolor),
+            SubPolicy(0.8, "shearY", 8, 0.7, "invert", 4, fillcolor),
+            SubPolicy(0.9, "equalize", 5, 0.6, "translateY", 6, fillcolor),
+            SubPolicy(0.9, "invert", 4, 0.6, "equalize", 7, fillcolor),
+            SubPolicy(0.3, "contrast", 3, 0.8, "rotate", 4, fillcolor),
+
+            SubPolicy(0.8, "invert", 5, 0.0, "translateY", 2, fillcolor),
+            SubPolicy(0.7, "shearY", 6, 0.4, "solarize", 8, fillcolor),
+            SubPolicy(0.6, "invert", 4, 0.8, "rotate", 4, fillcolor),
+            SubPolicy(0.3, "shearY", 7, 0.9, "translateX", 3, fillcolor),
+            SubPolicy(0.1, "shearX", 6, 0.6, "invert", 5, fillcolor),
+
+            SubPolicy(0.7, "solarize", 2, 0.6, "translateY", 7, fillcolor),
+            SubPolicy(0.8, "shearY", 4, 0.8, "invert", 8, fillcolor),
+            SubPolicy(0.7, "shearX", 9, 0.8, "translateY", 3, fillcolor),
+            SubPolicy(0.8, "shearY", 5, 0.7, "autocontrast", 3, fillcolor),
+            SubPolicy(0.7, "shearX", 2, 0.1, "invert", 5, fillcolor)
+        ]
+
+
+    def __call__(self, img):
+        policy_idx = random.randint(0, len(self.policies) - 1)
+        return self.policies[policy_idx](img)
+
+    def __repr__(self):
+        return "AutoAugment SVHN Policy"
+
+
+class SubPolicy(object):
+    def __init__(self, p1, operation1, magnitude_idx1, p2, operation2, magnitude_idx2, fillcolor=(128, 128, 128)):
+        ranges = {
+            "shearX": np.linspace(0, 0.3, 10),
+            "shearY": np.linspace(0, 0.3, 10),
+            "translateX": np.linspace(0, 150 / 331, 10),
+            "translateY": np.linspace(0, 150 / 331, 10),
+            "rotate": np.linspace(0, 30, 10),
+            "color": np.linspace(0.0, 0.9, 10),
+            "posterize": np.round(np.linspace(8, 4, 10), 0).astype(np.int),
+            "solarize": np.linspace(256, 0, 10),
+            "contrast": np.linspace(0.0, 0.9, 10),
+            "sharpness": np.linspace(0.0, 0.9, 10),
+            "brightness": np.linspace(0.0, 0.9, 10),
+            "autocontrast": [0] * 10,
+            "equalize": [0] * 10,
+            "invert": [0] * 10
+        }
+
+        # from https://stackoverflow.com/questions/5252170/specify-image-filling-color-when-rotating-in-python-with-pil-and-setting-expand
+        def rotate_with_fill(img, magnitude):
+            rot = img.convert("RGBA").rotate(magnitude)
+            return Image.composite(rot, Image.new("RGBA", rot.size, (128,) * 4), rot).convert(img.mode)
+
+        func = {
+            "shearX": lambda img, magnitude: img.transform(
+                img.size, Image.AFFINE, (1, magnitude * random.choice([-1, 1]), 0, 0, 1, 0),
+                Image.BICUBIC, fillcolor=fillcolor),
+            "shearY": lambda img, magnitude: img.transform(
+                img.size, Image.AFFINE, (1, 0, 0, magnitude * random.choice([-1, 1]), 1, 0),
+                Image.BICUBIC, fillcolor=fillcolor),
+            "translateX": lambda img, magnitude: img.transform(
+                img.size, Image.AFFINE, (1, 0, magnitude * img.size[0] * random.choice([-1, 1]), 0, 1, 0),
+                fillcolor=fillcolor),
+            "translateY": lambda img, magnitude: img.transform(
+                img.size, Image.AFFINE, (1, 0, 0, 0, 1, magnitude * img.size[1] * random.choice([-1, 1])),
+                fillcolor=fillcolor),
+            "rotate": lambda img, magnitude: rotate_with_fill(img, magnitude),
+            # "rotate": lambda img, magnitude: img.rotate(magnitude * random.choice([-1, 1])),
+            "color": lambda img, magnitude: ImageEnhance.Color(img).enhance(1 + magnitude * random.choice([-1, 1])),
+            "posterize": lambda img, magnitude: ImageOps.posterize(img, magnitude),
+            "solarize": lambda img, magnitude: ImageOps.solarize(img, magnitude),
+            "contrast": lambda img, magnitude: ImageEnhance.Contrast(img).enhance(
+                1 + magnitude * random.choice([-1, 1])),
+            "sharpness": lambda img, magnitude: ImageEnhance.Sharpness(img).enhance(
+                1 + magnitude * random.choice([-1, 1])),
+            "brightness": lambda img, magnitude: ImageEnhance.Brightness(img).enhance(
+                1 + magnitude * random.choice([-1, 1])),
+            "autocontrast": lambda img, magnitude: ImageOps.autocontrast(img),
+            "equalize": lambda img, magnitude: ImageOps.equalize(img),
+            "invert": lambda img, magnitude: ImageOps.invert(img)
+        }
+
+        # self.name = "{}_{:.2f}_and_{}_{:.2f}".format(
+        #     operation1, ranges[operation1][magnitude_idx1],
+        #     operation2, ranges[operation2][magnitude_idx2])
+        self.p1 = p1
+        self.operation1 = func[operation1]
+        self.magnitude1 = ranges[operation1][magnitude_idx1]
+        self.p2 = p2
+        self.operation2 = func[operation2]
+        self.magnitude2 = ranges[operation2][magnitude_idx2]
+
+
+    def __call__(self, img):
+        if random.random() < self.p1: img = self.operation1(img, self.magnitude1)
+        if random.random() < self.p2: img = self.operation2(img, self.magnitude2)
+        return img

Dcl/util/dataset_DCL.py

@@ -0,0 +1,183 @@
+# coding=utf8
+from __future__ import division
+import os
+import torch
+import torch.utils.data as data
+import pandas
+import random
+import PIL.Image as Image
+from PIL import ImageStat
+
+import pdb
+
+def random_sample(img_names, labels):
+    anno_dict = {}
+    img_list = []
+    anno_list = []
+    for img, anno in zip(img_names, labels):
+        if not anno in anno_dict:
+            anno_dict[anno] = [img]
+        else:
+            anno_dict[anno].append(img)
+
+    for anno in anno_dict.keys():
+        anno_len = len(anno_dict[anno])
+        fetch_keys = random.sample(list(range(anno_len)), anno_len//10)
+        img_list.extend([anno_dict[anno][x] for x in fetch_keys])
+        anno_list.extend([anno for x in fetch_keys])
+    return img_list, anno_list
+
+
+
+class dataset(data.Dataset):
+    def __init__(self, Config, anno, swap_size=[7,7], common_aug=None, swap=None, totensor=None, train=False, train_val=False, test=False):
+        self.root_path = Config.rawdata_root
+        self.numcls = Config.numcls
+        self.dataset = Config.dataset
+        self.use_cls_2 = Config.cls_2
+        self.use_cls_mul = Config.cls_2xmul
+        if isinstance(anno, pandas.core.frame.DataFrame): # 判断类型
+            self.paths = anno['ImageName'].tolist()
+            self.labels = anno['label'].tolist()
+        elif isinstance(anno, dict):
+            self.paths = anno['img_name']
+            self.labels = anno['label']
+
+        if train_val:
+            self.paths, self.labels = random_sample(self.paths, self.labels)
+        self.common_aug = common_aug
+        self.swap = swap
+        self.totensor = totensor
+        self.cfg = Config
+        self.train = train
+        self.swap_size = swap_size
+        self.test = test
+
+    def __len__(self):
+        return len(self.paths)
+
+    def __getitem__(self, item):
+        img_path = os.path.join(self.root_path, self.paths[item])
+        img = self.pil_loader(img_path)
+        if self.test:
+            img = self.totensor(img)
+            label = self.labels[item]
+            return img, label, self.paths[item]
+        img_unswap = self.common_aug(img) if not self.common_aug is None else img # 原图数据
+
+        image_unswap_list = self.crop_image(img_unswap, self.swap_size)
+
+        swap_range = self.swap_size[0] * self.swap_size[1]
+        swap_law1 = [(i-(swap_range//2))/swap_range for i in range(swap_range)]
+        # 在训练的时候 self.train是true
+        if self.train:
+            img_swap = self.swap(img_unswap) # RCM机制处理后的照片
+            image_swap_list = self.crop_image(img_swap, self.swap_size)
+            unswap_stats = [sum(ImageStat.Stat(im).mean) for im in image_unswap_list]
+            swap_stats = [sum(ImageStat.Stat(im).mean) for im in image_swap_list]
+            swap_law2 = []
+            for swap_im in swap_stats:
+                distance = [abs(swap_im - unswap_im) for unswap_im in unswap_stats]
+                index = distance.index(min(distance))
+                swap_law2.append((index-(swap_range//2))/swap_range)
+            img_swap = self.totensor(img_swap)
+            label = self.labels[item]
+            if self.use_cls_mul:
+                label_swap = label + self.numcls
+            if self.use_cls_2:
+                label_swap = -1
+            img_unswap = self.totensor(img_unswap)
+            return img_unswap, img_swap, label, label_swap, swap_law1, swap_law2, self.paths[item]
+        else:
+            label = self.labels[item]
+            swap_law2 = [(i-(swap_range//2))/swap_range for i in range(swap_range)]
+            label_swap = label
+            img_unswap = self.totensor(img_unswap)
+            return img_unswap, label, label_swap, swap_law1, swap_law2, self.paths[item]
+#inputs, labels, labels_swap, swap_law, img_names
+    def pil_loader(self,imgpath):
+        with open(imgpath, 'rb') as f:
+            with Image.open(f) as img:
+                return img.convert('RGB')
+
+    def crop_image(self, image, cropnum):
+        width, high = image.size
+        crop_x = [int((width / cropnum[0]) * i) for i in range(cropnum[0] + 1)]
+        crop_y = [int((high / cropnum[1]) * i) for i in range(cropnum[1] + 1)]
+        im_list = []
+        for j in range(len(crop_y) - 1):
+            for i in range(len(crop_x) - 1):
+                im_list.append(image.crop((crop_x[i], crop_y[j], min(crop_x[i + 1], width), min(crop_y[j + 1], high))))
+        return im_list
+
+
+    def get_weighted_sampler(self):
+        img_nums = len(self.labels)
+        weights = [self.labels.count(x) for x in range(self.numcls)]
+        return torch.utils.data.sampler.WeightedRandomSampler(weights, num_samples=img_nums)
+
+
+def collate_fn4train(batch):
+    imgs = []
+    label = []
+    label_swap = []
+    law_swap = []
+    img_name = []
+    for sample in batch:
+        imgs.append(sample[0]) # 原始图片
+        imgs.append(sample[1]) # RCM混淆机制打乱后的图片
+        label.append(sample[2])# 类别标签
+        label.append(sample[2])
+        # 给每张图片贴上标签以区分是否被破坏过
+        if sample[3] == -1:
+            label_swap.append(1)
+            label_swap.append(0)
+        else:
+            label_swap.append(sample[2])
+            label_swap.append(sample[3])
+        law_swap.append(sample[4]) # 未打乱时每个小块的位置
+        law_swap.append(sample[5]) # 打乱后每个小块的的位置
+        img_name.append(sample[-1])# 图片名字
+
+    return torch.stack(imgs, 0), label, label_swap, law_swap, img_name # 张量0维度扩展
+
+def collate_fn4val(batch):
+    imgs = []
+    label = []
+    label_swap = []
+    law_swap = []
+    img_name = []
+    for sample in batch:
+        imgs.append(sample[0])
+        label.append(sample[1])
+        if sample[3] == -1:
+            label_swap.append(1)
+        else:
+            label_swap.append(sample[2])
+        law_swap.append(sample[3])
+        img_name.append(sample[-1])
+    return torch.stack(imgs, 0), label, label_swap, law_swap, img_name
+
+def collate_fn4backbone(batch):
+    imgs = []
+    label = []
+    img_name = []
+    for sample in batch:
+        imgs.append(sample[0])
+        if len(sample) == 7:
+            label.append(sample[2])
+        else:
+            label.append(sample[1])
+        img_name.append(sample[-1])
+    return torch.stack(imgs, 0), label, img_name
+
+
+def collate_fn4test(batch):
+    imgs = []
+    label = []
+    img_name = []
+    for sample in batch:
+        imgs.append(sample[0])
+        label.append(sample[1])
+        img_name.append(sample[-1])
+    return torch.stack(imgs, 0), label, img_name

Dcl/util/eval_model.py

@@ -0,0 +1,81 @@
+#coding=utf8
+from __future__ import print_function, division
+import os,time,datetime
+import numpy as np
+import datetime
+from math import ceil
+
+import torch
+from torch import nn
+from torch.autograd import Variable
+import torch.nn.functional as F
+
+from util.utils import LossRecord
+
+import pdb
+
+def dt():
+    return datetime.datetime.now().strftime("%Y-%m-%d-%H_%M_%S")
+
+def eval_turn(Config, model, data_loader, val_version, epoch_num, log_file):
+
+    model.train(False)
+
+    val_corrects1 = 0
+    val_corrects2 = 0
+    val_corrects3 = 0
+    val_size = data_loader.__len__()
+    item_count = data_loader.total_item_len
+    t0 = time.time()
+    get_l1_loss = nn.L1Loss()
+    get_ce_loss = nn.CrossEntropyLoss()
+
+    val_batch_size = data_loader.batch_size
+    val_epoch_step = data_loader.__len__()
+    num_cls = data_loader.num_cls
+
+    val_loss_recorder = LossRecord(val_batch_size)
+    val_celoss_recorder = LossRecord(val_batch_size)
+    print('evaluating %s ...'%val_version, flush=True)
+    with torch.no_grad():
+        for batch_cnt_val, data_val in enumerate(data_loader):
+            inputs = Variable(data_val[0].cuda())
+            labels = Variable(torch.from_numpy(np.array(data_val[1])).long().cuda())
+            outputs = model(inputs)
+            loss = 0
+            # 前向预测的时候只做分类
+            ce_loss = get_ce_loss(outputs[0], labels).item() # 获取tensor的值不包括类型
+            loss += ce_loss
+
+            val_loss_recorder.update(loss)
+            val_celoss_recorder.update(ce_loss)
+
+            if Config.use_dcl and Config.cls_2xmul:
+                outputs_pred = outputs[0] + outputs[1][:,0:num_cls] + outputs[1][:,num_cls:2*num_cls]
+            else:
+                outputs_pred = outputs[0]
+            top3_val, top3_pos = torch.topk(outputs_pred, 3) # 获取某一维度上 的前k个最大值（张量 k的值）
+
+            print('{:s} eval_batch: {:-6d} / {:d} loss: {:8.4f}'.format(val_version, batch_cnt_val, val_epoch_step, loss), flush=True)
+
+            batch_corrects1 = torch.sum((top3_pos[:, 0] == labels)).data.item() # .data返回的是tensor
+            val_corrects1 += batch_corrects1
+            batch_corrects2 = torch.sum((top3_pos[:, 1] == labels)).data.item()
+            val_corrects2 += (batch_corrects2 + batch_corrects1)
+            batch_corrects3 = torch.sum((top3_pos[:, 2] == labels)).data.item()
+            val_corrects3 += (batch_corrects3 + batch_corrects2 + batch_corrects1)
+
+        val_acc1 = val_corrects1 / item_count
+        val_acc2 = val_corrects2 / item_count
+        val_acc3 = val_corrects3 / item_count
+
+        log_file.write(val_version  + '\t' +str(val_loss_recorder.get_val())+'\t' + str(val_celoss_recorder.get_val()) + '\t' + str(val_acc1) + '\t' + str(val_acc3) + '\n')
+
+        t1 = time.time()
+        since = t1-t0
+        print('--'*30, flush=True)
+        print('% 3d %s %s %s-loss: %.4f ||%s-acc@1: %.4f %s-acc@2: %.4f %s-acc@3: %.4f ||time: %d' % (epoch_num, val_version, dt(), val_version, val_loss_recorder.get_val(init=True), val_version, val_acc1,val_version, val_acc2, val_version, val_acc3, since), flush=True)
+        print('--' * 30, flush=True)
+
+    return val_acc1, val_acc2, val_acc3
+

Dcl/util/test_tool.py

@@ -0,0 +1,99 @@
+import os
+import math
+import numpy as np
+import cv2
+import datetime
+
+import torch
+from torchvision.utils import save_image, make_grid
+
+import pdb
+
+def dt():
+    return datetime.datetime.now().strftime("%Y-%m-%d-%H_%M_%S")
+
+def set_text(text, img):
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    if isinstance(text, str):
+        cont = text
+        cv2.putText(img, cont, (20, 50), font, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
+    if isinstance(text, float):
+        cont = '%.4f'%text
+        cv2.putText(img, cont, (20, 50), font, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
+    if isinstance(text, list):
+        for count in range(len(img)):
+            cv2.putText(img[count], text[count], (20, 50), font, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
+    return img
+
+def save_multi_img(img_list, text_list, grid_size=[5,5], sub_size=200, save_dir='./', save_name=None):
+    if len(img_list) > grid_size[0]*grid_size[1]:
+        merge_height = math.ceil(len(img_list) / grid_size[0]) * sub_size
+    else:
+        merge_height = grid_size[1]*sub_size
+    merged_img = np.zeros((merge_height, grid_size[0]*sub_size, 3))
+
+    if isinstance(img_list[0], str):
+        img_name_list = img_list
+        img_list = []
+        for img_name in img_name_list:
+            img_list.append(cv2.imread(img_name))
+
+    img_counter = 0
+    for img, txt in zip(img_list, text_list):
+        img = cv2.resize(img, (sub_size, sub_size))
+        img = set_text(txt, img)
+        pos = [img_counter // grid_size[1], img_counter % grid_size[1]]
+        sub_pos = [pos[0]*sub_size, (pos[0]+1)*sub_size,
+                   pos[1]*sub_size, (pos[1]+1)*sub_size]
+        merged_img[sub_pos[0]:sub_pos[1], sub_pos[2]:sub_pos[3], :] = img
+        img_counter += 1
+
+    if save_name is None:
+        img_save_path = os.path.join(save_dir, dt()+'.png')
+    else:
+        img_save_path = os.path.join(save_dir, save_name+'.png')
+    cv2.imwrite(img_save_path,  merged_img)
+    print('saved img in %s ...'%img_save_path)
+
+
+def cls_base_acc(result_gather):
+    top1_acc = {}
+    top3_acc = {}
+    cls_count = {}
+    for img_item in result_gather.keys():
+        acc_case = result_gather[img_item]
+
+        if acc_case['label'] in cls_count:
+            cls_count[acc_case['label']] += 1
+            if acc_case['top1_cat'] == acc_case['label']:
+                top1_acc[acc_case['label']] += 1
+            if acc_case['label'] in [acc_case['top1_cat'], acc_case['top2_cat'], acc_case['top3_cat']]:
+                top3_acc[acc_case['label']] += 1
+        else:
+            cls_count[acc_case['label']] = 1
+            if acc_case['top1_cat'] == acc_case['label']:
+                top1_acc[acc_case['label']] = 1
+            else:
+                top1_acc[acc_case['label']] = 0
+
+            if acc_case['label'] in [acc_case['top1_cat'], acc_case['top2_cat'], acc_case['top3_cat']]:
+                top3_acc[acc_case['label']] = 1
+            else:
+                top3_acc[acc_case['label']] = 0
+
+    for label_item in cls_count:
+        top1_acc[label_item] /= max(1.0*cls_count[label_item], 0.001)
+        top3_acc[label_item] /= max(1.0*cls_count[label_item], 0.001)
+
+    print('top1_acc:', top1_acc)
+    print('top3_acc:', top3_acc)
+    print('cls_count', cls_count)
+
+    return top1_acc, top3_acc, cls_count
+
+
+
+
+
+
+

Dcl/util/train_model.py

@@ -0,0 +1,164 @@
+#coding=utf8
+from __future__ import print_function, division
+
+import os,time,datetime
+import numpy as np
+from math import ceil
+import datetime
+
+import torch
+from torch import nn
+from torch.autograd import Variable
+#from torchvision.utils import make_grid, save_image
+
+from util.utils import LossRecord, clip_gradient
+from models.focal_loss import FocalLoss
+from util.eval_model import eval_turn
+from util.Asoftmax_loss import AngleLoss
+
+import pdb
+
+def dt():
+    return datetime.datetime.now().strftime("%Y-%m-%d-%H_%M_%S")
+
+
+def train(Config,
+          model,
+          epoch_num,
+          start_epoch,
+          optimizer,
+          exp_lr_scheduler,
+          data_loader,
+          save_dir,
+          data_size=448,
+          savepoint=500,
+          checkpoint=1000
+          ):
+    # savepoint: save without evalution
+    # checkpoint: save with evaluation
+    # print(epoch_num, start_epoch) # 360  0
+    step = 0
+    eval_train_flag = False
+    rec_loss = []
+    checkpoint_list = []
+
+    train_batch_size = data_loader['train'].batch_size
+    train_epoch_step = data_loader['train'].__len__()
+    train_loss_recorder = LossRecord(train_batch_size)
+    if savepoint > train_epoch_step:
+        savepoint = 1*train_epoch_step
+        checkpoint = savepoint
+
+    date_suffix = dt()
+    log_file = open(os.path.join(Config.log_folder, 'formal_log_r50_dcl_%s_%s.log'%(str(data_size), date_suffix)), 'a')
+
+    add_loss = nn.L1Loss()
+    get_ce_loss = nn.CrossEntropyLoss()
+    get_focal_loss = FocalLoss()
+    get_angle_loss = AngleLoss()
+
+    for epoch in range(start_epoch,epoch_num-1):
+        exp_lr_scheduler.step(epoch)
+        model.train(True)
+
+        save_grad = []
+        step = 0
+        for batch_cnt, data in enumerate(data_loader['train']):
+            step += 1
+            loss = 0
+            model.train(True)
+            if Config.use_backbone:
+                inputs, labels, img_names = data
+
+                inputs = Variable(inputs.cuda())
+                labels = Variable(torch.from_numpy(np.array(labels)).cuda())
+
+            if Config.use_dcl:
+                inputs, labels, labels_swap, swap_law, img_names = data
+                inputs = Variable(inputs.cuda())
+                labels = Variable(torch.from_numpy(np.array(labels)).long().cuda())
+                labels_swap = Variable(torch.from_numpy(np.array(labels_swap)).long().cuda())
+                swap_law = Variable(torch.from_numpy(np.array(swap_law)).float().cuda())
+            optimizer.zero_grad()
+            # print(inputs.size())
+            if inputs.size(0) < 2*train_batch_size:
+                outputs = model(inputs, inputs[0:-1:2])
+            else:
+                outputs = model(inputs, None)
+            if Config.use_focal_loss:
+                ce_loss = get_focal_loss(outputs[0], labels)
+            else:
+                ce_loss = get_ce_loss(outputs[0], labels) #交叉熵损失函数 用来分类
+
+            # if Config.use_Asoftmax:
+            #     fetch_batch = labels.size(0)
+            #     if batch_cnt % (train_epoch_step // 5) == 0:
+            #         angle_loss = get_angle_loss(outputs[3], labels[0:fetch_batch:2], decay=0.9)
+            #     else:
+            #         angle_loss = get_angle_loss(outputs[3], labels[0:fetch_batch:2])
+            #     loss += angle_loss
+
+            loss += ce_loss
+            alpha_ = 1
+            beta_ = 1
+            gamma_ = 0.01 if Config.dataset == 'STCAR' or Config.dataset == 'AIR' else 1
+            if Config.use_dcl:
+                swap_loss = get_ce_loss(outputs[1], labels_swap) * beta_ # 对抗学习
+                loss += swap_loss
+                law_loss = add_loss(outputs[2], swap_law) * gamma_ # 重建学习损失
+                loss += law_loss
+
+            loss.backward()
+            torch.cuda.synchronize() #统计时间
+
+            optimizer.step()
+            torch.cuda.synchronize()
+
+            if Config.use_dcl:
+                # 8d不足就右对齐 加上-左对齐
+                print('epoch: {} step: {:-8d} / {:d} loss=ce_loss+swap_loss+law_loss: {:6.4f} = {:6.4f} + {:6.4f} + {:6.4f} '.format(epoch, step, train_epoch_step, loss.detach().item(), ce_loss.detach().item(), swap_loss.detach().item(), law_loss.detach().item()), flush=True)
+                # print(savepoint)
+            if Config.use_backbone:
+                print('step: {:-8d} / {:d} loss=ce_loss+swap_loss+law_loss: {:6.4f} = {:6.4f} '.format(step, train_epoch_step, loss.detach().item(), ce_loss.detach().item()), flush=True)
+            rec_loss.append(loss.detach().item())# 对loss进行叠加的时候一定要加detach，返回的是叶子节点，不然会进行计算图的叠加，
+
+            train_loss_recorder.update(loss.detach().item())
+            # evaluation & save
+            # checkpoint = 5000
+            if step % checkpoint == 0:
+                rec_loss = []
+                print(32*'-', flush=True)
+                print('step: {:d} / {:d} global_step: {:8.2f} train_epoch: {:04d} rec_train_loss: {:6.4f}'.format(step, train_epoch_step, 1.0*step/train_epoch_step, epoch, train_loss_recorder.get_val()), flush=True)
+                print('current lr:%s' % exp_lr_scheduler.get_lr(), flush=True)
+                if eval_train_flag:
+                    trainval_acc1, trainval_acc2, trainval_acc3 = eval_turn(Config, model, data_loader['trainval'], 'trainval', epoch, log_file)
+                    if abs(trainval_acc1 - trainval_acc3) < 0.01:
+                        eval_train_flag = False
+
+                val_acc1, val_acc2, val_acc3 = eval_turn(Config, model, data_loader['val'], 'val', epoch, log_file)
+
+                save_path = os.path.join(save_dir, 'weights_%d_%d_%.4f_%.4f.pth'%(epoch, batch_cnt, val_acc1, val_acc3))
+                torch.cuda.synchronize()
+                torch.save(model.state_dict(), save_path)
+                print('saved model to %s' % (save_path), flush=True)
+                torch.cuda.empty_cache()
+
+            # save only
+            # elif step % savepoint == 0:
+            #     print('step')
+            #     train_loss_recorder.update(rec_loss)
+            #     rec_loss = []
+            #     save_path = os.path.join(save_dir, 'savepoint_weights-%d-%s.pth'%(step, dt()))
+            #
+            #     checkpoint_list.append(save_path)
+            #     if len(checkpoint_list) == 6:
+            #         os.remove(checkpoint_list[0])
+            #         del checkpoint_list[0]
+            #     torch.save(model.state_dict(), save_path)
+            #     torch.cuda.empty_cache()
+
+
+    log_file.close()
+
+
+

Dcl/util/utils.py

@@ -0,0 +1,124 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import pdb
+
+
+class LossRecord(object):
+    def __init__(self, batch_size):
+        self.rec_loss = 0
+        self.count = 0
+        self.batch_size = batch_size
+
+    def update(self, loss):
+        if isinstance(loss, list):
+            avg_loss = sum(loss)
+            avg_loss /= (len(loss)*self.batch_size)
+            self.rec_loss += avg_loss
+            self.count += 1
+        if isinstance(loss, float):
+            self.rec_loss += loss/self.batch_size
+            self.count += 1
+
+    def get_val(self, init=False):
+        pop_loss = self.rec_loss / self.count
+        if init:
+            self.rec_loss = 0
+            self.count = 0
+        return pop_loss
+
+
+def weights_normal_init(model, dev=0.01):
+    if isinstance(model, list):
+        for m in model:
+            weights_normal_init(m, dev)
+    else:
+        for m in model.modules():
+            if isinstance(m, nn.Conv2d):
+                m.weight.data.normal_(0.0, dev)
+            elif isinstance(m, nn.Linear):
+                m.weight.data.normal_(0.0, dev)
+
+
+def clip_gradient(model, clip_norm):
+    """Computes a gradient clipping coefficient based on gradient norm."""
+    totalnorm = 0
+    for p in model.parameters():
+        if p.requires_grad:
+            modulenorm = p.grad.data.norm()
+            totalnorm += modulenorm ** 2
+    totalnorm = torch.sqrt(totalnorm).item()
+    norm = (clip_norm / max(totalnorm, clip_norm))
+    for p in model.parameters():
+        if p.requires_grad:
+            p.grad.mul_(norm)
+
+
+def Linear(in_features, out_features, bias=True):
+    """Weight-normalized Linear layer (input: N x T x C)"""
+    m = nn.Linear(in_features, out_features, bias=bias)
+    m.weight.data.uniform_(-0.1, 0.1)
+    if bias:
+        m.bias.data.uniform_(-0.1, 0.1)
+    return m
+
+
+class convolution(nn.Module):
+    def __init__(self, k, inp_dim, out_dim, stride=1, with_bn=True):
+        super(convolution, self).__init__()
+
+        pad = (k - 1) // 2
+        self.conv = nn.Conv2d(inp_dim, out_dim, (k, k), padding=(pad, pad), stride=(stride, stride), bias=not with_bn)
+        self.bn   = nn.BatchNorm2d(out_dim) if with_bn else nn.Sequential()
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        conv = self.conv(x)
+        bn   = self.bn(conv)
+        relu = self.relu(bn)
+        return relu
+
+class fully_connected(nn.Module):
+    def __init__(self, inp_dim, out_dim, with_bn=True):
+        super(fully_connected, self).__init__()
+        self.with_bn = with_bn
+
+        self.linear = nn.Linear(inp_dim, out_dim)
+        if self.with_bn:
+            self.bn = nn.BatchNorm1d(out_dim)
+        self.relu   = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        linear = self.linear(x)
+        bn     = self.bn(linear) if self.with_bn else linear
+        relu   = self.relu(bn)
+        return relu
+
+class residual(nn.Module):
+    def __init__(self, k, inp_dim, out_dim, stride=1, with_bn=True):
+        super(residual, self).__init__()
+
+        self.conv1 = nn.Conv2d(inp_dim, out_dim, (3, 3), padding=(1, 1), stride=(stride, stride), bias=False)
+        self.bn1   = nn.BatchNorm2d(out_dim)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(out_dim, out_dim, (3, 3), padding=(1, 1), bias=False)
+        self.bn2   = nn.BatchNorm2d(out_dim)
+
+        self.skip  = nn.Sequential(
+            nn.Conv2d(inp_dim, out_dim, (1, 1), stride=(stride, stride), bias=False),
+            nn.BatchNorm2d(out_dim)
+        ) if stride != 1 or inp_dim != out_dim else nn.Sequential()
+        self.relu  = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        conv1 = self.conv1(x)
+        bn1   = self.bn1(conv1)
+        relu1 = self.relu1(bn1)
+
+        conv2 = self.conv2(relu1)
+        bn2   = self.bn2(conv2)
+
+        skip  = self.skip(x)
+        return self.relu(bn2 + skip)

DclTest.py

@@ -0,0 +1,159 @@
+# coding=utf-8
+
+'''
+Input:(Device, ImagePath, ResultVisualationPath)
+DclNumberToSpeciesPath:we need to convert number to related species,so need to find a train.txt to relieve this operation
+DclNumberToSpeciesPath:../data_MPKFNet/data/DCL/datasets/HCData/anno/train.txt
+ImagePath:../data_MPKFNet/data/DCL/datasets/HCData/data
+ResultVisualationPath:../data_MPKFNet/result/Dcl_FineGrainedRecognition
+
+How to use:
+Dcl = DclInference(devices, args.ImagePath, args.ResultVisualation)
+Dcl.NumbwerToSpecies()
+Dcl.parse_args()
+Species,  ConfidenceLevel= Dcl.DclInference()
+
+'''
+import os
+import json
+import time
+import cv2
+import sys
+import csv
+import pickle
+import shutil
+import argparse
+import pandas as pd
+import numpy as np
+from math import ceil
+from tqdm import tqdm
+import PIL.Image as Image
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+from torch.nn import CrossEntropyLoss
+from torchvision import datasets, models
+import torch.backends.cudnn as cudnn
+
+# Add the absolute path to find the package
+#sys.path.append('/data1/Hatcher/model/MPKFNet/Dcl')
+from Dcl.transforms import transforms
+from Dcl.models.LoadModel import MainModel
+from Dcl.DclConfig import LoadConfig, load_data_transformers
+from Dcl.util.test_tool import set_text, save_multi_img, cls_base_acc
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '6'
+devices = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+class DclInference:
+    def __init__(self, device, DclImagePath, DclResultVisualationPath):
+        self.device = device
+        self.ImagesPath = []
+        self.DclOutputSpecies = []
+        self.DclOutputConfidenceLevel = []
+        self.DclImagePath = DclImagePath
+        self.DataSpeciesClassier = dict()
+        self.DclResultVisualationPath = DclResultVisualationPath
+        self.DclNumberToSpeciesPath = '../data_MPKFNet/data/DCL/datasets/HCData/anno/train.txt'
+        self.DclPreWeightPath = '../data_MPKFNet/weigth/FinedGrainedRecognition/Test/weights_170_999_0.8298_0.9289.pth'
+
+    def DclSoftmax(DcloutputsMatrix, DclPredictedValue):
+
+        InputMax, InputIdx = DcloutputsMatrix.max(dim=1)  # row max
+        InputExp = torch.exp(DclPredictedValue - InputMax)
+        ConfidenceLevel = InputExp / torch.sum(torch.exp(DcloutputsMatrix - InputMax))
+        ConfidenceLevel = ConfidenceLevel.item()
+        return ConfidenceLevel
+
+    def NumbwerToSpecies(self):
+        Number_last = 0
+        with open(self.DclNumberToSpeciesPath) as file:
+            Data = file.readlines()
+            for line in Data:
+                line = line.strip('\n')
+                species = line.split(' ')[0]
+                Number = line.split(' ')[1]
+                # Judge Number
+                if Number != Number_last:
+                    self.DataSpeciesClassier[Number] = str(species.split('_')[0]) + '_' + str(species.split('_')[1])
+                Number_last = Number
+
+    def parse_args(self):
+        parser = argparse.ArgumentParser(description='dcl parameters')
+        parser.add_argument('--data', dest='dataset', default='HCData', type=str)
+        parser.add_argument('--backbone', dest='backbone', default='resnet50', type=str)
+        parser.add_argument('--b', dest='batch_size', default=1, type=int)
+        parser.add_argument('--nw', dest='num_workers', default=8, type=int)
+        parser.add_argument('--ver', dest='version', default='val', type=str)
+        parser.add_argument('--size', dest='resize_resolution', default=512, type=int)
+        parser.add_argument('--crop', dest='crop_resolution', default=448, type=int)
+        parser.add_argument('--ss', dest='save_suffix', default=None, type=str)
+        parser.add_argument('--acc_report', dest='acc_report', action='store_true')
+        parser.add_argument('--swap_num', default=[7, 7], nargs=2, metavar=('swap1', 'swap2'), type=int, help='specify a range')
+        self.args = parser.parse_args()
+
+    def DclInference(self):
+
+        # Traverse Images
+        if os.path.exists(self.DclImagePath):
+            for File in os.listdir(self.DclImagePath):
+                if File.endswith('jpg'):
+                    self.ImagesPath.append(os.path.join(self.DclImagePath, File))
+        print('Load data is ok')
+
+        # Data Pretreated
+        self.Config = LoadConfig(self.args, 's')
+        transformers = load_data_transformers(self.args.resize_resolution, self.args.crop_resolution, self.args.swap_num)
+
+        # load model
+        self.model = MainModel(self.Config)
+        model_dict = self.model.state_dict()
+        pretrained_dict = torch.load(self.DclPreWeightPath)
+        model_dict.update(pretrained_dict)
+        self.model.load_state_dict(model_dict)
+        print("load model weight is ok")
+
+        self.model.to(self.device)
+        self.model.train(False)
+        count_bar = tqdm(total = len(self.ImagesPath), desc='Dcl Inference')
+        with torch.no_grad():
+            for ImgPath in self.ImagesPath:
+                count_bar.update(1)
+                Img = cv2.imread(ImgPath)
+                Img = Image.fromarray(cv2.cvtColor(Img, cv2.COLOR_BGR2RGB))
+                Img = transformers['test_totensor'](Img)
+                Img = Img.unsqueeze(0)
+                inputs = Variable(Img.cuda())
+                outputs = self.model(inputs)
+                t1 = time.time()
+                top3_val, top3_pos = torch.topk(outputs[0], 1)  # topk返回的是一个元组，第一个元素是每一行最大的三个值（默认是按行取最大值），第二个元素是每一行最大的值的排序
+                Species = top3_pos.squeeze()
+                ConfidenceLevel = DclInference.DclSoftmax(outputs[0], top3_val[0][0])
+                t2 = time.time()
+                self.DclOutputSpecies.append(Species)
+                self.DclOutputConfidenceLevel.append(ConfidenceLevel)
+                print('Post Process Time:{:.2f} ConfidenceLevel:{:.2f}'.format((t2 - t1) * 1000, ConfidenceLevel))
+            return self.DclOutputSpecies, self.DclOutputConfidenceLevel
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='dcl parameters')
+    parser.add_argument('--save', dest='resume',
+                        default='../data_MPKFNet/weigth/FinedGrainedRecognition/Test/weights_170_999_0.8298_0.9289.pth', type=str)
+    parser.add_argument('--ResultVisualation',
+                        default='../data_MPKFNet/result/Dcl_FineGrainedRecognition', type=str)
+    parser.add_argument('--ImagePath',
+                        default='../data_MPKFNet/data/DCL/datasets/HCData/data', type=str)
+    parser.add_argument('--NumberToSpecies',
+                        default='../data_MPKFNet/data/DCL/datasets/HCData/anno/train.txt', type=str)
+    args = parser.parse_args()
+    return args
+
+if __name__ == '__main__':
+    args = parse_args()
+    Dcl = DclInference(devices, args.ImagePath, args.ResultVisualation)
+    Dcl.NumbwerToSpecies()
+    Dcl.parse_args()
+    Species, ConfidenceLevel = Dcl.DclInference()
+    print(Species)

DensityVideo_Video.py

@@ -0,0 +1,226 @@
+'''
+author ：  Hatcher W huimin
+Email： hcwang@semi.ac.cn
+use: 密度估计视频处理
+input:视频路径
+output：avi后缀的结果视频
+'''
+import torch
+import datetime
+import random
+import cv2
+import shutil
+import glob
+import json
+import os
+import numpy as np
+from tqdm import tqdm
+from configs.configs import *
+from configs.configs import model_head
+import matplotlib.pyplot as plt
+from torch.utils.data import DataLoader
+from torch.utils.data import Dataset
+from dataset import People_Dataset
+from network import CSRNet,MCNN,CANNet,PKSNet
+import argparse
+np.set_printoptions(threshold=np.inf)
+
+# Path Definition
+# RawVideoFilename = '../data_MPKFNet/rawdata/2021.2.20.mp4'
+# RawVideoExtrationSavedpath = '../data_MPKFNet/result_raw/' + "2021.2.22"
+# DensityEstimationResultPath = "../data_MPKFNet/result/2021.2.22midu"
+# DensityEstimationRawImgPath = '../data_MPKFNet/result_raw/' + "2021.2.22"
+# CombineVideoRawImgPath = "../data_MPKFNet/result/2021.2.22midu/"
+# CombineVideoResultPath = '../data_MPKFNet/result_vedio/2021.2.22midu.avi'
+ModelStateDict = '../data_MPKFNet/weigth/final.pth'
+
+# Variable Definition
+display = 5
+get_name = []
+
+'''
+Cdoe Area
+'''
+# First Part--Video Frame Extrstion
+def vedio_img(input,output,fps):
+
+    isExists = os.path.exists(output)
+    if not isExists:
+        os.makedirs(output)
+        print('path of %s is build' % (output))
+    else:
+        shutil.rmtree(output)
+        os.makedirs(output)
+        print('path of %s already exist and rebuild' % (output))
+    # 开始读视频
+    videoCapture = cv2.VideoCapture(input)
+    i = 0
+    j = 0
+    while True:
+        success, frame = videoCapture.read()
+        if not success:
+            print('video is all read')
+            break
+        i += 1
+        if (i % fps == 0):
+            # 保存图片
+            j += 1
+            savedname = output+"/" + "{}".format(j).zfill(7)+ '.jpg'
+            cv2.imwrite( savedname, frame)
+            print('image of %s is saved' % (savedname))
+
+
+# Second Part--Density Etimation
+class INf_Dataset(Dataset):
+
+    def __init__(self, DensityEstimationRawImgPath,img_list, train=True):
+        random.shuffle(img_list)
+        self.DensityEstimationRawImgPath=DensityEstimationRawImgPath
+        self.img_list = img_list
+        self.train = train
+
+    def __len__(self):
+        return len(self.img_list)
+
+    def __getitem__(self, index):
+        assert index <= len(self), 'index range error'
+        path_temp = self.DensityEstimationRawImgPath + "/" + self.img_list[index]
+        img = cv2.imread(path_temp)  # H read image
+        img = cv2.resize(img, (1024, 768)) / 255.0
+        img = np.transpose(img, [2, 0, 1])
+        torch.from_numpy(img).type(torch.FloatTensor)
+        return img, self.img_list[index]
+
+
+
+
+
+def inf(DensityEstimationResultPath,devices):
+
+    net = PKSNet().to(devices)
+    checkpoint = torch.load('../data_MPKFNet/weigth/final.pth')
+    net.load_state_dict(checkpoint['net'])
+    #net.load_state_dict(torch.load('../data_MPKFNet/weigth/sheep/bestmaesheep93,mse29262.pth'))
+    print("load model weight ok")
+    i = 0
+    for img, img_pathtemp in tqdm(data_loader):
+
+        img_pathtemp = img_pathtemp[0]
+        img = torch.tensor(img)
+        img = img.to(devices)
+        if model_head:
+            targets_pred, posture_pred, spice_pred = net(img)
+        else:
+            targets_pred = net(img)
+        t1 = datetime.datetime.now()
+        #num = targets_pred.data.sum()
+        targets_pred1 = targets_pred.cpu()
+        num = targets_pred1.data.sum().numpy()
+        #print(type(num))
+        targets_pred1 = targets_pred1.data.numpy()
+
+        t2 = datetime.datetime.now()
+        if i % display == 0:
+            num_display = num+15
+        data_dict = {"sum":str(num_display),"density":(str(targets_pred1[0, 0, :, :]))}
+        #print(data_dict)
+        json_outputpath = (DensityEstimationResultPath + "/" + str(img_pathtemp)).replace("/result/", "/result_json/").replace(".jpg",".json")
+        json_data = json.dumps(data_dict,ensure_ascii=False,indent=4)
+        with open (json_outputpath,'w',encoding='utf-8') as f:
+            f.write(str(json_data))
+
+        # plt.figure(num_display)
+        # plt.title("predict:{}".format(int(num_display)))
+        # plt.imshow(targets_pred1[0, 0, :, :])
+        # plt.savefig(DensityEstimationResultPath + "/" + str(img_pathtemp))
+        # plt.close()
+
+        i += 1
+
+
+# Three Part Combine Video
+def images_to_video(path,CombineVideoResultPath):
+
+    img_array = []
+    img_list = [os.path.join(path,nm) for nm in os.listdir(path) if 1]
+    img_list.sort()
+    i=0
+    for filename in img_list:
+        print(i)
+        i+=1
+        img = cv2.imread(filename)
+        if img is None:
+            print(filename + " is error!")
+            continue
+        img_array.append(img)
+    # 图片的大小需要一致
+    #img_array, size = resize(img_array, 'largest')
+    fps = 25
+    out = cv2.VideoWriter(CombineVideoResultPath, cv2.VideoWriter_fourcc(*'DIVX'), fps, (640,480))
+    for i in range(len(img_array)):
+        out.write(img_array[i])
+    out.release()
+
+if __name__ == '__main__':
+
+    # Build Ralted Saved Path
+
+    parser = argparse.ArgumentParser(description="DensityVideodemo.")
+    parser.add_argument("video_name", default='', type=str,
+                        help='Video name.')
+    parser.add_argument("skip_getimg", default=False, type=str,
+                        help='skip_getimg.')
+    parser.add_argument("GPU", default='0', type=str,
+                        help='GPU.')
+
+    parser.add_argument(
+        "opts",
+        help="Modify config options using the command-line",
+        default=None,
+        nargs=argparse.REMAINDER,
+    )
+    args = parser.parse_args()
+    print(args)
+    # Gpu Selection
+    os.environ['CUDA_VISIBLE_DEVICES'] = args.GPU
+    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    video_name_path = args.video_name
+    video_name_path = video_name_path.replace(".MP4","").replace(".mp4","").replace(".MOV","").replace(".mov","")
+    RawVideoFilename = '../data_MPKFNet/rawdata/'+args.video_name
+    RawVideoExtrationSavedpath = '../data_MPKFNet/result_raw/' + video_name_path
+    DensityEstimationResultPath = "../data_MPKFNet/result/"+video_name_path+"desity"
+    DensityEstimationRawImgPath = '../data_MPKFNet/result_raw/' + video_name_path
+    CombineVideoRawImgPath =      "../data_MPKFNet/result/"+video_name_path+"desity"
+    Json_ouputPath = "../data_MPKFNet/result_json/" + video_name_path + "desity"
+    CombineVideoResultPath =      '../data_MPKFNet/result_vedio/'+video_name_path+"desity.avi"
+    print(args.skip_getimg)
+    print(args.GPU)
+    if not args.skip_getimg:
+
+        vedio_img(RawVideoFilename, RawVideoExtrationSavedpath, 1)
+    get_name = []
+    for _, _, filenames in os.walk(RawVideoExtrationSavedpath):  # root是一个list是正在遍历的文件夹本身的地址 dirs文件夹中的目录 files是文件夹下的所有文件
+        for filename in filenames:
+            get_name.append(filename)
+    train = RawVideoExtrationSavedpath
+    trainfile = os.listdir(RawVideoExtrationSavedpath)
+    train_json = [file for file in trainfile]
+    train_json.sort()
+    with open(inf_json, 'w') as f:
+        json.dump(train_json, f)
+    with open(inf_json, 'r') as outfile:
+        inf_list = json.load(outfile)
+    inf_dataset = INf_Dataset(DensityEstimationRawImgPath,inf_list)
+    inf_dataset.img_list.sort()
+    data_loader = DataLoader(inf_dataset, 1)
+
+
+    if not os.path.exists(DensityEstimationResultPath) or not os.path.exists(DensityEstimationResultPath):
+        os.mkdir(DensityEstimationResultPath)
+    if not os.path.exists(Json_ouputPath) :
+        os.mkdir(Json_ouputPath)
+        #os.mkdir(RawVideoExtrationSavedpath)
+    # Video Extration  Density Estimation  Combine Image
+
+    inf(DensityEstimationResultPath,device)
+    images_to_video(CombineVideoRawImgPath,CombineVideoResultPath)

DetectionVideo_Video.py

@@ -0,0 +1,130 @@
+'''
+author ：  Hatcher W huimin
+Email： hcwang@semi.ac.cn
+use: 目标检测视频处理
+input:需要处理的视频路径
+output：avi后缀的视频处理结果
+'''
+import glob
+import os
+import time
+
+import torch
+from PIL import Image
+from vizer.draw import draw_boxes
+
+from ssd.config import cfg
+from ssd.data.datasets import COCODataset, VOCDataset
+import argparse
+import numpy as np
+
+from ssd.data.transforms import build_transforms
+from ssd.modeling.detector import build_detection_model
+from ssd.utils import mkdir
+from ssd.utils.checkpoint import CheckPointer
+from detection import  main
+from vediotoima import vedio_img
+from imatovideo import images_to_video,images_to_video_big
+from moviepy.editor import *
+import os
+os.environ['CUDA_VISIBLE_DEVICES'] = '2'
+
+def parse_args():
+
+    parser = argparse.ArgumentParser(description='Detection Video to Video')
+    parser.add_argument(
+        "--config-file",
+        default="ssd/config/hat-512.yaml",
+        metavar="FILE",
+        help="path to config file",
+        type=str,
+    )
+    parser.add_argument("--ckpt", type=str, default="../data_MPKFNet/weigth/detection/detection.pth",
+                        help="Trained weights.")
+    parser.add_argument("--score_threshold", type=float, default=0.3)
+    parser.add_argument("--video_name", default='000120210501080936.mp4', type=str, help='Video name.')
+    parser.add_argument("--images_dir", default='../data_MPKFNet/result_raw/yang/', type=str,
+                        help='Specify a image dir to do prediction.')
+    parser.add_argument("--output_dir", default='../data_MPKFNet/result/yangd/', type=str,
+                        help='Specify a image dir to save predicted images.')
+    parser.add_argument("--dataset_type", default="voc", type=str,
+                        help='Specify dataset type. Currently support voc and coco.')
+
+    parser.add_argument(
+        "opts",
+        help="Modify config options using the command-line",
+        default=None,
+        nargs=argparse.REMAINDER,
+    )
+    args = parser.parse_args()
+    return args
+
+
+# if __name__ == '__main__':
+#     args = parse_args()
+#     path = "/data/Hatcher/output1/"
+#     pathlists = os.listdir(path)
+#     for pathlist in pathlists:
+#         if os.path.exists('../data_MPKFNet/rawdata/' + pathlist + ".mp4"):
+#             print(pathlist," is skipping")
+#             continue
+#         print(pathlist," is starting")
+#         L = []
+#         for root, dirs, files in os.walk(os.path.join(path,pathlist)):
+#             files.sort()
+#             for file in files:
+#                 if os.path.splitext(file)[1] in [".avi",".mp4"]:
+#                     filepath = os.path.join(root, file)
+#                     video = VideoFileClip(filepath)
+#                     L.append(video)
+#         if len(L) == 0:
+#             continue
+#         final_clip = concatenate_videoclips(L)
+#         final_clip.to_videofile(os.path.join('../data_MPKFNet/rawdata/', "{}.mp4".format(pathlist)), fps=25, remove_temp=False)
+#         video_name_path = pathlist
+#         video_name_path = video_name_path.replace(".MP4", "").replace(".mp4", "").replace(".MOV", "").replace(".mov", "")
+#         RawVideoFilename = '../data_MPKFNet/rawdata/' + video_name_path + ".mp4"
+#         RawVideoExtrationSavedpath = '../data_MPKFNet/result_raw/' + video_name_path
+#         args.images_dir =  '../data_MPKFNet/result_raw/' + video_name_path
+#         args.output_dir = "../data_MPKFNet/result/" + video_name_path + "Detection"
+#         CombineVideoRawImgPath = "../data_MPKFNet/result/" + video_name_path + "Detection"
+#         CombineVideoResultPath = '../data_MPKFNet/result_vedio11/'  + "{}.avi".format(video_name_path)
+#
+#         # vedio_img(RawVideoFilename, RawVideoExtrationSavedpath, 1)
+#         # main(args)
+#         # images_to_video(CombineVideoRawImgPath, CombineVideoResultPath)
+#         if os.path.exists(RawVideoExtrationSavedpath):
+#             print("skip video to image")
+#             pass
+#         else:
+#             vedio_img(RawVideoFilename, RawVideoExtrationSavedpath, 1)
+#         if os.path.exists(args.output_dir):
+#             print("skip deal")
+#             pass
+#         else:
+#             main(args)
+#         images_to_video(CombineVideoRawImgPath, CombineVideoResultPath)
+
+
+#zhendui danyi shipin
+if __name__ == '__main__':
+    args = parse_args()
+    video_name_path = args.video_name
+    video_name_path = video_name_path.replace(".MP4", "").replace(".mp4", "").replace(".MOV", "").replace(".mov", "")
+    RawVideoFilename = '../data_MPKFNet/rawdata/1/'+args.video_name
+    RawVideoExtrationSavedpath = '../data_MPKFNet/result_raw/' + video_name_path
+    args.images_dir =  '../data_MPKFNet/result_raw/' + video_name_path
+    args.output_dir = "../data_MPKFNet/result/" + video_name_path + "Detection"
+    CombineVideoRawImgPath = "../data_MPKFNet/result/" + video_name_path + "Detection"
+    CombineVideoResultPath = '../data_MPKFNet/result_vedio/' + video_name_path + "Detection.avi"
+    if os.path.exists(RawVideoExtrationSavedpath):
+        print("skip video to image")
+        pass
+    else:
+        vedio_img(RawVideoFilename, RawVideoExtrationSavedpath, 1)
+    if os.path.exists(args.output_dir):
+        print("skip deal")
+        pass
+    else:
+        main(args)
+    images_to_video(CombineVideoRawImgPath, CombineVideoResultPath)

GYD_RedCamera.py

@@ -0,0 +1,44 @@
+import os,torch,cv2
+from System_I import  Detection_S,Fine_grained
+from vediotoima import vedio_img
+def GYD_RedCameraDeal (root,DetectionCalss,Fine_grainedClass):
+    CameraLists = os.listdir(root)
+    for CameraIMEI in CameraLists:
+        if not os.path.exists(os.path.join(root+"_Result",CameraIMEI)):
+            os.mkdir(os.path.join(root+"_Result",CameraIMEI))
+        image_videolists = os.listdir(os.path.join(root,CameraIMEI))
+        for image_video in image_videolists:
+            if "JPG" in image_video:
+                pass
+                # print(image_video, " is starting")
+                # image = cv2.imread(os.path.join(root,CameraIMEI,image_video))
+                # DetectionCalss.inference_Red(image,os.path.join(root,CameraIMEI,image_video),Fine_grainedClass)
+                # print(image_video," is ok")
+            else:
+                video_path = os.path.join(root, CameraIMEI,image_video)
+                video_path_out = os.path.join(root + "_Result", CameraIMEI, "RAWImage",image_video.replace(".MP4",""))
+                vedio_img(video_path,video_path_out,25)
+                isExists = os.path.exists(os.path.join(root + "_Result", CameraIMEI, "VideoResult",image_video.replace(".MP4","")))
+                if not isExists:
+                    os.makedirs(os.path.join(root + "_Result", CameraIMEI, "VideoResult",image_video.replace(".MP4","")))
+                img_lists = os.listdir(video_path_out)
+                img_lists.sort()
+       
+                for img in img_lists:
+
+                    image = cv2.imread(os.path.join(video_path_out, img))
+                    DetectionCalss.inference_Red(image, os.path.join(root, CameraIMEI, "VideoResult",image_video.replace(".MP4",""),img),
+                                                 Fine_grainedClass)
+
+
+
+
+
+
+if __name__ == '__main__':
+    os.environ['CUDA_VISIBLE_DEVICES'] = '4'
+    devices = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    root = "/data1/Hatcher/ftp"
+    Detection1 = Detection_S(devices, root)
+    Fine_grained1 = Fine_grained(devices, root)
+    GYD_RedCameraDeal(root,Detection1,Fine_grained1)

Hatcher_cam.py

@@ -0,0 +1,56 @@
+
+import os
+import torch
+from System_I import Density,Detection_Y5,Fine_grained,ThreadMonitor_mul,isopenmul
+import cv2
+import time
+import redis
+from rtsp_process import rtspThread,cusLogger
+from common import global_para as gp
+import globalvar as gl
+
+r_db = None
+#if __name__ == '__main__':
+def Hatcher():
+    gl._init()
+    log = cusLogger("rtsp_process").getLogger()
+
+    Density1 = Density(gp.devices,gp.root)
+    Detection1 = Detection_Y5(gp.devices, gp.root)
+    Fine_grained1 = Fine_grained(gp.devices,gp.root)
+    #     Density1.inference(image,image_time,str(1))
+    #     Detection1.inference(image,image_time,str(1),Fine_grained1)
+    try:
+        # 如果有多个流要同时取存，可以创建多个线程
+        thread1 = rtspThread("thread_vs",
+                             "rtsp://admin:glodon123@192.168.116.13:554/Streaming/Channels/101",
+                             "ch1_",
+                             "",# 为空串时，默认保存在当前目录
+                             25,
+                             (3840, 2160),
+                             #(3264,2448),
+                             150,
+                             30,
+                             "key",
+                             1,
+                             Density1,
+                             Detection1,
+                             Fine_grained1
+                             )
+
+        thread1.start()
+        # ThreadMonitor_mul1 = ThreadMonitor_mul("0001","ThreadMonitor_mul")   #开启线程管理
+        # ThreadMonitor_mul1.start()
+
+
+
+
+
+    except Exception as e:
+        log.info(e)
+
+
+    #
+    #     Density1.inference(image,image_time,str(1))
+    #     Detection1.inference(image,image_time,str(1),Fine_grained1)
+    #

Hatcher_is_cool.py

@@ -0,0 +1,58 @@
+'''
+author ：  Hatcher W
+Email： hcwang@semi.ac.cn
+use: 递归一个文件下面所有的视频和图像的路径
+input:一个文件夹目录
+output：两个json文件，分别存储视频和图像的路径
+'''
+import os
+import shutil
+import  json
+
+vediofile = []
+imagefile = []
+vedio_format = [".MP4" , ".mp4" , ".MOV" , ".mov"]
+image_format = [".JPG" , ".JPEG" , ".jpg"]
+def eachfile(filepath):
+    pathdir = os.listdir(filepath)
+    for s in pathdir:
+        newdir = os.path.join(filepath,s)#将文件名加到当前文件路径后面
+        if os.path.isfile(newdir):  #判断是不是文件
+            print(os.path.splitext(newdir)[1])
+            if os.path.splitext(newdir)[1] in vedio_format:
+                vediofile.append(newdir)
+            elif os.path.splitext(newdir)[1]in image_format:
+                imagefile.append(newdir)
+            elif os.path.isdir(newdir):
+                eachfile(newdir)
+        else:
+            pathdir_1 = os.listdir(newdir)   #取出下一层的路径
+            for s_1 in pathdir_1:           #遍历下一层的路径
+                newdir_1 = os.path.join(newdir,s_1)     #取出其中一个路径
+                if os.path.isfile(newdir_1):  # 判断是不是文件
+                    print(os.path.splitext(newdir_1)[1])
+                    if os.path.splitext(newdir_1)[1] in vedio_format:
+                        vediofile.append(newdir_1)
+                    elif os.path.splitext(newdir_1)[1] in image_format:
+                        imagefile.append(newdir_1)
+                    elif os.path.isdir(newdir_1):
+                        eachfile(newdir_1)
+                else:eachfile(newdir_1)
+    return vediofile,imagefile
+
+fp = "/data1/Hatcher/DATA/Hat_Raw_Data"
+
+vedio_f,image_f = eachfile(fp)
+print(vedio_f)
+print(image_f)
+with open ("./vedioroot.json",'w') as f:
+    json.dump(vedio_f,f)
+    f.close()
+with open ("./imageroot.json",'w') as g:
+    json.dump(image_f,g)
+    g.close()
+for i in range(len(vedio_f)):
+    pcm_file = vedio_f[i]
+    path_0 = os.path.split(vedio_f[i])[0]
+    path_1 = os.path.split(vedio_f[i])[1]
+    #print(path_1)

Hatcher_red.py

@@ -0,0 +1,331 @@
+import os,torch,cv2,time
+from System_I import  Detection_Y5,Fine_grained
+from vediotoima import vedio_img
+from common import global_para as gp
+from video_notice import VideoNotice
+#!/usr/bin/python
+# -*- coding: UTF-8 -*-
+
+from ftplib import FTP, error_perm
+import socket
+import sys
+import time
+import os
+def GYD_RedCameraDeal (root,DetectionCalss,Fine_grainedClass):
+    CameraLists = os.listdir(root)
+    print(CameraLists)
+    for CameraIMEI in CameraLists:
+        if not os.path.exists(os.path.join(root+"_Result",CameraIMEI)):
+            os.makedirs(os.path.join(root+"_Result",CameraIMEI))
+        image_videolists = os.listdir(os.path.join(root,CameraIMEI))
+        for image_video in image_videolists:
+            print(image_videolists)
+            if "JPG" in image_video:
+                pass
+                # print(image_video, " is starting")
+                # image = cv2.imread(os.path.join(root,CameraIMEI,image_video))
+                # DetectionCalss.inference_Red(image,os.path.join(root,CameraIMEI,image_video),Fine_grainedClass)
+                # print(image_video," is ok")
+            else:
+                video_path = os.path.join(root, CameraIMEI,image_video)
+                video_path_out = os.path.join(root + "_Result", CameraIMEI, "RAWImage",image_video.replace(".MP4",""))
+                vedio_img(video_path,video_path_out,25)
+                isExists = os.path.exists(os.path.join(root + "_Result", CameraIMEI, "VideoResult",image_video.replace(".MP4","")))
+                if not isExists:
+                    os.makedirs(os.path.join(root + "_Result", CameraIMEI, "VideoResult",image_video.replace(".MP4","")))
+                img_lists = os.listdir(video_path_out)
+                img_lists.sort()
+
+                for img in img_lists:
+
+                    image = cv2.imread(os.path.join(video_path_out, img))
+                    DetectionCalss.inference_Red(image, os.path.join(root, CameraIMEI, "VideoResult",image_video.replace(".MP4",""),img),
+                                                 Fine_grainedClass)
+
+
+
+
+
+
+#if __name__ == '__main__':
+
+
+
+
+
+class MyFTP:
+    def __init__(self, host, port):
+        """ 初始化 FTP 客户端
+        参数:
+                 host:ip地址
+
+                 port:端口号
+        """
+        # print("__init__()---> host = %s ,port = %s" % (host, port))
+
+        self.host = host
+        self.port = port
+        self.ftp = FTP()
+        # 重新设置下编码方式
+        self.ftp.encoding = 'gbk'
+        #self.log_file = open("/home/gyd/ftp/log/log.txt", "a")
+        self.file_list = []
+
+
+    def login(self, username, password):
+        """ 初始化 FTP 客户端
+            参数:
+                username: 用户名
+
+                password: 密码
+            """
+        try:
+            timeout = 1000
+            socket.setdefaulttimeout(timeout)
+            # 0主动模式 1 #被动模式
+            self.ftp.set_pasv(True)
+            # 打开调试级别2，显示详细信息
+            # self.ftp.set_debuglevel(2)
+
+            self.debug_print('开始尝试连接到 %s' % self.host)
+            self.ftp.connect(self.host, self.port)
+            self.debug_print('成功连接到 %s' % self.host)
+
+            self.debug_print('开始尝试登录到 %s' % self.host)
+            self.ftp.login(username, password)
+            self.debug_print('成功登录到 %s' % self.host)
+
+            self.debug_print(self.ftp.welcome)
+        except Exception as err:
+            self.debug_print("FTP 连接或登录失败 ，错误描述为：%s" % err)
+            pass
+
+    def close(self):
+        """ 退出ftp
+        """
+        self.debug_print("close()---> FTP退出")
+        self.ftp.quit()
+
+    def get_file_list(self, line):
+        """ 获取文件列表
+            参数：
+                line：
+        """
+        file_arr = self.get_file_name(line)
+        # 去除  . 和  ..
+        if file_arr[1] not in ['.', '..']:
+            self.file_list.append(file_arr)
+
+    def get_file_name(self, line):
+        """ 获取文件名
+            参数：
+                line：
+        """
+        pos = line.rfind(':')
+        while (line[pos] != ' '):
+            pos += 1
+        while (line[pos] == ' '):
+            pos += 1
+        file_arr = [line[0], line[pos:]]
+        return file_arr
+
+    def download_file(self, local_file, remote_file):
+        """从ftp下载文件
+            参数:
+                local_file: 本地文件
+
+                remote_file: 远程文件
+        """
+        self.debug_print("download_file()---> local_path = %s ,remote_path = %s" % (local_file, remote_file))
+
+        if self.is_same_size(local_file, remote_file):
+            self.debug_print('%s 文件大小相同，无需下载' % local_file)
+            return
+        else:
+            try:
+                self.debug_print('>>>>>>>>>>>>下载文件 %s ... ...' % local_file)
+                buf_size = 1024
+                file_handler = open(local_file, 'wb')
+                self.ftp.retrbinary('RETR %s' % remote_file, file_handler.write, buf_size)
+                file_handler.close()
+            except Exception as err:
+                self.debug_print('下载文件出错，出现异常：%s ' % err)
+                return
+
+    def download_file_tree(self, local_path, remote_path):
+        """从远程目录下载多个文件到本地目录
+                       参数:
+                         local_path: 本地路径
+
+                         remote_path: 远程路径
+                """
+        print("download_file_tree()--->  local_path = %s ,remote_path = %s" % (local_path, remote_path))
+        try:
+            self.ftp.cwd(remote_path)
+        except Exception as err:
+            self.debug_print('远程目录%s不存在，继续...' % remote_path + " ,具体错误描述为：%s" % err)
+            return
+
+        if not os.path.isdir(local_path):
+            self.debug_print('本地目录%s不存在，先创建本地目录' % local_path)
+            os.makedirs(local_path)
+
+        self.debug_print('切换至目录: %s' % self.ftp.pwd())
+
+        self.file_list = []
+        # 方法回调
+        self.ftp.dir(self.get_file_list)
+
+        remote_names = self.file_list
+        self.debug_print('远程目录 列表: %s' % remote_names)
+        for item in remote_names:
+            file_type = item[0]
+            file_name = item[1]
+            local = os.path.join(local_path, file_name)
+            if file_type == 'd':
+                print("download_file_tree()---> 下载目录： %s" % file_name)
+                self.download_file_tree(local, file_name)
+                self.ftp.cwd("..")
+                self.debug_print('返回上层目录 %s' % self.ftp.pwd())
+            elif file_type == '-':
+                print("download_file()---> 下载文件： %s" % file_name)
+                self.download_file(local, file_name)
+            """
+            """
+        return True
+
+    def upload_file(self, local_file, remote_file):
+        """从本地上传文件到ftp
+
+           参数:
+             local_path: 本地文件
+
+             remote_path: 远程文件
+        """
+        if not os.path.isfile(local_file):
+            self.debug_print('%s 不存在' % local_file)
+            return
+
+        if self.is_same_size(local_file, remote_file):
+            self.debug_print('跳过相等的文件: %s' % local_file)
+            return
+
+        buf_size = 1024
+        file_handler = open(local_file, 'rb')
+        self.ftp.storbinary('STOR %s' % remote_file, file_handler, buf_size)
+        file_handler.close()
+        self.debug_print('上传: %s' % local_file + "成功!")
+
+    def upload_file_tree(self, local_path, remote_path):
+        """从本地上传目录下多个文件到ftp
+           参数:
+
+             local_path: 本地路径
+
+             remote_path: 远程路径
+        """
+        if not os.path.isdir(local_path):
+            self.debug_print('本地目录 %s 不存在' % local_path)
+            return
+
+        try:
+            self.ftp.cwd(remote_path)
+            self.debug_print('切换至远程目录: %s' % self.ftp.pwd())
+        except error_perm:
+            self.debug_print('远程目录%s不存在，先创建远程目录' % remote_path)
+            self.ftp.mkd(remote_path)
+            self.ftp.cwd(remote_path)
+            self.debug_print('切换至远程目录: %s' % self.ftp.pwd())
+
+        local_name_list = os.listdir(local_path)
+        for local_name in local_name_list:
+            src = os.path.join(local_path, local_name)
+            if os.path.isdir(src):
+                try:
+                    self.ftp.mkd(local_name)
+                except Exception as err:
+                    self.debug_print("目录已存在 %s ,具体错误描述为：%s" % (local_name, err))
+                self.debug_print("upload_file_tree()---> 上传目录： %s" % local_name)
+                self.upload_file_tree(src, local_name)
+            else:
+                self.debug_print("upload_file_tree()---> 上传文件： %s" % local_name)
+                self.upload_file(src, local_name)
+        self.ftp.cwd("..")
+
+    def debug_print(self, s):
+        """ 打印日志
+        """
+        self.write_log(s)
+
+    def write_log(self, log_str):
+        """ 记录日志
+            参数：
+                log_str：日志
+        """
+        time_now = time.localtime()
+        date_now = time.strftime('%Y-%m-%d %H:%M:%S', time_now)
+        format_log_str = "%s ---> %s \n " % (date_now, log_str)
+        print(format_log_str)
+        #self.log_file.write(format_log_str)
+
+    def deal_error(self, e):
+        """ 处理错误异常
+            参数：
+                e：异常
+        """
+        log_str = '发生错误: %s' % e
+        self.write_log(log_str)
+        sys.exit()
+
+    def is_same_size(self, local_file, remote_file):
+        """判断远程文件和本地文件大小是否一致
+
+           参数:
+             local_file: 本地文件
+
+             remote_file: 远程文件
+        """
+        try:
+            remote_file_size = self.ftp.size(remote_file)
+        except Exception as err:
+            # self.debug_print("is_same_size() 错误描述为：%s" % err)
+            remote_file_size = -1
+
+        try:
+            local_file_size = os.path.getsize(local_file)
+        except Exception as err:
+            # self.debug_print("is_same_size() 错误描述为：%s" % err)
+            local_file_size = -1
+
+        self.debug_print('local_file_size:%d  , remote_file_size:%d' % (local_file_size, remote_file_size))
+        if remote_file_size == local_file_size:
+            return 1
+        else:
+            return 0
+
+
+def Hatcher_R():
+    print("start system")
+    data_today = time.strftime('%Y%m%d',time.localtime(time.time()))
+    my_ftp = MyFTP("113.204.49.10", 802)
+    my_ftp.login("camera", "Birdsdata2020")
+    my_ftp.download_file_tree(os.path.join("/home/gyd/Infrared",data_today,"ftp","860548046389083"),"/home/camera/860548046389083/")
+    print("860548046389083 down")
+    my_ftp.download_file_tree(os.path.join("/home/gyd/Infrared", data_today,"ftp", "860548046386865"), "/home/camera/860548046386865/")
+    print("860548046386865 down")
+    my_ftp.download_file_tree(os.path.join("/home/gyd/Infrared", data_today,"ftp", "860548042913449"), "/home/camera/860548042913449/")
+    print("860548042913449 down")
+    my_ftp.download_file_tree(os.path.join("/home/gyd/Infrared", data_today,"ftp", "860548046390768"), "/home/camera/860548046390768/")
+    print("860548046390768 down")
+    #my_ftp.upload_file_tree("/data1/changhao/ftp/download/20210422/", "/home/camera/test/")
+    my_ftp.close()
+    root = os.path.join("/home/gyd/Infrared", data_today,"ftp")
+    Detection1 = Detection_Y5(gp.devices, root)
+    Fine_grained1 = Fine_grained(gp.devices, root)
+    GYD_RedCameraDeal(root,Detection1,Fine_grained1)
+    time.sleep(100)
+    notice_model = VideoNotice('http://192.168.200.29:1001/guangyd/videoNotice', "Infrared")
+    content = notice_model.send_notice()
+    print("一批数据处理完毕，已发送信号至服务端来取")
+# import time
+# print(time.strftime('%Y%m%d',time.localtime(time.time())))

PlayVideo.py

@@ -0,0 +1,26 @@
+import cv2
+import argparse
+parser = argparse.ArgumentParser(description="DensityVideodemo.")
+parser.add_argument("video_path", default='', type=str,
+                    help='Video path.')
+
+
+parser.add_argument(
+    "opts",
+    help="Modify config options using the command-line",
+    default=None,
+    nargs=argparse.REMAINDER,
+)
+args = parser.parse_args()
+print(args)
+
+#filepath = "/data1/Hatcher/model/data_MPKFNet/rawdata/1_3.mp4"
+cap = cv2.VideoCapture(args.video_path)
+while (cap.isOpened()):
+    ret,frame = cap.read()
+    cv2.imshow('image',frame)
+    k = cv2.waitKey(10)
+    if(k&0xff == ord('q')):
+        break
+cap.release()
+cv2.destroyAllWindows()

System_I.py

@@ -0,0 +1,941 @@
+"""
+密度估计
+"""
+
+import cv2
+import json
+import random
+from configs.configs import model_head
+from network import PKSNet
+import os,sys
+import torch
+from PIL import Image
+from vizer.draw import draw_boxes
+from ssd.config import cfg
+from ssd.data.datasets import  VOCDataset
+
+import numpy as np
+from ssd.data.transforms import build_transforms
+from ssd.modeling.detector import build_detection_model
+from ssd.utils.checkpoint import CheckPointer
+from creat_xml_fordetection import make_xml
+from Dcl.models.LoadModel import MainModel
+from Dcl.DclConfig import LoadConfig, load_data_transformers
+from torch.autograd import Variable
+import threading
+from common import global_para as  gp
+import time
+from numpy import random
+root_path = os.getcwd()
+sys.path.insert(0,root_path+"/Y5")
+from Y5.models.experimental import attempt_load
+from Y5.utils.datasets import  letterbox
+from Y5.utils.general import   non_max_suppression, \
+    scale_coords
+from Y5.utils.plots import plot_one_box
+
+
+
+
+
+#np.set_printoptions(threshold=np.inf)
+class Density:
+    def __init__(self, device, root):
+            self.device = device
+            self.root = root
+            self.net = PKSNet().to(self.device)
+            self.pth = "../data_MPKFNet/weigth/final.pth"
+           # print("creat model is ", (t4 - t3) * 1000)
+           #  self.checkpointer = CheckPointer(self.net)
+           #  self.net.load_state_dict(self.checkpointer['net'])
+           #  self.checkpointer.load(self.checkpointer['net'])
+            checkpoint = torch.load(self.pth)
+            self.net.load_state_dict(checkpoint['net'])
+            if gp.DebugPrint_flag:
+                print("load model weight ok")
+
+            np.set_printoptions(threshold=np.inf)
+
+
+
+    def inference(self,image,imagetime,index):
+        rootpath = os.path.join(self.root, imagetime)
+        if not os.path.exists(rootpath):
+            os.mkdir(rootpath)
+        if not  os.path.exists(rootpath + "/" + "Density"+"/"):
+            os.mkdir(rootpath + "/" + "Density"+"/")
+        target_h, target_w = image.shape[0], image.shape[1]
+        t1 = time.time()
+        image = cv2.resize(image, (1024, 768)) / 255.0
+        image = np.transpose(image, [2, 0, 1])
+        image = np.expand_dims(image, axis=0)
+        image = torch.tensor(image)
+        image = image.to(self.device)
+        t2 = time.time()
+        if gp.DebugPrint_flag:
+            print("Density data deal is ", (t2 - t1) * 1000)
+
+       # print("data dell is ", (t7 - t6) * 1000)
+
+        if model_head:
+            targets_pred, posture_pred, spice_pred = self.net(image)
+        else:
+            targets_pred = self.net(image)
+        t3=time.time()
+        if gp.DebugPrint_flag:
+            print("Density inference is ", (t3 - t2) * 1000)
+        targets_pred1 = targets_pred.cpu()
+        num = targets_pred1.data.sum().numpy()
+        targets_pred1 = targets_pred1.data.numpy()
+        t4=time.time()
+        if gp.DebugPrint_flag:
+            print("Density final is ", (t4 - t3) * 1000)
+
+        data_dict = {"sum": str(num), "density": (str(targets_pred1[0, 0, :, :]))}
+        json_outputpath = (rootpath + "/" + "Density"+"/"+str(index).zfill(3)+".json")
+        json_data = json.dumps(data_dict, ensure_ascii=True, indent=4)
+        with open (json_outputpath,'w',encoding='utf-8') as f:
+            f.write(str(json_data))
+        t5=time.time()
+        if gp.DebugPrint_flag:
+            print("Density save is ", (t5 - t4) * 1000)
+
+    def inference_init(self, image):
+        image = cv2.resize(image, (1024, 768)) / 255.0
+        image = np.transpose(image, [2, 0, 1])
+        image = np.expand_dims(image, axis=0)
+        image = torch.tensor(image)
+        image = image.to(self.device)
+
+        if model_head:
+            targets_pred, posture_pred, spice_pred = self.net(image)
+        else:
+            targets_pred = self.net(image)
+
+        targets_pred1 = targets_pred.cpu()
+        num = targets_pred1.data.sum().numpy()
+        targets_pred1 = targets_pred1.data.numpy()
+        if gp.DebugPrint_flag:
+            print("Density initing")
+
+"""
+目标检测
+"""
+class Detection_S:
+    def __init__(self,device,root):
+        self.device = device
+        self.root = root
+        cfg.merge_from_file("ssd/config/hat-512.yaml")
+        cfg.freeze()
+        self.net =build_detection_model(cfg).to(self.device)
+        self.checkpointer = CheckPointer(self.net, save_dir=cfg.OUTPUT_DIR)
+        self.ckpt = "../data_MPKFNet/weigth/detection/detection.pth"
+        self.checkpointer.load(self.ckpt, use_latest=self.ckpt is None)
+        self.transforms = build_transforms(cfg, is_train=False)
+        self.threshold = gp.detect_thresh
+        self.net.eval()
+    def inference(self,image,imagetime,index,fine_grained):
+        with torch.no_grad():
+            class_names = VOCDataset.class_names
+            rootpath = os.path.join(self.root, imagetime)
+            if not os.path.exists(rootpath):
+                os.mkdir(rootpath)
+            if not os.path.exists(rootpath + "/" + "Detection" + "/"):
+                os.mkdir(rootpath + "/" + "Detection" + "/")
+            image_or = image
+            height, width = image.shape[:2]
+            t1 = time.time()
+            image = self.transforms(image)[0].unsqueeze(0)
+            t2 = time.time()
+            if gp.DebugPrint_flag:
+                print("Detection data deal is ",(t2-t1)*1000)
+            result = self.net(image.to(self.device))[0]
+            t3 = time.time()
+            if gp.DebugPrint_flag:
+                print("Detection inference is ",(t3-t2)*1000)
+            result = result.to(torch.device("cpu"))
+            result = result.resize((width, height)).numpy()
+            t5 = time.time()
+            if gp.DebugPrint_flag:
+                print("Detection get result is ", (t5 - t3) * 1000)
+            boxes, labels, scores = result['boxes'], result['labels'], result['scores']
+            indices_score = (scores > self.threshold)
+            indices_label_left =  labels <4
+            indices_label_right = labels >2
+            indices = []
+            for i in range(len(indices_score)):
+                indices.append(indices_score[i] and indices_label_left[i] and indices_label_right[i])
+            boxes = boxes[indices]
+            labels = labels[indices]
+            scores = scores[indices]
+            for i in range(len(boxes)):
+                if boxes[i][0]<0:
+                    boxes[i][0] = 0
+                if boxes[i][1]<0:
+                    boxes[i][1] = 0
+                if boxes[i][2]<0:
+                    boxes[i][2] = 0
+                if boxes[i][3]<0:
+                    boxes[i][3] = 0
+            for i in range(len(labels)):
+                labels[i] = 99999
+            boxes_fine, labels_fine, scores_fine = fine_grained.inference(image_or, boxes, labels, scores)
+            t4 = time.time()
+            if gp.DebugPrint_flag:
+                print("Detection final is ", (t4 - t3) * 1000)
+            #numb = 2
+            #fine_grained.inference(image_or,boxes,labels,scores)
+            # for box in boxes:
+            #     img1 = image_or.copy()[int(box[1]):int(box[3]), int(box[0]):int(box[2]), :]
+            #     cv2.imwrite(rootpath + "/" + "Detection" + "/"+"{}.jpg".format(numb), img1)
+            #     numb+=1
+            #drawn_image = draw_boxes(image_or, boxes_fine, labels_fine, scores_fine).astype(np.uint8)
+            xml_outputpath = rootpath + "/" + "Detection"+"/" + index+".xml"
+            make_xml(str(index),image_or.shape[0],image_or.shape[1],boxes_fine,labels_fine,scores_fine,xml_outputpath,class_names)
+            #Image.fromarray(drawn_image).save(xml_outputpath.replace(".xml",".jpg"))
+            t5 = time.time()
+            if gp.DebugPrint_flag:
+                print("Detection save is ", (t5 - t4) * 1000)
+    def inference_Red(self,image,ImagePath,fine_grained):
+        with torch.no_grad():
+            class_names = VOCDataset.class_names
+            image_or = image
+            height, width = image.shape[:2]
+            t1 = time.time()
+            image = self.transforms(image)[0].unsqueeze(0)
+            t2 = time.time()
+            if gp.DebugPrint_flag:
+                print("data deal is ",(t2-t1)*1000)
+            result = self.net(image.to(self.device))[0]
+            t3 = time.time()
+            if gp.DebugPrint_flag:
+                print("inference is ",(t3-t2)*1000)
+            result = result.to(torch.device("cpu"))
+            result = result.resize((width, height)).numpy()
+            boxes, labels, scores = result['boxes'], result['labels'], result['scores']
+            indices_score = (scores > self.threshold)
+            indices_label_left =  labels <4
+            indices_label_right = labels >2
+            indices = []
+            for i in range(len(indices_score)):
+                indices.append(indices_score[i] and indices_label_left[i] and indices_label_right[i])
+            boxes = boxes[indices]
+            for i in range(len(boxes)):
+                if boxes[i][0]<0:
+                    boxes[i][0] = 0
+                if boxes[i][1]<0:
+                    boxes[i][1] = 0
+                if boxes[i][2]<0:
+                    boxes[i][2] = 0
+                if boxes[i][3]<0:
+                    boxes[i][3] = 0
+
+            labels = labels[indices]
+            scores = scores[indices]
+            for i in range(len(labels)):
+                labels[i] = 99999
+            boxes_fine, labels_fine, scores_fine = fine_grained.inference(image_or, boxes, labels, scores)
+            t4 = time.time()
+            if gp.DebugPrint_flag:
+                print("final is ", (t4 - t3) * 1000)
+            #numb = 2
+            #fine_grained.inference(image_or,boxes,labels,scores)
+            # for box in boxes:
+            #     img1 = image_or.copy()[int(box[1]):int(box[3]), int(box[0]):int(box[2]), :]
+            #     cv2.imwrite(rootpath + "/" + "Detection" + "/"+"{}.jpg".format(numb), img1)
+            #     numb+=1
+            #drawn_image = draw_boxes(image_or, boxes_fine, labels_fine, scores_fine).astype(np.uint8)
+            xml_outputpath = ImagePath.replace("/ftp/", "/ftp_Result/").replace(".JPG", ".xml").replace(".jpg",".xml")
+            make_xml(str(ImagePath),image_or.shape[0],image_or.shape[1],boxes_fine,labels_fine,scores_fine,xml_outputpath)
+            #Image.fromarray(drawn_image).save(xml_outputpath.replace(".xml",".jpg"))
+            t5 = time.time()
+            if gp.DebugPrint_flag:
+                print("save is ", (t5 - t4) * 1000)
+
+    def check_object(self,image):
+        flag = False
+        with torch.no_grad():
+            height, width = image.shape[:2]
+            img_temp = image
+            t1 = time.time()
+            image = self.transforms(image)[0].unsqueeze(0)
+            t2 = time.time()
+            if gp.DebugPrint_flag:
+                print("check datadeal is ",(t2-t1)*1000)
+            result = self.net(image.to(self.device))[0]
+            t3 = time.time()
+            if gp.DebugPrint_flag:
+                print("check inference is ", (t3 - t2) * 1000)
+            result = result.to(torch.device("cpu"))
+            t4 = time.time()
+            if gp.DebugPrint_flag:
+                print("check result deal is ", (t4 - t3) * 1000)
+            result = result.resize((width, height)).numpy()
+
+            boxes, labels, scores = result['boxes'], result['labels'], result['scores']
+            indices_score = (scores > self.threshold)
+            indices_label_left =  labels <4
+            indices_label_right = labels >2
+            indices = []
+            for i in range(len(indices_score)):
+                indices.append(indices_score[i] and indices_label_left[i] and indices_label_right[i])
+            labels = labels[indices]
+            boxes = boxes[indices]
+            scores=scores[indices]
+            if 3 in labels:
+                flag = True
+            if gp.DebugPrint_flag:
+                print("boxes = ",boxes)
+                print("check flag is  ",flag)
+            t5 = time.time()
+            if gp.DebugPrint_flag:
+                print("check final is ", (t5 - t4) * 1000)
+        return flag,labels,boxes,scores
+class Detection_Y5:
+    def __init__(self,device,root):
+        self.device = device
+        self.root = root
+        self.img_size = 640
+        self.conf_thres = gp.detect_thresh
+        self.iou_thres = 0.4
+        self.stride = 1
+        if gp.red_flag:
+            self.weights = '../data_MPKFNet/weigth/detection/Y/best_r.pt'
+        else:
+            self.weights = '../data_MPKFNet/weigth/detection/Y/best.pt'
+        # self.device = select_device(self.device)
+        model = attempt_load(self.weights, map_location=self.device)
+        model.to(self.device).eval()
+        model.half()
+        self.model = model
+        self.names = model.module.names if hasattr(
+            model, 'module') else model.names
+    def preprocess(self, img):
+
+        img0 = img.copy()
+        #img = cv2.resize(img, (1024, 768))
+        img = letterbox(img, new_shape=self.img_size, stride=self.stride)[0]
+        img = img[:, :, ::-1].transpose(2, 0, 1)
+        img = np.ascontiguousarray(img)
+        img = torch.from_numpy(img).to(self.device)
+        img = img.half()  # °ë¾«¶È
+        img /= 255.0  # ͼÏñ¹éÒ»»¯
+        if img.ndimension() == 3:
+            img = img.unsqueeze(0)
+
+        return img0, img
+    def inference(self,image,imagetime,index,fine_grained):
+        with torch.no_grad():
+            rootpath = os.path.join(self.root, imagetime)
+            if not os.path.exists(rootpath):
+                os.mkdir(rootpath)
+            if not os.path.exists(rootpath + "/" + "Detection" + "/"):
+                os.mkdir(rootpath + "/" + "Detection" + "/")
+            image_or = image
+            img_detect = cv2.resize(image, (1024, 768))
+            t1 = time.time()
+            image_scale,image = self.preprocess(img_detect)
+            t2 = time.time()
+            if gp.DebugPrint_flag:
+                print("Detection data deal is ",(t2-t1)*1000)
+            pred = self.model(image, augment=False)[0]
+            pred = pred.float()
+            pred = non_max_suppression(pred, self.conf_thres, self.iou_thres)
+            t3 = time.time()
+            if gp.DebugPrint_flag:
+                print("Detection inference is ",(t3-t2)*1000)
+            boxes = []
+            scores = []
+            labels = []
+            for det in pred:
+                if det is not None and len(det):
+                    det[:, :4] = scale_coords(
+                        image.shape[2:], det[:, :4], image_scale.shape).round()
+
+                    boxes.append(det[:, :4])
+                    scores.append(det[:, 4])
+                    labels.append((det[:, -1]))
+                    boxes = boxes[0]
+                    scores = scores[0]
+                    labels = labels[0]
+            t5 = time.time()
+            if gp.DebugPrint_flag:
+                print("Detection get result is ", (t5 - t3) * 1000)
+#cichuxuyaoxiugai   ruguo meiyoujiance daomubiao jiuhuibaocuo
+            if len(labels) > 0:
+
+                indices_label_left = labels < 4
+                indices_label_right = labels > 2
+                indices = []
+                for i in range(len(indices_label_left)):
+                    indices.append(indices_label_left[i] and indices_label_right[i])
+
+                boxes = boxes[indices]
+                labels = labels[indices]
+                scores = scores[indices]  # 因为莫名奇妙多了一个list维度，所以加了这个0
+                boxes = boxes[0]
+                labels = labels[0]
+                scores = scores[0]
+                ratio_x = (3840.0 / 1024.0)
+                ratio_y = (2160.0 / 768.0)
+
+
+                for i in range(len(boxes)):
+                    boxes[i][0]=boxes[i][0] * ratio_x
+                    boxes[i][2]=boxes[i][2] * ratio_x
+                    boxes[i][1] = boxes[i][1] * ratio_y
+                    boxes[i][3] = boxes[i][3] * ratio_y
+                    if boxes[i][0]<0:
+                        boxes[i][0] = 0
+                    if boxes[i][1]<0:
+                        boxes[i][1] = 0
+                    if boxes[i][2]<0:
+                        boxes[i][2] = 0
+                    if boxes[i][3]<0:
+                        boxes[i][3] = 0
+                for i in range(len(labels)):
+                    labels[i] = 99999
+                boxes_fine, labels_fine, scores_fine = fine_grained.inference(image_or, boxes, labels, scores)
+                t4 = time.time()
+                if gp.DebugPrint_flag:
+                    print("Detection final is ", (t4 - t3) * 1000)
+            else:
+                boxes_fine, labels_fine, scores_fine = [],[],[]
+                if gp.DebugPrint_flag:
+                    print("Detection is nothong")
+            #numb = 2
+            #fine_grained.inference(image_or,boxes,labels,scores)
+            # for box in boxes:
+            #     img1 = image_or.copy()[int(box[1]):int(box[3]), int(box[0]):int(box[2]), :]
+            #     cv2.imwrite(rootpath + "/" + "Detection" + "/"+"{}.jpg".format(numb), img1)
+            #     numb+=1
+            #drawn_image = draw_boxes(image_or, boxes_fine, labels_fine, scores_fine).astype(np.uint8)
+            #labels_fine = labels_fine.numpy()
+            t6 = time.time()
+            xml_outputpath = rootpath + "/" + "Detection"+"/" + index+".xml"
+            make_xml(str(index),image_or.shape[0],image_or.shape[1],boxes_fine,labels_fine,scores_fine,xml_outputpath)
+            #Image.fromarray(drawn_image).save(xml_outputpath.replace(".xml",".jpg"))
+            t5 = time.time()
+            if gp.DebugPrint_flag:
+                print("Detection save is ", (t5 - t6) * 1000)
+    def inference_Red(self,image,ImagePath,fine_grained):
+        with torch.no_grad():
+            image_or = image
+            height, width = image.shape[:2]
+            img_detect = cv2.resize(image, (1024, 768))
+            t1 = time.time()
+            image_scale, image = self.preprocess(img_detect)
+            t2 = time.time()
+            if gp.DebugPrint_flag:
+                print("Detection data deal is ", (t2 - t1) * 1000)
+            pred = self.model(image, augment=False)[0]
+            pred = pred.float()
+            pred = non_max_suppression(pred, self.conf_thres, self.iou_thres)
+            t3 = time.time()
+            if gp.DebugPrint_flag:
+                print("Detection inference is ", (t3 - t2) * 1000)
+            boxes = []
+            scores = []
+            labels = []
+            for det in pred:
+
+                if det is not None and len(det):
+                    det[:, :4] = scale_coords(
+                        image.shape[2:], det[:, :4], image_scale.shape).round()
+
+                    boxes.append(det[:, :4])
+                    scores.append(det[:, 4])
+                    labels.append((det[:, -1]))
+                    boxes = boxes[0]
+                    scores = scores[0]
+                    labels = labels[0]
+            t5 = time.time()
+            if gp.DebugPrint_flag:
+                print("Detection get result is ", (t5 - t3) * 1000)
+            if len(labels) > 0:
+
+                indices_label_left = labels < 4
+                indices_label_right = labels > 2
+                indices = []
+                for i in range(len(indices_label_left)):
+                    indices.append(indices_label_left[i] and indices_label_right[i])
+
+                boxes = boxes[indices]
+                labels = labels[indices]
+                scores = scores[indices]  # 因为莫名奇妙多了一个list维度，所以加了这个0
+                boxes = boxes[0]
+                labels = labels[0]
+                scores = scores[0]
+                ratio_x = (width / 1024.0)
+                ratio_y = (height / 768.0)
+
+                for i in range(len(boxes)):
+                    boxes[i][0] = boxes[i][0] * ratio_x
+                    boxes[i][2] = boxes[i][2] * ratio_x
+                    boxes[i][1] = boxes[i][1] * ratio_y
+                    boxes[i][3] = boxes[i][3] * ratio_y
+                    if boxes[i][0] < 0:
+                        boxes[i][0] = 0
+                    if boxes[i][0] > 3840:
+                        boxes[i][0] = 3839
+                    if boxes[i][1] < 0:
+                        boxes[i][1] = 0
+                    if boxes[i][1] > 2160:
+                        boxes[i][1] = 2159
+                    if boxes[i][2] < 0:
+                        boxes[i][2] = 0
+                    if boxes[i][2]> 3840:
+                        boxes[i][2] = 3839
+                    if boxes[i][3] < 0:
+                        boxes[i][3] = 0
+                    if boxes[i][3] > 2160:
+                        boxes[i][3] = 2159
+                for i in range(len(labels)):
+                    labels[i] = 99999
+                boxes_fine, labels_fine, scores_fine = fine_grained.inference(image_or, boxes, labels, scores)
+                t4 = time.time()
+                if gp.DebugPrint_flag:
+                    print("Detection final is ", (t4 - t3) * 1000)
+                t6 = time.time()
+                xml_outputpath = ImagePath.replace("/ftp/", "/ftp_Result/").replace(".JPG", ".xml").replace(".jpg",
+                                                                                                            ".xml")
+                make_xml(str(ImagePath), image_or.shape[0], image_or.shape[1], boxes_fine, labels_fine, scores_fine,
+                         xml_outputpath)
+                # Image.fromarray(drawn_image).save(xml_outputpath.replace(".xml",".jpg"))
+                t5 = time.time()
+                if gp.DebugPrint_flag:
+                    print("Detection save is ", (t5 - t6) * 1000)
+            else:
+                boxes_fine, labels_fine, scores_fine = [], [], []
+                if gp.DebugPrint_flag:
+                    print("Detection is nothong")
+
+
+    def check_object(self,img_o):
+        flag = False
+        with torch.no_grad():
+            # path = '/data/Hatcher/ouput1/0001_20210414151147/RawImage/0099.jpg'
+            # img = cv2.imread(path)
+            # img = cv2.resize()
+            height, width = img_o.shape[:2]
+            img = cv2.resize(img_o,(1024,768))
+            t1 = time.time()
+            image_scale, image = self.preprocess(img)
+            t2 = time.time()
+            if gp.DebugPrint_flag:
+                print("Check data deal is ", (t2 - t1) * 1000)
+            pred = self.model(image, augment=False)[0]
+            pred = pred.float()
+            pred = non_max_suppression(pred, self.conf_thres, self.iou_thres)
+            t3 = time.time()
+            if gp.DebugPrint_flag:
+                print("Check inference is ", (t3 - t2) * 1000)
+            boxes = []
+            scores = []
+            labels = []
+            for det in pred:
+
+                if det is not None and len(det):
+                    det[:, :4] = scale_coords(
+                        image.shape[2:], det[:, :4], image_scale.shape).round()
+
+                    boxes.append(det[:, :4])
+                    scores.append(det[:, 4])
+                    labels.append((det[:, -1]))
+                    boxes = boxes[0]
+                    scores = scores[0]
+                    labels = labels[0]
+            t5 = time.time()
+            if gp.DebugPrint_flag:
+                print("Check get result is ", (t5 - t3) * 1000)
+
+            if len(labels) > 0:
+                indices_label_left = labels < 4
+                indices_label_right = labels > 2
+                indices = []
+                for i in range(len(indices_label_left)):
+                    indices.append(indices_label_left[i] and indices_label_right[i])
+
+                boxes = boxes[indices]
+                labels = labels[indices]
+                scores = scores[indices]  # 因为莫名奇妙多了一个list维度，所以加了这个0
+                boxes = boxes[0]
+                labels = labels[0]
+                scores = scores[0]
+                ratio_x = (width / 1024.0)
+                ratio_y = (height / 768.0)
+                for i in range(len(boxes)):
+                    boxes[i][0] = boxes[i][0] * ratio_x
+                    boxes[i][2] = boxes[i][2] * ratio_x
+                    boxes[i][1] = boxes[i][1] * ratio_y
+                    boxes[i][3] = boxes[i][3] * ratio_y
+                    if boxes[i][0] < 0:
+                        boxes[i][0] = 0
+                    if boxes[i][0] > 3840:
+                        boxes[i][0] = 3839
+                    if boxes[i][1] < 0:
+                        boxes[i][1] = 0
+                    if boxes[i][1] > 2160:
+                        boxes[i][1] = 2159
+                    if boxes[i][2] < 0:
+                        boxes[i][2] = 0
+                    if boxes[i][2] > 3840:
+                        boxes[i][2] = 3839
+                    if boxes[i][3] < 0:
+                        boxes[i][3] = 0
+                    if boxes[i][3] > 2160:
+                        boxes[i][3] = 2159
+            else:
+                pass
+            if 3 in labels:
+                flag = True
+            if gp.DebugPrint_flag:
+                print("boxes = ",boxes)
+                print("check flag is  ",flag)
+        return flag,labels,boxes,scores
+"""
+细粒度识别
+"""
+class args_temp():
+    dataset = 'HCData'
+    backbone = 'resnet50'
+    batch_size = 1
+    num_workers = 8
+    version = 'val'
+    resize_resolution = 512
+    crop_resolution = 448
+    save_suffix = None
+    swap_num = [7,7]
+class Fine_grained:
+    def __init__(self,device,root):
+        self.device = device
+        self.root = root
+        self.DclPreWeightPath = '../data_MPKFNet/weigth/FinedGrainedRecognition/Test/weights_19_499_0.9998_1.0000.pth'
+        # parser = argparse.ArgumentParser(description='dcl parameters')
+        # parser.add_argument('--data', dest='dataset', default='HCData', type=str)
+        # parser.add_argument('--backbone', dest='backbone', default='resnet50', type=str)
+        # parser.add_argument('--b', dest='batch_size', default=1, type=int)
+        # parser.add_argument('--nw', dest='num_workers', default=8, type=int)
+        # parser.add_argument('--ver', dest='version', default='val', type=str)
+        # parser.add_argument('--size', dest='resize_resolution', default=512, type=int)
+        # parser.add_argument('--crop', dest='crop_resolution', default=448, type=int)
+        # parser.add_argument('--ss', dest='save_suffix', default=None, type=str)
+        # parser.add_argument('--acc_report', dest='acc_report', action='store_true')
+        # parser.add_argument('--swap_num', default=[7, 7], nargs=2, metavar=('swap1', 'swap2'), type=int,
+        #                     help='specify a range')
+        self.args = args_temp()
+        self.Config = LoadConfig(self.args, 's')
+        self.transformers = load_data_transformers(self.args.resize_resolution, self.args.crop_resolution, self.args.swap_num)
+        self.net = MainModel(self.Config).to(device)
+        model_dict = self.net.state_dict()
+        pretrained_dict = torch.load(self.DclPreWeightPath)
+        model_dict.update(pretrained_dict)
+        self.net.load_state_dict(model_dict)
+        if gp.DebugPrint_flag:
+            print("load model weight is ok")
+        self.net.train(False)
+    def inference(self,img,imagelist,labellist,scorelist):
+        with torch.no_grad():
+            for i,box in enumerate(imagelist):
+                Img = img.copy()[int(box[1]):int(box[3]), int(box[0]):int(box[2]), :]
+                if (int(box[3])-int(box[1])>100) and (int(box[2])-int(box[0])>100):
+                    t1 = time.time()
+                    Img = Image.fromarray(cv2.cvtColor(Img, cv2.COLOR_BGR2RGB))
+                    Img = self.transformers['test_totensor'](Img)
+                    Img = Img.unsqueeze(0)
+                    inputs = Variable(Img.cuda())
+                    t2 = time.time()
+                    if gp.DebugPrint_flag:
+                        print("fine data deal is ", (t2 - t1) * 1000)
+                    outputs = self.net(inputs)
+                    t3 = time.time()
+                    if gp.DebugPrint_flag:
+                        print("fine inference is ", (t3 - t2) * 1000)
+                    top3_val, top3_pos = torch.topk(outputs[0], 3)  # topk返回的是一个元组，第一个元素是每一行最大的三个值（默认是按行取最大值），第二个元素是每一行最大的值的排序
+                    Species = top3_pos.cpu().numpy()[0][0]
+                    ConfidenceLevel = top3_val.cpu().numpy()[0][0]
+                    ConfidenceLevel = ConfidenceLevel/26.0
+                    if ConfidenceLevel >1:
+                        ConfidenceLevel = 1.0* random.uniform(0, 1)
+                    if gp.DebugPrint_flag:
+                        print("hhhhhhhhhhhhhhhhhhhhhhh is ",ConfidenceLevel)
+                    labellist[i] = Species
+                    scorelist[i] = ConfidenceLevel
+                    t4 = time.time()
+                    if gp.DebugPrint_flag:
+                        print("fine data final is ", (t4 - t3) * 1000)
+
+            return imagelist, labellist, scorelist
+class Check_start (threading.Thread):
+    def __init__(self, threadID, name, img,DetectionClass):
+        threading.Thread.__init__(self)
+        self.threadID = threadID
+        self.name = name
+        self.img = img
+        self.DetectionClass = DetectionClass
+    def run(self):
+        if gp.DebugPrint_flag:
+            print(self.name + " is start")
+        gp.is_ex_bird ,labels, boxes ,scores= self.DetectionClass.check_object(self.img)
+
+        if gp.is_ex_bird :
+            if gp.DebugPrint_flag:
+                print(gp.is_ex_bird)
+            gp.image_time = gp.Camera_device +time.strftime('%Y%m%d%H%M%S',time.localtime(time.time()))
+        if gp.DebugPrint_flag:
+            print(self.name + " is finish")
+class Check_start_fordataselect (threading.Thread):
+    def __init__(self, threadID, name, img,DetectionClass,fine_grained):
+        threading.Thread.__init__(self)
+        self.threadID = threadID
+        self.name = name
+        self.img = img
+        self.DetectionClass = DetectionClass
+        self.fine_grained = fine_grained
+    def run(self):
+        if gp.DebugPrint_flag:
+            print(self.name + " is start")
+        is_ex_bird ,labels, boxes ,scores= self.DetectionClass.check_object(self.img)
+        image_or = self.img.copy()
+        if is_ex_bird :
+            for i in range(len(boxes)):
+                if boxes[i][0] < 0:
+                    boxes[i][0] = 0
+                if boxes[i][1] < 0:
+                    boxes[i][1] = 0
+                if boxes[i][2] < 0:
+                    boxes[i][2] = 0
+                if boxes[i][3] < 0:
+                    boxes[i][3] = 0
+            for i in range(len(labels)):
+                labels[i] = 99999
+            boxes_fine, labels_fine, scores_fine = self.fine_grained.inference(self.img, boxes, labels, scores)
+            for i in range(len(scores)):
+                xyxy = boxes_fine[i]
+                conf = scores_fine[i]
+                cls = labels_fine[i]
+                label = f'{int(cls)} {conf:.2f}'
+                palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1)
+                colors = [int((p * (cls ** 2 - cls + 1)) % 255) for p in palette]
+                plot_one_box(xyxy, self.img, label=label, color=colors, line_thickness=1)
+
+            cv2.imwrite("../data_MPKFNet/RTSP/" + "{}.jpg".format(time.strftime('%Y%m%d%H%M%S', time.localtime(time.time()))),
+                        image_or)
+            cv2.imwrite(
+                "../data_MPKFNet/RTSP/result/" + "{}.jpg".format(time.strftime('%Y%m%d%H%M%S', time.localtime(time.time()))),
+                self.img)
+        if gp.DebugPrint_flag:
+            print(self.name + " is finish")
+
+class Density_mul (threading.Thread):
+    def __init__(self, threadID, name,imagetime,DensityClass):
+        threading.Thread.__init__(self)
+        self.threadID = threadID
+        self.name = name
+
+        self.imagetime = imagetime
+        self.DensityClass = DensityClass
+    def run(self):
+        gp.Density_thread_flag = True
+
+        if gp.DebugPrint_flag:
+            print(self.name +self.imagetime," is starting")
+        i = 0
+        while (i<=29):
+            # if(len(gp.ImgDealQueue)<=i):
+            #     gp.event.wait()
+            # if ((i+1)%25==0):   #说明是第25*n帧,间隔1s处理一次
+            if gp.DebugPrint_flag:
+                print("Density is1 ", len(gp.ImgDealQueue), "  ", i)
+                print("Density is waiting")
+            if (len(gp.ImgDealQueue) <= i):
+                gp.eventDensity.wait()
+            if gp.DebugPrint_flag:
+                print("Density is2 ", len(gp.ImgDealQueue), "  ", i)
+                print("Density is not waiting")
+            self.DensityClass.inference(gp.ImgDealQueue[i],self.imagetime,str(i).zfill(4))
+            gp.eventDensity.clear()
+            i += 1
+        gp.Density_thread_flag = False
+        gp.is_ex_bird = False
+        gp.destory_image = True
+        if gp.DebugPrint_flag:
+            print(self.name + self.imagetime," is over")
+class Detecntion_mul (threading.Thread):
+    def __init__(self, threadID, name,imagetime,DetectionClass,Fine_grainedClass):
+        threading.Thread.__init__(self)
+        self.threadID = threadID
+        self.name = name
+        self.imagetime = imagetime
+        self.DetectionClass = DetectionClass
+        self.Fine_grainedClass = Fine_grainedClass
+    def run(self):
+        gp.Detection_thread_flag = True
+
+        if gp.DebugPrint_flag:
+            print(self.name ,self.imagetime," is starting")
+        i = 0
+        while (i<=29):
+            if gp.DebugPrint_flag:
+                print("Detection is1 ",len(gp.ImgDealQueue),"  ",i)
+            # if (len(gp.ImgDealQueue) <= i):
+            #     gp.event.wait()
+            # if ((i+1)%25==0):   #说明是第25*n帧,间隔1s处理一次
+                print("Detection is waiting")
+            if (len(gp.ImgDealQueue) <= i):
+                gp.eventDetection.wait()
+            if gp.DebugPrint_flag:
+                print("Detection is2 ", len(gp.ImgDealQueue), "  ", i)
+                print("Detection is not waiting")
+            self.DetectionClass.inference(gp.ImgDealQueue[i],self.imagetime,str(i).zfill(4),self.Fine_grainedClass)
+            gp.eventDetection.clear()
+
+            i+=1
+        gp.Detection_thread_flag = False
+        gp.is_ex_bird = False
+        gp.destory_image = True
+        if gp.DebugPrint_flag:
+            print(self.name ,self.imagetime," is over")
+class SaveImage_mul (threading.Thread):
+    def __init__(self, threadID, name,imagetime,root):
+        threading.Thread.__init__(self)
+        self.threadID = threadID
+        self.name = name
+        self.imagetime = imagetime
+        self.root = root
+        self.rootpath = os.path.join(self.root, self.imagetime)
+    def run(self):
+        gp.SaveImage_thread_flag = True
+
+        if not os.path.exists(self.rootpath):
+            os.mkdir(self.rootpath)
+        if not os.path.exists(self.rootpath + "/" + "RawImage" + "/"):
+            os.mkdir(self.rootpath + "/" + "RawImage" + "/")
+        if gp.DebugPrint_flag:
+            print(self.name ,self.imagetime," is starting")
+        i = 0
+        while (i<=29):
+            # if(len(gp.ImgDealQueue)<=i):
+            #     gp.event.wait()
+            # if ((i+1)%25==0):   #说明是第25*n帧,间隔1s处理一次
+            if gp.DebugPrint_flag:
+                print("Saveimage is1 ", len(gp.ImgDealQueue), "  ", i)
+                print("Saveimage is waiting")
+            if (len(gp.ImgDealQueue) <= i):
+                gp.eventSaveimage.wait()
+            if gp.DebugPrint_flag:
+                print("Saveimage is2 ", len(gp.ImgDealQueue), "  ", i)
+                print("Saveimage is not waiting")
+            cv2.imwrite(os.path.join(self.rootpath,"RawImage","{}.jpg".format(str(i).zfill(4))),gp.ImgDealQueue[i])
+            gp.eventSaveimage.clear()
+            if gp.DebugPrint_flag:
+                print("Saveimage finish once")
+            i+=1
+        gp.SaveImage_thread_flag = False
+        gp.is_ex_bird = False
+        gp.destory_image = True
+        if gp.DebugPrint_flag:
+            print(self.name ,self.imagetime," is over")
+class ThreadMonitor_mul (threading.Thread):
+    def __init__(self, threadID, name):
+        threading.Thread.__init__(self)
+        self.threadID = threadID
+        self.name = name
+
+    def run(self):
+        if gp.DebugPrint_flag:
+            print(self.name," is starting")
+        while 1:
+           # print(gp.Density_thread_flag,gp.Detection_thread_flag,gp.SaveImage_thread_flag)
+            if (gp.Density_thread_flag) == False and (gp.Detection_thread_flag == False) and (gp.SaveImage_thread_flag == False):  #说明三个进程都已经结束了
+                gp.ImgDealQueue = []
+class isopenmul (threading.Thread):
+    def __init__(self, threadID, name,DetectionClass):
+        threading.Thread.__init__(self)
+        self.threadID = threadID
+        self.name = name
+        self.DetectionClass = DetectionClass
+    def run(self):
+        while 1:
+           # print(gp.Density_thread_flag,gp.Detection_thread_flag,gp.SaveImage_thread_flag)
+           if (gp.Vedio_thread_flag == False) and (gp.Density_thread_flag == False) and (
+                   gp.Detection_thread_flag == False) and (gp.SaveImage_thread_flag == False) and (
+                   gp.is_ex_bird == False) and gp.num_detection > 24:
+               gp.num_detection = 0
+
+               if gp.DebugPrint_flag:
+                   print( "check is start")
+               gp.is_ex_bird, labels, boxes, scores = self.DetectionClass.check_object(gp.frame_openmul)
+
+               if gp.is_ex_bird:
+                   if gp.DebugPrint_flag:
+                       print(gp.is_ex_bird)
+                   gp.image_time = gp.Camera_device + time.strftime('%Y%m%d%H%M%S', time.localtime(time.time()))
+               if gp.DebugPrint_flag:
+                   print("check is finish")
+
+               # Check_start1 = Check_start("0002", "check_start",  gp.frame_openmul, self.DetectionClass)
+               # Check_start1.start()
+
+
+
+
+
+
+
+
+
+
+
+
+
+if __name__ == '__main__':
+    t1 = time.time()
+    image = cv2.imread("/home/wanghc/output1/0001_20210413155842/RawImage/0174.jpg")
+    t2 = time.time()
+    camera_num = "0001"
+    image_time = time.strftime('%Y%m%d%H%M%S',time.localtime(time.time()))
+    image_time = camera_num + image_time
+    root = "/home/wanghc/output"
+    Density1 = Density(gp.devices,root)
+    Detection1 = Detection_Y5(gp.devices,root)
+    Fine_grained1 = Fine_grained(gp.devices,root)
+    os.environ['CUDA_VISIBLE_DEVICES'] = '4'
+    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    path = '/home/wanghc/output1/001020210509163621/RawImage/0000.jpg'
+    img = cv2.imread(path)  # BGR
+    Detection1.inference(img, image_time, str(1), Fine_grained1)
+
+    # root = "/home/wanghc/output"
+    # detector = Detection_Y5(gp.devices, root)
+    # print("ok")
+    # path = '/home/wanghc/output1/001020210509163621/RawImage/0000.jpg'
+    # img = cv2.imread(path)  # BGR
+    # img_show = img#cv2.resize(img, (1024, 768))
+    # names = detector.model.module.names if hasattr(detector.model, 'module') else detector.model.names
+    # colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
+    # flag, labels,pred_boxes, scores  = detector.check_object(img_show)
+    # print(len(scores))
+    # #img_show = cv2.resize(img, (768, 1024))
+    # for i in range(len(scores)):
+    #     xyxy = pred_boxes[i]
+    #     conf = scores[i]
+    #     cls = labels[i]
+    #     label = f'{names[int(cls)]} {conf:.2f}'
+    #     plot_one_box(xyxy, img_show, label=label, color=colors[int(cls)], line_thickness=1)
+    # cv2.imshow(str(path), img_show)
+    # t1 = time.time()
+    # cv2.imwrite("./a.jpg", img_show)
+    # t2 = time.time()
+    # print((t2 - t1) * 1000)
+    # cv2.waitKey(-1)
+    # while(1):
+    #     pass
+
+
+
+
+
+
+
+

Y5/data/argoverse_hd.yaml

@@ -0,0 +1,21 @@
+# Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/
+# Train command: python train.py --data argoverse_hd.yaml
+# Default dataset location is next to /yolov5:
+#   /parent_folder
+#     /argoverse
+#     /yolov5
+
+
+# download command/URL (optional)
+download: bash data/scripts/get_argoverse_hd.sh
+
+# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
+train: ../argoverse/Argoverse-1.1/images/train/  # 39384 images
+val: ../argoverse/Argoverse-1.1/images/val/  # 15062 iamges
+test: ../argoverse/Argoverse-1.1/images/test/  # Submit to: https://eval.ai/web/challenges/challenge-page/800/overview
+
+# number of classes
+nc: 8
+
+# class names
+names: [ 'person',  'bicycle',  'car',  'motorcycle',  'bus',  'truck',  'traffic_light',  'stop_sign' ]

Y5/data/coco.yaml

@@ -0,0 +1,35 @@
+# COCO 2017 dataset http://cocodataset.org
+# Train command: python train.py --data coco.yaml
+# Default dataset location is next to /yolov5:
+#   /parent_folder
+#     /coco
+#     /yolov5
+
+
+# download command/URL (optional)
+download: bash data/scripts/get_coco.sh
+
+# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
+train: ../coco/train2017.txt  # 118287 images
+val: ../coco/val2017.txt  # 5000 images
+test: ../coco/test-dev2017.txt  # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794
+
+# number of classes
+nc: 80
+
+# class names
+names: [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', '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', 'couch',
+         'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
+         'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
+         'hair drier', 'toothbrush' ]
+
+# Print classes
+# with open('data/coco.yaml') as f:
+#   d = yaml.load(f, Loader=yaml.FullLoader)  # dict
+#   for i, x in enumerate(d['names']):
+#     print(i, x)

Y5/data/coco128.yaml

@@ -0,0 +1,28 @@
+# COCO 2017 dataset http://cocodataset.org - first 128 training images
+# Train command: python train.py --data coco128.yaml
+# Default dataset location is next to /yolov5:
+#   /parent_folder
+#     /coco128
+#     /yolov5
+
+
+# download command/URL (optional)
+download: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip
+
+# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
+train: ../coco128/images/train2017/  # 128 images
+val: ../coco128/images/train2017/  # 128 images
+
+# number of classes
+nc: 80
+
+# class names
+names: [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', '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', 'couch',
+         'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
+         'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
+         'hair drier', 'toothbrush' ]

Y5/data/hyp.finetune.yaml

@@ -0,0 +1,38 @@
+# Hyperparameters for VOC finetuning
+# python train.py --batch 64 --weights yolov5m.pt --data voc.yaml --img 512 --epochs 50
+# See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials
+
+
+# Hyperparameter Evolution Results
+# Generations: 306
+#                   P         R     mAP.5 mAP.5:.95       box       obj       cls
+# Metrics:        0.6     0.936     0.896     0.684    0.0115   0.00805   0.00146
+
+lr0: 0.0032
+lrf: 0.12
+momentum: 0.843
+weight_decay: 0.00036
+warmup_epochs: 2.0
+warmup_momentum: 0.5
+warmup_bias_lr: 0.05
+box: 0.0296
+cls: 0.243
+cls_pw: 0.631
+obj: 0.301
+obj_pw: 0.911
+iou_t: 0.2
+anchor_t: 2.91
+# anchors: 3.63
+fl_gamma: 0.0
+hsv_h: 0.0138
+hsv_s: 0.664
+hsv_v: 0.464
+degrees: 0.373
+translate: 0.245
+scale: 0.898
+shear: 0.602
+perspective: 0.0
+flipud: 0.00856
+fliplr: 0.5
+mosaic: 1.0
+mixup: 0.243

Y5/data/hyp.scratch.yaml

@@ -0,0 +1,33 @@
+# Hyperparameters for COCO training from scratch
+# python train.py --batch 40 --cfg yolov5m.yaml --weights '' --data coco.yaml --img 640 --epochs 300
+# See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials
+
+
+lr0: 0.01  # initial learning rate (SGD=1E-2, Adam=1E-3)
+lrf: 0.2  # final OneCycleLR learning rate (lr0 * lrf)
+momentum: 0.937  # SGD momentum/Adam beta1
+weight_decay: 0.0005  # optimizer weight decay 5e-4
+warmup_epochs: 3.0  # warmup epochs (fractions ok)
+warmup_momentum: 0.8  # warmup initial momentum
+warmup_bias_lr: 0.1  # warmup initial bias lr
+box: 0.05  # box loss gain
+cls: 0.5  # cls loss gain
+cls_pw: 1.0  # cls BCELoss positive_weight
+obj: 1.0  # obj loss gain (scale with pixels)
+obj_pw: 1.0  # obj BCELoss positive_weight
+iou_t: 0.20  # IoU training threshold
+anchor_t: 4.0  # anchor-multiple threshold
+# anchors: 3  # anchors per output layer (0 to ignore)
+fl_gamma: 0.0  # focal loss gamma (efficientDet default gamma=1.5)
+hsv_h: 0.015  # image HSV-Hue augmentation (fraction)
+hsv_s: 0.7  # image HSV-Saturation augmentation (fraction)
+hsv_v: 0.4  # image HSV-Value augmentation (fraction)
+degrees: 0.0  # image rotation (+/- deg)
+translate: 0.1  # image translation (+/- fraction)
+scale: 0.5  # image scale (+/- gain)
+shear: 0.0  # image shear (+/- deg)
+perspective: 0.0  # image perspective (+/- fraction), range 0-0.001
+flipud: 0.0  # image flip up-down (probability)
+fliplr: 0.5  # image flip left-right (probability)
+mosaic: 1.0  # image mosaic (probability)
+mixup: 0.0  # image mixup (probability)

Y5/data/scripts/get_argoverse_hd.sh

@@ -0,0 +1,62 @@
+#!/bin/bash
+# Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/
+# Download command: bash data/scripts/get_argoverse_hd.sh
+# Train command: python train.py --data argoverse_hd.yaml
+# Default dataset location is next to /yolov5:
+#   /parent_folder
+#     /argoverse
+#     /yolov5
+
+# Download/unzip images
+d='../argoverse/' # unzip directory
+mkdir $d
+url=https://argoverse-hd.s3.us-east-2.amazonaws.com/
+f=Argoverse-HD-Full.zip
+wget $url$f -O $f && unzip -q $f -d $d && rm $f &# download, unzip, remove in background
+wait                                              # finish background tasks
+
+cd ../argoverse/Argoverse-1.1/
+ln -s tracking images
+
+cd ../Argoverse-HD/annotations/
+
+python3 - "$@" <<END
+import json
+from pathlib import Path
+
+annotation_files = ["train.json", "val.json"]
+print("Converting annotations to YOLOv5 format...")
+
+for val in annotation_files:
+    a = json.load(open(val, "rb"))
+
+    label_dict = {}
+    for annot in a['annotations']:
+        img_id = annot['image_id']
+        img_name = a['images'][img_id]['name']
+        img_label_name = img_name[:-3] + "txt"
+
+        obj_class = annot['category_id']
+        x_center, y_center, width, height = annot['bbox']
+        x_center = (x_center + width / 2) / 1920.  # offset and scale
+        y_center = (y_center + height / 2) / 1200.  # offset and scale
+        width /= 1920.  # scale
+        height /= 1200.  # scale
+
+        img_dir = "./labels/" + a['seq_dirs'][a['images'][annot['image_id']]['sid']]
+
+        Path(img_dir).mkdir(parents=True, exist_ok=True)
+
+        if img_dir + "/" + img_label_name not in label_dict:
+            label_dict[img_dir + "/" + img_label_name] = []
+
+        label_dict[img_dir + "/" + img_label_name].append(f"{obj_class} {x_center} {y_center} {width} {height}\n")
+
+    for filename in label_dict:
+        with open(filename, "w") as file:
+            for string in label_dict[filename]:
+                file.write(string)
+
+END
+
+mv ./labels ../../Argoverse-1.1/

Y5/data/scripts/get_coco.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+# COCO 2017 dataset http://cocodataset.org
+# Download command: bash data/scripts/get_coco.sh
+# Train command: python train.py --data coco.yaml
+# Default dataset location is next to /yolov5:
+#   /parent_folder
+#     /coco
+#     /yolov5
+
+# Download/unzip labels
+d='../' # unzip directory
+url=https://github.com/ultralytics/yolov5/releases/download/v1.0/
+f='coco2017labels.zip' # or 'coco2017labels-segments.zip', 68 MB
+echo 'Downloading' $url$f ' ...'
+curl -L $url$f -o $f && unzip -q $f -d $d && rm $f & # download, unzip, remove in background
+
+# Download/unzip images
+d='../coco/images' # unzip directory
+url=http://images.cocodataset.org/zips/
+f1='train2017.zip' # 19G, 118k images
+f2='val2017.zip'   # 1G, 5k images
+f3='test2017.zip'  # 7G, 41k images (optional)
+for f in $f1 $f2; do
+  echo 'Downloading' $url$f '...'
+  curl -L $url$f -o $f && unzip -q $f -d $d && rm $f & # download, unzip, remove in background
+done
+wait # finish background tasks

Y5/data/scripts/get_voc.sh

@@ -0,0 +1,139 @@
+#!/bin/bash
+# PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC/
+# Download command: bash data/scripts/get_voc.sh
+# Train command: python train.py --data voc.yaml
+# Default dataset location is next to /yolov5:
+#   /parent_folder
+#     /VOC
+#     /yolov5
+
+start=$(date +%s)
+mkdir -p ../tmp
+cd ../tmp/
+
+# Download/unzip images and labels
+d='.' # unzip directory
+url=https://github.com/ultralytics/yolov5/releases/download/v1.0/
+f1=VOCtrainval_06-Nov-2007.zip # 446MB, 5012 images
+f2=VOCtest_06-Nov-2007.zip     # 438MB, 4953 images
+f3=VOCtrainval_11-May-2012.zip # 1.95GB, 17126 images
+for f in $f3 $f2 $f1; do
+  echo 'Downloading' $url$f '...' 
+  curl -L $url$f -o $f && unzip -q $f -d $d && rm $f & # download, unzip, remove in background
+done
+wait # finish background tasks
+
+end=$(date +%s)
+runtime=$((end - start))
+echo "Completed in" $runtime "seconds"
+
+echo "Splitting dataset..."
+python3 - "$@" <<END
+import xml.etree.ElementTree as ET
+import pickle
+import os
+from os import listdir, getcwd
+from os.path import join
+
+sets=[('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
+
+classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
+
+
+def convert(size, box):
+    dw = 1./(size[0])
+    dh = 1./(size[1])
+    x = (box[0] + box[1])/2.0 - 1
+    y = (box[2] + box[3])/2.0 - 1
+    w = box[1] - box[0]
+    h = box[3] - box[2]
+    x = x*dw
+    w = w*dw
+    y = y*dh
+    h = h*dh
+    return (x,y,w,h)
+
+def convert_annotation(year, image_id):
+    in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml'%(year, image_id))
+    out_file = open('VOCdevkit/VOC%s/labels/%s.txt'%(year, image_id), 'w')
+    tree=ET.parse(in_file)
+    root = tree.getroot()
+    size = root.find('size')
+    w = int(size.find('width').text)
+    h = int(size.find('height').text)
+
+    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 = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
+        bb = convert((w,h), b)
+        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
+
+wd = getcwd()
+
+for year, image_set in sets:
+    if not os.path.exists('VOCdevkit/VOC%s/labels/'%(year)):
+        os.makedirs('VOCdevkit/VOC%s/labels/'%(year))
+    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:
+        list_file.write('%s/VOCdevkit/VOC%s/JPEGImages/%s.jpg\n'%(wd, year, image_id))
+        convert_annotation(year, image_id)
+    list_file.close()
+
+END
+
+cat 2007_train.txt 2007_val.txt 2012_train.txt 2012_val.txt >train.txt
+cat 2007_train.txt 2007_val.txt 2007_test.txt 2012_train.txt 2012_val.txt >train.all.txt
+
+python3 - "$@" <<END
+
+import shutil
+import os
+os.system('mkdir ../VOC/')
+os.system('mkdir ../VOC/images')
+os.system('mkdir ../VOC/images/train')
+os.system('mkdir ../VOC/images/val')
+
+os.system('mkdir ../VOC/labels')
+os.system('mkdir ../VOC/labels/train')
+os.system('mkdir ../VOC/labels/val')
+
+import os
+print(os.path.exists('../tmp/train.txt'))
+f = open('../tmp/train.txt', 'r')
+lines = f.readlines()
+
+for line in lines:
+    line = "/".join(line.split('/')[-5:]).strip()
+    if (os.path.exists("../" + line)):
+        os.system("cp ../"+ line + " ../VOC/images/train")
+        
+    line = line.replace('JPEGImages', 'labels')
+    line = line.replace('jpg', 'txt')
+    if (os.path.exists("../" + line)):
+        os.system("cp ../"+ line + " ../VOC/labels/train")
+
+
+print(os.path.exists('../tmp/2007_test.txt'))
+f = open('../tmp/2007_test.txt', 'r')
+lines = f.readlines()
+
+for line in lines:
+    line = "/".join(line.split('/')[-5:]).strip()
+    if (os.path.exists("../" + line)):
+        os.system("cp ../"+ line + " ../VOC/images/val")
+        
+    line = line.replace('JPEGImages', 'labels')
+    line = line.replace('jpg', 'txt')
+    if (os.path.exists("../" + line)):
+        os.system("cp ../"+ line + " ../VOC/labels/val")
+
+END
+
+rm -rf ../tmp # remove temporary directory
+echo "VOC download done."

Y5/data/voc.yaml

@@ -0,0 +1,21 @@
+# PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC/
+# Train command: python train.py --data voc.yaml
+# Default dataset location is next to /yolov5:
+#   /parent_folder
+#     /VOC
+#     /yolov5
+
+
+# download command/URL (optional)
+#download: bash data/scripts/get_voc.sh
+
+# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
+train: /data1/Hatcher/model/data_MPKFNet/data/VOC/VOCdevkit/VOC6/images/train/  # 16551 images
+val: /data1/Hatcher/model/data_MPKFNet/data/VOC/VOCdevkit/VOC6/images/test/  # 4952 images
+
+# number of classes
+nc: 20
+
+# class names
+names: [ 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
+         'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor' ]

Y5/models/common.py

@@ -0,0 +1,309 @@
+# This file contains modules common to various models
+
+import math
+from pathlib import Path
+import sys
+sys.path.append("..")
+import numpy as np
+import requests
+import torch
+import torch.nn as nn
+from PIL import Image
+
+from Y5.utils.datasets import letterbox
+from Y5.utils.general import non_max_suppression, make_divisible, scale_coords, xyxy2xywh
+from Y5.utils.plots import color_list, plot_one_box
+
+
+def autopad(k, p=None):  # kernel, padding
+    # Pad to 'same'
+    if p is None:
+        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
+    return p
+
+
+def DWConv(c1, c2, k=1, s=1, act=True):
+    # Depthwise convolution
+    return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act)
+
+
+class Conv(nn.Module):
+    # Standard convolution
+    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
+        super(Conv, self).__init__()
+        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
+        self.bn = nn.BatchNorm2d(c2)
+        self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
+
+    def forward(self, x):
+        return self.act(self.bn(self.conv(x)))
+
+    def fuseforward(self, x):
+        return self.act(self.conv(x))
+
+
+class Bottleneck(nn.Module):
+    # Standard bottleneck
+    def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, shortcut, groups, expansion
+        super(Bottleneck, self).__init__()
+        c_ = int(c2 * e)  # hidden channels
+        self.cv1 = Conv(c1, c_, 1, 1)
+        self.cv2 = Conv(c_, c2, 3, 1, g=g)
+        self.add = shortcut and c1 == c2
+
+    def forward(self, x):
+        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
+
+
+class BottleneckCSP(nn.Module):
+    # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
+    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
+        super(BottleneckCSP, self).__init__()
+        c_ = int(c2 * e)  # hidden channels
+        self.cv1 = Conv(c1, c_, 1, 1)
+        self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
+        self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
+        self.cv4 = Conv(2 * c_, c2, 1, 1)
+        self.bn = nn.BatchNorm2d(2 * c_)  # applied to cat(cv2, cv3)
+        self.act = nn.LeakyReLU(0.1, inplace=True)
+        self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
+
+    def forward(self, x):
+        y1 = self.cv3(self.m(self.cv1(x)))
+        y2 = self.cv2(x)
+        return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1))))
+
+
+class C3(nn.Module):
+    # CSP Bottleneck with 3 convolutions
+    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
+        super(C3, self).__init__()
+        c_ = int(c2 * e)  # hidden channels
+        self.cv1 = Conv(c1, c_, 1, 1)
+        self.cv2 = Conv(c1, c_, 1, 1)
+        self.cv3 = Conv(2 * c_, c2, 1)  # act=FReLU(c2)
+        self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
+        # self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)])
+
+    def forward(self, x):
+        return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1))
+
+
+class SPP(nn.Module):
+    # Spatial pyramid pooling layer used in YOLOv3-SPP
+    def __init__(self, c1, c2, k=(5, 9, 13)):
+        super(SPP, self).__init__()
+        c_ = c1 // 2  # hidden channels
+        self.cv1 = Conv(c1, c_, 1, 1)
+        self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
+        self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
+
+    def forward(self, x):
+        x = self.cv1(x)
+        return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
+
+
+class Focus(nn.Module):
+    # Focus wh information into c-space
+    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
+        super(Focus, self).__init__()
+        self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
+        # self.contract = Contract(gain=2)
+
+    def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
+        return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
+        # return self.conv(self.contract(x))
+
+
+class Contract(nn.Module):
+    # Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40)
+    def __init__(self, gain=2):
+        super().__init__()
+        self.gain = gain
+
+    def forward(self, x):
+        N, C, H, W = x.size()  # assert (H / s == 0) and (W / s == 0), 'Indivisible gain'
+        s = self.gain
+        x = x.view(N, C, H // s, s, W // s, s)  # x(1,64,40,2,40,2)
+        x = x.permute(0, 3, 5, 1, 2, 4).contiguous()  # x(1,2,2,64,40,40)
+        return x.view(N, C * s * s, H // s, W // s)  # x(1,256,40,40)
+
+
+class Expand(nn.Module):
+    # Expand channels into width-height, i.e. x(1,64,80,80) to x(1,16,160,160)
+    def __init__(self, gain=2):
+        super().__init__()
+        self.gain = gain
+
+    def forward(self, x):
+        N, C, H, W = x.size()  # assert C / s ** 2 == 0, 'Indivisible gain'
+        s = self.gain
+        x = x.view(N, s, s, C // s ** 2, H, W)  # x(1,2,2,16,80,80)
+        x = x.permute(0, 3, 4, 1, 5, 2).contiguous()  # x(1,16,80,2,80,2)
+        return x.view(N, C // s ** 2, H * s, W * s)  # x(1,16,160,160)
+
+
+class Concat(nn.Module):
+    # Concatenate a list of tensors along dimension
+    def __init__(self, dimension=1):
+        super(Concat, self).__init__()
+        self.d = dimension
+
+    def forward(self, x):
+        return torch.cat(x, self.d)
+
+
+class NMS(nn.Module):
+    # Non-Maximum Suppression (NMS) module
+    conf = 0.25  # confidence threshold
+    iou = 0.45  # IoU threshold
+    classes = None  # (optional list) filter by class
+
+    def __init__(self):
+        super(NMS, self).__init__()
+
+    def forward(self, x):
+        return non_max_suppression(x[0], conf_thres=self.conf, iou_thres=self.iou, classes=self.classes)
+
+
+class autoShape(nn.Module):
+    # input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
+    img_size = 640  # inference size (pixels)
+    conf = 0.25  # NMS confidence threshold
+    iou = 0.45  # NMS IoU threshold
+    classes = None  # (optional list) filter by class
+
+    def __init__(self, model):
+        super(autoShape, self).__init__()
+        self.model = model.eval()
+
+    def autoshape(self):
+        print('autoShape already enabled, skipping... ')  # model already converted to model.autoshape()
+        return self
+
+    def forward(self, imgs, size=640, augment=False, profile=False):
+        # Inference from various sources. For height=720, width=1280, RGB images example inputs are:
+        #   filename:   imgs = 'data/samples/zidane.jpg'
+        #   URI:             = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg'
+        #   OpenCV:          = cv2.imread('image.jpg')[:,:,::-1]  # HWC BGR to RGB x(720,1280,3)
+        #   PIL:             = Image.open('image.jpg')  # HWC x(720,1280,3)
+        #   numpy:           = np.zeros((720,1280,3))  # HWC
+        #   torch:           = torch.zeros(16,3,720,1280)  # BCHW
+        #   multiple:        = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...]  # list of images
+
+        p = next(self.model.parameters())  # for device and type
+        if isinstance(imgs, torch.Tensor):  # torch
+            return self.model(imgs.to(p.device).type_as(p), augment, profile)  # inference
+
+        # Pre-process
+        n, imgs = (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs])  # number of images, list of images
+        shape0, shape1, files = [], [], []  # image and inference shapes, filenames
+        for i, im in enumerate(imgs):
+            if isinstance(im, str):  # filename or uri
+                im, f = Image.open(requests.get(im, stream=True).raw if im.startswith('http') else im), im  # open
+                im.filename = f  # for uri
+            files.append(Path(im.filename).with_suffix('.jpg').name if isinstance(im, Image.Image) else f'image{i}.jpg')
+            im = np.array(im)  # to numpy
+            if im.shape[0] < 5:  # image in CHW
+                im = im.transpose((1, 2, 0))  # reverse dataloader .transpose(2, 0, 1)
+            im = im[:, :, :3] if im.ndim == 3 else np.tile(im[:, :, None], 3)  # enforce 3ch input
+            s = im.shape[:2]  # HWC
+            shape0.append(s)  # image shape
+            g = (size / max(s))  # gain
+            shape1.append([y * g for y in s])
+            imgs[i] = im  # update
+        shape1 = [make_divisible(x, int(self.stride.max())) for x in np.stack(shape1, 0).max(0)]  # inference shape
+        x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs]  # pad
+        x = np.stack(x, 0) if n > 1 else x[0][None]  # stack
+        x = np.ascontiguousarray(x.transpose((0, 3, 1, 2)))  # BHWC to BCHW
+        x = torch.from_numpy(x).to(p.device).type_as(p) / 255.  # uint8 to fp16/32
+
+        # Inference
+        with torch.no_grad():
+            y = self.model(x, augment, profile)[0]  # forward
+        y = non_max_suppression(y, conf_thres=self.conf, iou_thres=self.iou, classes=self.classes)  # NMS
+
+        # Post-process
+        for i in range(n):
+            scale_coords(shape1, y[i][:, :4], shape0[i])
+
+        return Detections(imgs, y, files, self.names)
+
+
+class Detections:
+    # detections class for YOLOv5 inference results
+    def __init__(self, imgs, pred, files, names=None):
+        super(Detections, self).__init__()
+        d = pred[0].device  # device
+        gn = [torch.tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.], device=d) for im in imgs]  # normalizations
+        self.imgs = imgs  # list of images as numpy arrays
+        self.pred = pred  # list of tensors pred[0] = (xyxy, conf, cls)
+        self.names = names  # class names
+        self.files = files  # image filenames
+        self.xyxy = pred  # xyxy pixels
+        self.xywh = [xyxy2xywh(x) for x in pred]  # xywh pixels
+        self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)]  # xyxy normalized
+        self.xywhn = [x / g for x, g in zip(self.xywh, gn)]  # xywh normalized
+        self.n = len(self.pred)
+
+    def display(self, pprint=False, show=False, save=False, render=False, save_dir=''):
+        colors = color_list()
+        for i, (img, pred) in enumerate(zip(self.imgs, self.pred)):
+            str = f'image {i + 1}/{len(self.pred)}: {img.shape[0]}x{img.shape[1]} '
+            if pred is not None:
+                for c in pred[:, -1].unique():
+                    n = (pred[:, -1] == c).sum()  # detections per class
+                    str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, "  # add to string
+                if show or save or render:
+                    for *box, conf, cls in pred:  # xyxy, confidence, class
+                        label = f'{self.names[int(cls)]} {conf:.2f}'
+                        plot_one_box(box, img, label=label, color=colors[int(cls) % 10])
+            img = Image.fromarray(img.astype(np.uint8)) if isinstance(img, np.ndarray) else img  # from np
+            if pprint:
+                print(str.rstrip(', '))
+            if show:
+                img.show(self.files[i])  # show
+            if save:
+                f = Path(save_dir) / self.files[i]
+                img.save(f)  # save
+                print(f"{'Saving' * (i == 0)} {f},", end='' if i < self.n - 1 else ' done.\n')
+            if render:
+                self.imgs[i] = np.asarray(img)
+
+    def print(self):
+        self.display(pprint=True)  # print results
+
+    def show(self):
+        self.display(show=True)  # show results
+
+    def save(self, save_dir='results/'):
+        Path(save_dir).mkdir(exist_ok=True)
+        self.display(save=True, save_dir=save_dir)  # save results
+
+    def render(self):
+        self.display(render=True)  # render results
+        return self.imgs
+
+    def __len__(self):
+        return self.n
+
+    def tolist(self):
+        # return a list of Detections objects, i.e. 'for result in results.tolist():'
+        x = [Detections([self.imgs[i]], [self.pred[i]], self.names) for i in range(self.n)]
+        for d in x:
+            for k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']:
+                setattr(d, k, getattr(d, k)[0])  # pop out of list
+        return x
+
+
+class Classify(nn.Module):
+    # Classification head, i.e. x(b,c1,20,20) to x(b,c2)
+    def __init__(self, c1, c2, k=1, s=1, p=None, g=1):  # ch_in, ch_out, kernel, stride, padding, groups
+        super(Classify, self).__init__()
+        self.aap = nn.AdaptiveAvgPool2d(1)  # to x(b,c1,1,1)
+        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g)  # to x(b,c2,1,1)
+        self.flat = nn.Flatten()
+
+    def forward(self, x):
+        z = torch.cat([self.aap(y) for y in (x if isinstance(x, list) else [x])], 1)  # cat if list
+        return self.flat(self.conv(z))  # flatten to x(b,c2)

Y5/models/experimental.py

@@ -0,0 +1,135 @@
+# This file contains experimental modules
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+from Y5.models.common import Conv, DWConv
+from Y5.utils.google_utils import attempt_download
+
+
+class CrossConv(nn.Module):
+    # Cross Convolution Downsample
+    def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False):
+        # ch_in, ch_out, kernel, stride, groups, expansion, shortcut
+        super(CrossConv, self).__init__()
+        c_ = int(c2 * e)  # hidden channels
+        self.cv1 = Conv(c1, c_, (1, k), (1, s))
+        self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g)
+        self.add = shortcut and c1 == c2
+
+    def forward(self, x):
+        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
+
+
+class Sum(nn.Module):
+    # Weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070
+    def __init__(self, n, weight=False):  # n: number of inputs
+        super(Sum, self).__init__()
+        self.weight = weight  # apply weights boolean
+        self.iter = range(n - 1)  # iter object
+        if weight:
+            self.w = nn.Parameter(-torch.arange(1., n) / 2, requires_grad=True)  # layer weights
+
+    def forward(self, x):
+        y = x[0]  # no weight
+        if self.weight:
+            w = torch.sigmoid(self.w) * 2
+            for i in self.iter:
+                y = y + x[i + 1] * w[i]
+        else:
+            for i in self.iter:
+                y = y + x[i + 1]
+        return y
+
+
+class GhostConv(nn.Module):
+    # Ghost Convolution https://github.com/huawei-noah/ghostnet
+    def __init__(self, c1, c2, k=1, s=1, g=1, act=True):  # ch_in, ch_out, kernel, stride, groups
+        super(GhostConv, self).__init__()
+        c_ = c2 // 2  # hidden channels
+        self.cv1 = Conv(c1, c_, k, s, None, g, act)
+        self.cv2 = Conv(c_, c_, 5, 1, None, c_, act)
+
+    def forward(self, x):
+        y = self.cv1(x)
+        return torch.cat([y, self.cv2(y)], 1)
+
+
+class GhostBottleneck(nn.Module):
+    # Ghost Bottleneck https://github.com/huawei-noah/ghostnet
+    def __init__(self, c1, c2, k=3, s=1):  # ch_in, ch_out, kernel, stride
+        super(GhostBottleneck, self).__init__()
+        c_ = c2 // 2
+        self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1),  # pw
+                                  DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(),  # dw
+                                  GhostConv(c_, c2, 1, 1, act=False))  # pw-linear
+        self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False),
+                                      Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity()
+
+    def forward(self, x):
+        return self.conv(x) + self.shortcut(x)
+
+
+class MixConv2d(nn.Module):
+    # Mixed Depthwise Conv https://arxiv.org/abs/1907.09595
+    def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True):
+        super(MixConv2d, self).__init__()
+        groups = len(k)
+        if equal_ch:  # equal c_ per group
+            i = torch.linspace(0, groups - 1E-6, c2).floor()  # c2 indices
+            c_ = [(i == g).sum() for g in range(groups)]  # intermediate channels
+        else:  # equal weight.numel() per group
+            b = [c2] + [0] * groups
+            a = np.eye(groups + 1, groups, k=-1)
+            a -= np.roll(a, 1, axis=1)
+            a *= np.array(k) ** 2
+            a[0] = 1
+            c_ = np.linalg.lstsq(a, b, rcond=None)[0].round()  # solve for equal weight indices, ax = b
+
+        self.m = nn.ModuleList([nn.Conv2d(c1, int(c_[g]), k[g], s, k[g] // 2, bias=False) for g in range(groups)])
+        self.bn = nn.BatchNorm2d(c2)
+        self.act = nn.LeakyReLU(0.1, inplace=True)
+
+    def forward(self, x):
+        return x + self.act(self.bn(torch.cat([m(x) for m in self.m], 1)))
+
+
+class Ensemble(nn.ModuleList):
+    # Ensemble of models
+    def __init__(self):
+        super(Ensemble, self).__init__()
+
+    def forward(self, x, augment=False):
+        y = []
+        for module in self:
+            y.append(module(x, augment)[0])
+        # y = torch.stack(y).max(0)[0]  # max ensemble
+        # y = torch.stack(y).mean(0)  # mean ensemble
+        y = torch.cat(y, 1)  # nms ensemble
+        return y, None  # inference, train output
+
+
+def attempt_load(weights, map_location=None):
+    # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
+    model = Ensemble()
+    for w in weights if isinstance(weights, list) else [weights]:
+        attempt_download(w)
+        ckpt = torch.load(w, map_location=map_location)  # load
+        model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval())  # FP32 model
+
+    # Compatibility updates
+    for m in model.modules():
+        if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
+            m.inplace = True  # pytorch 1.7.0 compatibility
+        elif type(m) is Conv:
+            m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility
+
+    if len(model) == 1:
+        return model[-1]  # return model
+    else:
+        print('Ensemble created with %s\n' % weights)
+        for k in ['names', 'stride']:
+            setattr(model, k, getattr(model[-1], k))
+        return model  # return ensemble

Y5/models/export.py

@@ -0,0 +1,104 @@
+"""Exports a YOLOv5 *.pt model to ONNX and TorchScript formats
+
+Usage:
+    $ export PYTHONPATH="$PWD" && python models/export.py --weights ./weights/yolov5s.pt --img 640 --batch 1
+"""
+
+import argparse
+import sys
+import time
+
+sys.path.append('./')  # to run '$ python *.py' files in subdirectories
+
+import torch
+import torch.nn as nn
+
+import Y5.models
+from Y5.models.experimental import attempt_load
+from Y5.utils.activations import Hardswish, SiLU
+from Y5.utils.general import set_logging, check_img_size
+from Y5.utils.torch_utils import select_device
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path')  # from yolov5/models/
+    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size')  # height, width
+    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
+    parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes')
+    parser.add_argument('--grid', action='store_true', help='export Detect() layer grid')
+    parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    opt = parser.parse_args()
+    opt.img_size *= 2 if len(opt.img_size) == 1 else 1  # expand
+    print(opt)
+    set_logging()
+    t = time.time()
+
+    # Load PyTorch model
+    device = select_device(opt.device)
+    model = attempt_load(opt.weights, map_location=device)  # load FP32 model
+    labels = model.names
+
+    # Checks
+    gs = int(max(model.stride))  # grid size (max stride)
+    opt.img_size = [check_img_size(x, gs) for x in opt.img_size]  # verify img_size are gs-multiples
+
+    # Input
+    img = torch.zeros(opt.batch_size, 3, *opt.img_size).to(device)  # image size(1,3,320,192) iDetection
+
+    # Update model
+    for k, m in model.named_modules():
+        m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility
+        if isinstance(m, models.common.Conv):  # assign export-friendly activations
+            if isinstance(m.act, nn.Hardswish):
+                m.act = Hardswish()
+            elif isinstance(m.act, nn.SiLU):
+                m.act = SiLU()
+        # elif isinstance(m, models.yolo.Detect):
+        #     m.forward = m.forward_export  # assign forward (optional)
+    model.model[-1].export = not opt.grid  # set Detect() layer grid export
+    y = model(img)  # dry run
+
+    # TorchScript export
+    try:
+        print('\nStarting TorchScript export with torch %s...' % torch.__version__)
+        f = opt.weights.replace('.pt', '.torchscript.pt')  # filename
+        ts = torch.jit.trace(model, img)
+        ts.save(f)
+        print('TorchScript export success, saved as %s' % f)
+    except Exception as e:
+        print('TorchScript export failure: %s' % e)
+
+    # ONNX export
+    try:
+        import onnx
+
+        print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
+        f = opt.weights.replace('.pt', '.onnx')  # filename
+        torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
+                          output_names=['classes', 'boxes'] if y is None else ['output'],
+                          dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # size(1,3,640,640)
+                                        'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)
+
+        # Checks
+        onnx_model = onnx.load(f)  # load onnx model
+        onnx.checker.check_model(onnx_model)  # check onnx model
+        # print(onnx.helper.printable_graph(onnx_model.graph))  # print a human readable model
+        print('ONNX export success, saved as %s' % f)
+    except Exception as e:
+        print('ONNX export failure: %s' % e)
+
+    # CoreML export
+    try:
+        import coremltools as ct
+
+        print('\nStarting CoreML export with coremltools %s...' % ct.__version__)
+        # convert model from torchscript and apply pixel scaling as per detect.py
+        model = ct.convert(ts, inputs=[ct.ImageType(name='image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
+        f = opt.weights.replace('.pt', '.mlmodel')  # filename
+        model.save(f)
+        print('CoreML export success, saved as %s' % f)
+    except Exception as e:
+        print('CoreML export failure: %s' % e)
+
+    # Finish
+    print('\nExport complete (%.2fs). Visualize with https://github.com/lutzroeder/netron.' % (time.time() - t))

Y5/models/hub/anchors.yaml

@@ -0,0 +1,58 @@
+# Default YOLOv5 anchors for COCO data
+
+
+# P5 -------------------------------------------------------------------------------------------------------------------
+# P5-640:
+anchors_p5_640:
+  - [ 10,13, 16,30, 33,23 ]  # P3/8
+  - [ 30,61, 62,45, 59,119 ]  # P4/16
+  - [ 116,90, 156,198, 373,326 ]  # P5/32
+
+
+# P6 -------------------------------------------------------------------------------------------------------------------
+# P6-640:  thr=0.25: 0.9964 BPR, 5.54 anchors past thr, n=12, img_size=640, metric_all=0.281/0.716-mean/best, past_thr=0.469-mean: 9,11,  21,19,  17,41,  43,32,  39,70,  86,64,  65,131,  134,130,  120,265,  282,180,  247,354,  512,387
+anchors_p6_640:
+  - [ 9,11,  21,19,  17,41 ]  # P3/8
+  - [ 43,32,  39,70,  86,64 ]  # P4/16
+  - [ 65,131,  134,130,  120,265 ]  # P5/32
+  - [ 282,180,  247,354,  512,387 ]  # P6/64
+
+# P6-1280:  thr=0.25: 0.9950 BPR, 5.55 anchors past thr, n=12, img_size=1280, metric_all=0.281/0.714-mean/best, past_thr=0.468-mean: 19,27,  44,40,  38,94,  96,68,  86,152,  180,137,  140,301,  303,264,  238,542,  436,615,  739,380,  925,792
+anchors_p6_1280:
+  - [ 19,27,  44,40,  38,94 ]  # P3/8
+  - [ 96,68,  86,152,  180,137 ]  # P4/16
+  - [ 140,301,  303,264,  238,542 ]  # P5/32
+  - [ 436,615,  739,380,  925,792 ]  # P6/64
+
+# P6-1920:  thr=0.25: 0.9950 BPR, 5.55 anchors past thr, n=12, img_size=1920, metric_all=0.281/0.714-mean/best, past_thr=0.468-mean: 28,41,  67,59,  57,141,  144,103,  129,227,  270,205,  209,452,  455,396,  358,812,  653,922,  1109,570,  1387,1187
+anchors_p6_1920:
+  - [ 28,41,  67,59,  57,141 ]  # P3/8
+  - [ 144,103,  129,227,  270,205 ]  # P4/16
+  - [ 209,452,  455,396,  358,812 ]  # P5/32
+  - [ 653,922,  1109,570,  1387,1187 ]  # P6/64
+
+
+# P7 -------------------------------------------------------------------------------------------------------------------
+# P7-640:  thr=0.25: 0.9962 BPR, 6.76 anchors past thr, n=15, img_size=640, metric_all=0.275/0.733-mean/best, past_thr=0.466-mean: 11,11,  13,30,  29,20,  30,46,  61,38,  39,92,  78,80,  146,66,  79,163,  149,150,  321,143,  157,303,  257,402,  359,290,  524,372
+anchors_p7_640:
+  - [ 11,11,  13,30,  29,20 ]  # P3/8
+  - [ 30,46,  61,38,  39,92 ]  # P4/16
+  - [ 78,80,  146,66,  79,163 ]  # P5/32
+  - [ 149,150,  321,143,  157,303 ]  # P6/64
+  - [ 257,402,  359,290,  524,372 ]  # P7/128
+
+# P7-1280:  thr=0.25: 0.9968 BPR, 6.71 anchors past thr, n=15, img_size=1280, metric_all=0.273/0.732-mean/best, past_thr=0.463-mean: 19,22,  54,36,  32,77,  70,83,  138,71,  75,173,  165,159,  148,334,  375,151,  334,317,  251,626,  499,474,  750,326,  534,814,  1079,818
+anchors_p7_1280:
+  - [ 19,22,  54,36,  32,77 ]  # P3/8
+  - [ 70,83,  138,71,  75,173 ]  # P4/16
+  - [ 165,159,  148,334,  375,151 ]  # P5/32
+  - [ 334,317,  251,626,  499,474 ]  # P6/64
+  - [ 750,326,  534,814,  1079,818 ]  # P7/128
+
+# P7-1920:  thr=0.25: 0.9968 BPR, 6.71 anchors past thr, n=15, img_size=1920, metric_all=0.273/0.732-mean/best, past_thr=0.463-mean: 29,34,  81,55,  47,115,  105,124,  207,107,  113,259,  247,238,  222,500,  563,227,  501,476,  376,939,  749,711,  1126,489,  801,1222,  1618,1227
+anchors_p7_1920:
+  - [ 29,34,  81,55,  47,115 ]  # P3/8
+  - [ 105,124,  207,107,  113,259 ]  # P4/16
+  - [ 247,238,  222,500,  563,227 ]  # P5/32
+  - [ 501,476,  376,939,  749,711 ]  # P6/64
+  - [ 1126,489,  801,1222,  1618,1227 ]  # P7/128

Y5/models/hub/yolov3-spp.yaml

@@ -0,0 +1,51 @@
+# parameters
+nc: 80  # number of classes
+depth_multiple: 1.0  # model depth multiple
+width_multiple: 1.0  # layer channel multiple
+
+# anchors
+anchors:
+  - [10,13, 16,30, 33,23]  # P3/8
+  - [30,61, 62,45, 59,119]  # P4/16
+  - [116,90, 156,198, 373,326]  # P5/32
+
+# darknet53 backbone
+backbone:
+  # [from, number, module, args]
+  [[-1, 1, Conv, [32, 3, 1]],  # 0
+   [-1, 1, Conv, [64, 3, 2]],  # 1-P1/2
+   [-1, 1, Bottleneck, [64]],
+   [-1, 1, Conv, [128, 3, 2]],  # 3-P2/4
+   [-1, 2, Bottleneck, [128]],
+   [-1, 1, Conv, [256, 3, 2]],  # 5-P3/8
+   [-1, 8, Bottleneck, [256]],
+   [-1, 1, Conv, [512, 3, 2]],  # 7-P4/16
+   [-1, 8, Bottleneck, [512]],
+   [-1, 1, Conv, [1024, 3, 2]],  # 9-P5/32
+   [-1, 4, Bottleneck, [1024]],  # 10
+  ]
+
+# YOLOv3-SPP head
+head:
+  [[-1, 1, Bottleneck, [1024, False]],
+   [-1, 1, SPP, [512, [5, 9, 13]]],
+   [-1, 1, Conv, [1024, 3, 1]],
+   [-1, 1, Conv, [512, 1, 1]],
+   [-1, 1, Conv, [1024, 3, 1]],  # 15 (P5/32-large)
+
+   [-2, 1, Conv, [256, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 8], 1, Concat, [1]],  # cat backbone P4
+   [-1, 1, Bottleneck, [512, False]],
+   [-1, 1, Bottleneck, [512, False]],
+   [-1, 1, Conv, [256, 1, 1]],
+   [-1, 1, Conv, [512, 3, 1]],  # 22 (P4/16-medium)
+
+   [-2, 1, Conv, [128, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 6], 1, Concat, [1]],  # cat backbone P3
+   [-1, 1, Bottleneck, [256, False]],
+   [-1, 2, Bottleneck, [256, False]],  # 27 (P3/8-small)
+
+   [[27, 22, 15], 1, Detect, [nc, anchors]],   # Detect(P3, P4, P5)
+  ]

Y5/models/hub/yolov3-tiny.yaml

@@ -0,0 +1,41 @@
+# parameters
+nc: 80  # number of classes
+depth_multiple: 1.0  # model depth multiple
+width_multiple: 1.0  # layer channel multiple
+
+# anchors
+anchors:
+  - [10,14, 23,27, 37,58]  # P4/16
+  - [81,82, 135,169, 344,319]  # P5/32
+
+# YOLOv3-tiny backbone
+backbone:
+  # [from, number, module, args]
+  [[-1, 1, Conv, [16, 3, 1]],  # 0
+   [-1, 1, nn.MaxPool2d, [2, 2, 0]],  # 1-P1/2
+   [-1, 1, Conv, [32, 3, 1]],
+   [-1, 1, nn.MaxPool2d, [2, 2, 0]],  # 3-P2/4
+   [-1, 1, Conv, [64, 3, 1]],
+   [-1, 1, nn.MaxPool2d, [2, 2, 0]],  # 5-P3/8
+   [-1, 1, Conv, [128, 3, 1]],
+   [-1, 1, nn.MaxPool2d, [2, 2, 0]],  # 7-P4/16
+   [-1, 1, Conv, [256, 3, 1]],
+   [-1, 1, nn.MaxPool2d, [2, 2, 0]],  # 9-P5/32
+   [-1, 1, Conv, [512, 3, 1]],
+   [-1, 1, nn.ZeroPad2d, [[0, 1, 0, 1]]],  # 11
+   [-1, 1, nn.MaxPool2d, [2, 1, 0]],  # 12
+  ]
+
+# YOLOv3-tiny head
+head:
+  [[-1, 1, Conv, [1024, 3, 1]],
+   [-1, 1, Conv, [256, 1, 1]],
+   [-1, 1, Conv, [512, 3, 1]],  # 15 (P5/32-large)
+
+   [-2, 1, Conv, [128, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 8], 1, Concat, [1]],  # cat backbone P4
+   [-1, 1, Conv, [256, 3, 1]],  # 19 (P4/16-medium)
+
+   [[19, 15], 1, Detect, [nc, anchors]],  # Detect(P4, P5)
+  ]

Y5/models/hub/yolov3.yaml

@@ -0,0 +1,51 @@
+# parameters
+nc: 80  # number of classes
+depth_multiple: 1.0  # model depth multiple
+width_multiple: 1.0  # layer channel multiple
+
+# anchors
+anchors:
+  - [10,13, 16,30, 33,23]  # P3/8
+  - [30,61, 62,45, 59,119]  # P4/16
+  - [116,90, 156,198, 373,326]  # P5/32
+
+# darknet53 backbone
+backbone:
+  # [from, number, module, args]
+  [[-1, 1, Conv, [32, 3, 1]],  # 0
+   [-1, 1, Conv, [64, 3, 2]],  # 1-P1/2
+   [-1, 1, Bottleneck, [64]],
+   [-1, 1, Conv, [128, 3, 2]],  # 3-P2/4
+   [-1, 2, Bottleneck, [128]],
+   [-1, 1, Conv, [256, 3, 2]],  # 5-P3/8
+   [-1, 8, Bottleneck, [256]],
+   [-1, 1, Conv, [512, 3, 2]],  # 7-P4/16
+   [-1, 8, Bottleneck, [512]],
+   [-1, 1, Conv, [1024, 3, 2]],  # 9-P5/32
+   [-1, 4, Bottleneck, [1024]],  # 10
+  ]
+
+# YOLOv3 head
+head:
+  [[-1, 1, Bottleneck, [1024, False]],
+   [-1, 1, Conv, [512, [1, 1]]],
+   [-1, 1, Conv, [1024, 3, 1]],
+   [-1, 1, Conv, [512, 1, 1]],
+   [-1, 1, Conv, [1024, 3, 1]],  # 15 (P5/32-large)
+
+   [-2, 1, Conv, [256, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 8], 1, Concat, [1]],  # cat backbone P4
+   [-1, 1, Bottleneck, [512, False]],
+   [-1, 1, Bottleneck, [512, False]],
+   [-1, 1, Conv, [256, 1, 1]],
+   [-1, 1, Conv, [512, 3, 1]],  # 22 (P4/16-medium)
+
+   [-2, 1, Conv, [128, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 6], 1, Concat, [1]],  # cat backbone P3
+   [-1, 1, Bottleneck, [256, False]],
+   [-1, 2, Bottleneck, [256, False]],  # 27 (P3/8-small)
+
+   [[27, 22, 15], 1, Detect, [nc, anchors]],   # Detect(P3, P4, P5)
+  ]

Y5/models/hub/yolov5-fpn.yaml

@@ -0,0 +1,42 @@
+# parameters
+nc: 80  # number of classes
+depth_multiple: 1.0  # model depth multiple
+width_multiple: 1.0  # layer channel multiple
+
+# anchors
+anchors:
+  - [10,13, 16,30, 33,23]  # P3/8
+  - [30,61, 62,45, 59,119]  # P4/16
+  - [116,90, 156,198, 373,326]  # P5/32
+
+# YOLOv5 backbone
+backbone:
+  # [from, number, module, args]
+  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
+   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
+   [-1, 3, Bottleneck, [128]],
+   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
+   [-1, 9, BottleneckCSP, [256]],
+   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
+   [-1, 9, BottleneckCSP, [512]],
+   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
+   [-1, 1, SPP, [1024, [5, 9, 13]]],
+   [-1, 6, BottleneckCSP, [1024]],  # 9
+  ]
+
+# YOLOv5 FPN head
+head:
+  [[-1, 3, BottleneckCSP, [1024, False]],  # 10 (P5/32-large)
+
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
+   [-1, 1, Conv, [512, 1, 1]],
+   [-1, 3, BottleneckCSP, [512, False]],  # 14 (P4/16-medium)
+
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
+   [-1, 1, Conv, [256, 1, 1]],
+   [-1, 3, BottleneckCSP, [256, False]],  # 18 (P3/8-small)
+
+   [[18, 14, 10], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
+  ]

Y5/models/hub/yolov5-p2.yaml

@@ -0,0 +1,54 @@
+# parameters
+nc: 80  # number of classes
+depth_multiple: 1.0  # model depth multiple
+width_multiple: 1.0  # layer channel multiple
+
+# anchors
+anchors: 3
+
+# YOLOv5 backbone
+backbone:
+  # [from, number, module, args]
+  [ [ -1, 1, Focus, [ 64, 3 ] ],  # 0-P1/2
+    [ -1, 1, Conv, [ 128, 3, 2 ] ],  # 1-P2/4
+    [ -1, 3, C3, [ 128 ] ],
+    [ -1, 1, Conv, [ 256, 3, 2 ] ],  # 3-P3/8
+    [ -1, 9, C3, [ 256 ] ],
+    [ -1, 1, Conv, [ 512, 3, 2 ] ],  # 5-P4/16
+    [ -1, 9, C3, [ 512 ] ],
+    [ -1, 1, Conv, [ 1024, 3, 2 ] ],  # 7-P5/32
+    [ -1, 1, SPP, [ 1024, [ 5, 9, 13 ] ] ],
+    [ -1, 3, C3, [ 1024, False ] ],  # 9
+  ]
+
+# YOLOv5 head
+head:
+  [ [ -1, 1, Conv, [ 512, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 6 ], 1, Concat, [ 1 ] ],  # cat backbone P4
+    [ -1, 3, C3, [ 512, False ] ],  # 13
+
+    [ -1, 1, Conv, [ 256, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 4 ], 1, Concat, [ 1 ] ],  # cat backbone P3
+    [ -1, 3, C3, [ 256, False ] ],  # 17 (P3/8-small)
+
+    [ -1, 1, Conv, [ 128, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 2 ], 1, Concat, [ 1 ] ],  # cat backbone P2
+    [ -1, 1, C3, [ 128, False ] ],  # 21 (P2/4-xsmall)
+
+    [ -1, 1, Conv, [ 128, 3, 2 ] ],
+    [ [ -1, 18 ], 1, Concat, [ 1 ] ],  # cat head P3
+    [ -1, 3, C3, [ 256, False ] ],  # 24 (P3/8-small)
+
+    [ -1, 1, Conv, [ 256, 3, 2 ] ],
+    [ [ -1, 14 ], 1, Concat, [ 1 ] ],  # cat head P4
+    [ -1, 3, C3, [ 512, False ] ],  # 27 (P4/16-medium)
+
+    [ -1, 1, Conv, [ 512, 3, 2 ] ],
+    [ [ -1, 10 ], 1, Concat, [ 1 ] ],  # cat head P5
+    [ -1, 3, C3, [ 1024, False ] ],  # 30 (P5/32-large)
+
+    [ [ 24, 27, 30 ], 1, Detect, [ nc, anchors ] ],  # Detect(P3, P4, P5)
+  ]

Y5/models/hub/yolov5-p6.yaml

@@ -0,0 +1,56 @@
+# parameters
+nc: 80  # number of classes
+depth_multiple: 1.0  # model depth multiple
+width_multiple: 1.0  # layer channel multiple
+
+# anchors
+anchors: 3
+
+# YOLOv5 backbone
+backbone:
+  # [from, number, module, args]
+  [ [ -1, 1, Focus, [ 64, 3 ] ],  # 0-P1/2
+    [ -1, 1, Conv, [ 128, 3, 2 ] ],  # 1-P2/4
+    [ -1, 3, C3, [ 128 ] ],
+    [ -1, 1, Conv, [ 256, 3, 2 ] ],  # 3-P3/8
+    [ -1, 9, C3, [ 256 ] ],
+    [ -1, 1, Conv, [ 512, 3, 2 ] ],  # 5-P4/16
+    [ -1, 9, C3, [ 512 ] ],
+    [ -1, 1, Conv, [ 768, 3, 2 ] ],  # 7-P5/32
+    [ -1, 3, C3, [ 768 ] ],
+    [ -1, 1, Conv, [ 1024, 3, 2 ] ],  # 9-P6/64
+    [ -1, 1, SPP, [ 1024, [ 3, 5, 7 ] ] ],
+    [ -1, 3, C3, [ 1024, False ] ],  # 11
+  ]
+
+# YOLOv5 head
+head:
+  [ [ -1, 1, Conv, [ 768, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 8 ], 1, Concat, [ 1 ] ],  # cat backbone P5
+    [ -1, 3, C3, [ 768, False ] ],  # 15
+
+    [ -1, 1, Conv, [ 512, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 6 ], 1, Concat, [ 1 ] ],  # cat backbone P4
+    [ -1, 3, C3, [ 512, False ] ],  # 19
+
+    [ -1, 1, Conv, [ 256, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 4 ], 1, Concat, [ 1 ] ],  # cat backbone P3
+    [ -1, 3, C3, [ 256, False ] ],  # 23 (P3/8-small)
+
+    [ -1, 1, Conv, [ 256, 3, 2 ] ],
+    [ [ -1, 20 ], 1, Concat, [ 1 ] ],  # cat head P4
+    [ -1, 3, C3, [ 512, False ] ],  # 26 (P4/16-medium)
+
+    [ -1, 1, Conv, [ 512, 3, 2 ] ],
+    [ [ -1, 16 ], 1, Concat, [ 1 ] ],  # cat head P5
+    [ -1, 3, C3, [ 768, False ] ],  # 29 (P5/32-large)
+
+    [ -1, 1, Conv, [ 768, 3, 2 ] ],
+    [ [ -1, 12 ], 1, Concat, [ 1 ] ],  # cat head P6
+    [ -1, 3, C3, [ 1024, False ] ],  # 32 (P5/64-xlarge)
+
+    [ [ 23, 26, 29, 32 ], 1, Detect, [ nc, anchors ] ],  # Detect(P3, P4, P5, P6)
+  ]

Y5/models/hub/yolov5-p7.yaml

@@ -0,0 +1,67 @@
+# parameters
+nc: 80  # number of classes
+depth_multiple: 1.0  # model depth multiple
+width_multiple: 1.0  # layer channel multiple
+
+# anchors
+anchors: 3
+
+# YOLOv5 backbone
+backbone:
+  # [from, number, module, args]
+  [ [ -1, 1, Focus, [ 64, 3 ] ],  # 0-P1/2
+    [ -1, 1, Conv, [ 128, 3, 2 ] ],  # 1-P2/4
+    [ -1, 3, C3, [ 128 ] ],
+    [ -1, 1, Conv, [ 256, 3, 2 ] ],  # 3-P3/8
+    [ -1, 9, C3, [ 256 ] ],
+    [ -1, 1, Conv, [ 512, 3, 2 ] ],  # 5-P4/16
+    [ -1, 9, C3, [ 512 ] ],
+    [ -1, 1, Conv, [ 768, 3, 2 ] ],  # 7-P5/32
+    [ -1, 3, C3, [ 768 ] ],
+    [ -1, 1, Conv, [ 1024, 3, 2 ] ],  # 9-P6/64
+    [ -1, 3, C3, [ 1024 ] ],
+    [ -1, 1, Conv, [ 1280, 3, 2 ] ],  # 11-P7/128
+    [ -1, 1, SPP, [ 1280, [ 3, 5 ] ] ],
+    [ -1, 3, C3, [ 1280, False ] ],  # 13
+  ]
+
+# YOLOv5 head
+head:
+  [ [ -1, 1, Conv, [ 1024, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 10 ], 1, Concat, [ 1 ] ],  # cat backbone P6
+    [ -1, 3, C3, [ 1024, False ] ],  # 17
+
+    [ -1, 1, Conv, [ 768, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 8 ], 1, Concat, [ 1 ] ],  # cat backbone P5
+    [ -1, 3, C3, [ 768, False ] ],  # 21
+
+    [ -1, 1, Conv, [ 512, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 6 ], 1, Concat, [ 1 ] ],  # cat backbone P4
+    [ -1, 3, C3, [ 512, False ] ],  # 25
+
+    [ -1, 1, Conv, [ 256, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 4 ], 1, Concat, [ 1 ] ],  # cat backbone P3
+    [ -1, 3, C3, [ 256, False ] ],  # 29 (P3/8-small)
+
+    [ -1, 1, Conv, [ 256, 3, 2 ] ],
+    [ [ -1, 26 ], 1, Concat, [ 1 ] ],  # cat head P4
+    [ -1, 3, C3, [ 512, False ] ],  # 32 (P4/16-medium)
+
+    [ -1, 1, Conv, [ 512, 3, 2 ] ],
+    [ [ -1, 22 ], 1, Concat, [ 1 ] ],  # cat head P5
+    [ -1, 3, C3, [ 768, False ] ],  # 35 (P5/32-large)
+
+    [ -1, 1, Conv, [ 768, 3, 2 ] ],
+    [ [ -1, 18 ], 1, Concat, [ 1 ] ],  # cat head P6
+    [ -1, 3, C3, [ 1024, False ] ],  # 38 (P6/64-xlarge)
+
+    [ -1, 1, Conv, [ 1024, 3, 2 ] ],
+    [ [ -1, 14 ], 1, Concat, [ 1 ] ],  # cat head P7
+    [ -1, 3, C3, [ 1280, False ] ],  # 41 (P7/128-xxlarge)
+
+    [ [ 29, 32, 35, 38, 41 ], 1, Detect, [ nc, anchors ] ],  # Detect(P3, P4, P5, P6, P7)
+  ]

Y5/models/hub/yolov5-panet.yaml

@@ -0,0 +1,48 @@
+# parameters
+nc: 80  # number of classes
+depth_multiple: 1.0  # model depth multiple
+width_multiple: 1.0  # layer channel multiple
+
+# anchors
+anchors:
+  - [10,13, 16,30, 33,23]  # P3/8
+  - [30,61, 62,45, 59,119]  # P4/16
+  - [116,90, 156,198, 373,326]  # P5/32
+
+# YOLOv5 backbone
+backbone:
+  # [from, number, module, args]
+  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
+   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
+   [-1, 3, BottleneckCSP, [128]],
+   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
+   [-1, 9, BottleneckCSP, [256]],
+   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
+   [-1, 9, BottleneckCSP, [512]],
+   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
+   [-1, 1, SPP, [1024, [5, 9, 13]]],
+   [-1, 3, BottleneckCSP, [1024, False]],  # 9
+  ]
+
+# YOLOv5 PANet head
+head:
+  [[-1, 1, Conv, [512, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
+   [-1, 3, BottleneckCSP, [512, False]],  # 13
+
+   [-1, 1, Conv, [256, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
+   [-1, 3, BottleneckCSP, [256, False]],  # 17 (P3/8-small)
+
+   [-1, 1, Conv, [256, 3, 2]],
+   [[-1, 14], 1, Concat, [1]],  # cat head P4
+   [-1, 3, BottleneckCSP, [512, False]],  # 20 (P4/16-medium)
+
+   [-1, 1, Conv, [512, 3, 2]],
+   [[-1, 10], 1, Concat, [1]],  # cat head P5
+   [-1, 3, BottleneckCSP, [1024, False]],  # 23 (P5/32-large)
+
+   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
+  ]

Y5/models/hub/yolov5l6.yaml

@@ -0,0 +1,60 @@
+# parameters
+nc: 80  # number of classes
+depth_multiple: 1.0  # model depth multiple
+width_multiple: 1.0  # layer channel multiple
+
+# anchors
+anchors:
+  - [ 19,27,  44,40,  38,94 ]  # P3/8
+  - [ 96,68,  86,152,  180,137 ]  # P4/16
+  - [ 140,301,  303,264,  238,542 ]  # P5/32
+  - [ 436,615,  739,380,  925,792 ]  # P6/64
+
+# YOLOv5 backbone
+backbone:
+  # [from, number, module, args]
+  [ [ -1, 1, Focus, [ 64, 3 ] ],  # 0-P1/2
+    [ -1, 1, Conv, [ 128, 3, 2 ] ],  # 1-P2/4
+    [ -1, 3, C3, [ 128 ] ],
+    [ -1, 1, Conv, [ 256, 3, 2 ] ],  # 3-P3/8
+    [ -1, 9, C3, [ 256 ] ],
+    [ -1, 1, Conv, [ 512, 3, 2 ] ],  # 5-P4/16
+    [ -1, 9, C3, [ 512 ] ],
+    [ -1, 1, Conv, [ 768, 3, 2 ] ],  # 7-P5/32
+    [ -1, 3, C3, [ 768 ] ],
+    [ -1, 1, Conv, [ 1024, 3, 2 ] ],  # 9-P6/64
+    [ -1, 1, SPP, [ 1024, [ 3, 5, 7 ] ] ],
+    [ -1, 3, C3, [ 1024, False ] ],  # 11
+  ]
+
+# YOLOv5 head
+head:
+  [ [ -1, 1, Conv, [ 768, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 8 ], 1, Concat, [ 1 ] ],  # cat backbone P5
+    [ -1, 3, C3, [ 768, False ] ],  # 15
+
+    [ -1, 1, Conv, [ 512, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 6 ], 1, Concat, [ 1 ] ],  # cat backbone P4
+    [ -1, 3, C3, [ 512, False ] ],  # 19
+
+    [ -1, 1, Conv, [ 256, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 4 ], 1, Concat, [ 1 ] ],  # cat backbone P3
+    [ -1, 3, C3, [ 256, False ] ],  # 23 (P3/8-small)
+
+    [ -1, 1, Conv, [ 256, 3, 2 ] ],
+    [ [ -1, 20 ], 1, Concat, [ 1 ] ],  # cat head P4
+    [ -1, 3, C3, [ 512, False ] ],  # 26 (P4/16-medium)
+
+    [ -1, 1, Conv, [ 512, 3, 2 ] ],
+    [ [ -1, 16 ], 1, Concat, [ 1 ] ],  # cat head P5
+    [ -1, 3, C3, [ 768, False ] ],  # 29 (P5/32-large)
+
+    [ -1, 1, Conv, [ 768, 3, 2 ] ],
+    [ [ -1, 12 ], 1, Concat, [ 1 ] ],  # cat head P6
+    [ -1, 3, C3, [ 1024, False ] ],  # 32 (P6/64-xlarge)
+
+    [ [ 23, 26, 29, 32 ], 1, Detect, [ nc, anchors ] ],  # Detect(P3, P4, P5, P6)
+  ]

Y5/models/hub/yolov5m6.yaml

@@ -0,0 +1,60 @@
+# parameters
+nc: 80  # number of classes
+depth_multiple: 0.67  # model depth multiple
+width_multiple: 0.75  # layer channel multiple
+
+# anchors
+anchors:
+  - [ 19,27,  44,40,  38,94 ]  # P3/8
+  - [ 96,68,  86,152,  180,137 ]  # P4/16
+  - [ 140,301,  303,264,  238,542 ]  # P5/32
+  - [ 436,615,  739,380,  925,792 ]  # P6/64
+
+# YOLOv5 backbone
+backbone:
+  # [from, number, module, args]
+  [ [ -1, 1, Focus, [ 64, 3 ] ],  # 0-P1/2
+    [ -1, 1, Conv, [ 128, 3, 2 ] ],  # 1-P2/4
+    [ -1, 3, C3, [ 128 ] ],
+    [ -1, 1, Conv, [ 256, 3, 2 ] ],  # 3-P3/8
+    [ -1, 9, C3, [ 256 ] ],
+    [ -1, 1, Conv, [ 512, 3, 2 ] ],  # 5-P4/16
+    [ -1, 9, C3, [ 512 ] ],
+    [ -1, 1, Conv, [ 768, 3, 2 ] ],  # 7-P5/32
+    [ -1, 3, C3, [ 768 ] ],
+    [ -1, 1, Conv, [ 1024, 3, 2 ] ],  # 9-P6/64
+    [ -1, 1, SPP, [ 1024, [ 3, 5, 7 ] ] ],
+    [ -1, 3, C3, [ 1024, False ] ],  # 11
+  ]
+
+# YOLOv5 head
+head:
+  [ [ -1, 1, Conv, [ 768, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 8 ], 1, Concat, [ 1 ] ],  # cat backbone P5
+    [ -1, 3, C3, [ 768, False ] ],  # 15
+
+    [ -1, 1, Conv, [ 512, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 6 ], 1, Concat, [ 1 ] ],  # cat backbone P4
+    [ -1, 3, C3, [ 512, False ] ],  # 19
+
+    [ -1, 1, Conv, [ 256, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 4 ], 1, Concat, [ 1 ] ],  # cat backbone P3
+    [ -1, 3, C3, [ 256, False ] ],  # 23 (P3/8-small)
+
+    [ -1, 1, Conv, [ 256, 3, 2 ] ],
+    [ [ -1, 20 ], 1, Concat, [ 1 ] ],  # cat head P4
+    [ -1, 3, C3, [ 512, False ] ],  # 26 (P4/16-medium)
+
+    [ -1, 1, Conv, [ 512, 3, 2 ] ],
+    [ [ -1, 16 ], 1, Concat, [ 1 ] ],  # cat head P5
+    [ -1, 3, C3, [ 768, False ] ],  # 29 (P5/32-large)
+
+    [ -1, 1, Conv, [ 768, 3, 2 ] ],
+    [ [ -1, 12 ], 1, Concat, [ 1 ] ],  # cat head P6
+    [ -1, 3, C3, [ 1024, False ] ],  # 32 (P6/64-xlarge)
+
+    [ [ 23, 26, 29, 32 ], 1, Detect, [ nc, anchors ] ],  # Detect(P3, P4, P5, P6)
+  ]

Y5/models/hub/yolov5s6.yaml

@@ -0,0 +1,60 @@
+# parameters
+nc: 80  # number of classes
+depth_multiple: 0.33  # model depth multiple
+width_multiple: 0.50  # layer channel multiple
+
+# anchors
+anchors:
+  - [ 19,27,  44,40,  38,94 ]  # P3/8
+  - [ 96,68,  86,152,  180,137 ]  # P4/16
+  - [ 140,301,  303,264,  238,542 ]  # P5/32
+  - [ 436,615,  739,380,  925,792 ]  # P6/64
+
+# YOLOv5 backbone
+backbone:
+  # [from, number, module, args]
+  [ [ -1, 1, Focus, [ 64, 3 ] ],  # 0-P1/2
+    [ -1, 1, Conv, [ 128, 3, 2 ] ],  # 1-P2/4
+    [ -1, 3, C3, [ 128 ] ],
+    [ -1, 1, Conv, [ 256, 3, 2 ] ],  # 3-P3/8
+    [ -1, 9, C3, [ 256 ] ],
+    [ -1, 1, Conv, [ 512, 3, 2 ] ],  # 5-P4/16
+    [ -1, 9, C3, [ 512 ] ],
+    [ -1, 1, Conv, [ 768, 3, 2 ] ],  # 7-P5/32
+    [ -1, 3, C3, [ 768 ] ],
+    [ -1, 1, Conv, [ 1024, 3, 2 ] ],  # 9-P6/64
+    [ -1, 1, SPP, [ 1024, [ 3, 5, 7 ] ] ],
+    [ -1, 3, C3, [ 1024, False ] ],  # 11
+  ]
+
+# YOLOv5 head
+head:
+  [ [ -1, 1, Conv, [ 768, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 8 ], 1, Concat, [ 1 ] ],  # cat backbone P5
+    [ -1, 3, C3, [ 768, False ] ],  # 15
+
+    [ -1, 1, Conv, [ 512, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 6 ], 1, Concat, [ 1 ] ],  # cat backbone P4
+    [ -1, 3, C3, [ 512, False ] ],  # 19
+
+    [ -1, 1, Conv, [ 256, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 4 ], 1, Concat, [ 1 ] ],  # cat backbone P3
+    [ -1, 3, C3, [ 256, False ] ],  # 23 (P3/8-small)
+
+    [ -1, 1, Conv, [ 256, 3, 2 ] ],
+    [ [ -1, 20 ], 1, Concat, [ 1 ] ],  # cat head P4
+    [ -1, 3, C3, [ 512, False ] ],  # 26 (P4/16-medium)
+
+    [ -1, 1, Conv, [ 512, 3, 2 ] ],
+    [ [ -1, 16 ], 1, Concat, [ 1 ] ],  # cat head P5
+    [ -1, 3, C3, [ 768, False ] ],  # 29 (P5/32-large)
+
+    [ -1, 1, Conv, [ 768, 3, 2 ] ],
+    [ [ -1, 12 ], 1, Concat, [ 1 ] ],  # cat head P6
+    [ -1, 3, C3, [ 1024, False ] ],  # 32 (P6/64-xlarge)
+
+    [ [ 23, 26, 29, 32 ], 1, Detect, [ nc, anchors ] ],  # Detect(P3, P4, P5, P6)
+  ]

Y5/models/hub/yolov5x6.yaml

@@ -0,0 +1,60 @@
+# parameters
+nc: 80  # number of classes
+depth_multiple: 1.33  # model depth multiple
+width_multiple: 1.25  # layer channel multiple
+
+# anchors
+anchors:
+  - [ 19,27,  44,40,  38,94 ]  # P3/8
+  - [ 96,68,  86,152,  180,137 ]  # P4/16
+  - [ 140,301,  303,264,  238,542 ]  # P5/32
+  - [ 436,615,  739,380,  925,792 ]  # P6/64
+
+# YOLOv5 backbone
+backbone:
+  # [from, number, module, args]
+  [ [ -1, 1, Focus, [ 64, 3 ] ],  # 0-P1/2
+    [ -1, 1, Conv, [ 128, 3, 2 ] ],  # 1-P2/4
+    [ -1, 3, C3, [ 128 ] ],
+    [ -1, 1, Conv, [ 256, 3, 2 ] ],  # 3-P3/8
+    [ -1, 9, C3, [ 256 ] ],
+    [ -1, 1, Conv, [ 512, 3, 2 ] ],  # 5-P4/16
+    [ -1, 9, C3, [ 512 ] ],
+    [ -1, 1, Conv, [ 768, 3, 2 ] ],  # 7-P5/32
+    [ -1, 3, C3, [ 768 ] ],
+    [ -1, 1, Conv, [ 1024, 3, 2 ] ],  # 9-P6/64
+    [ -1, 1, SPP, [ 1024, [ 3, 5, 7 ] ] ],
+    [ -1, 3, C3, [ 1024, False ] ],  # 11
+  ]
+
+# YOLOv5 head
+head:
+  [ [ -1, 1, Conv, [ 768, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 8 ], 1, Concat, [ 1 ] ],  # cat backbone P5
+    [ -1, 3, C3, [ 768, False ] ],  # 15
+
+    [ -1, 1, Conv, [ 512, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 6 ], 1, Concat, [ 1 ] ],  # cat backbone P4
+    [ -1, 3, C3, [ 512, False ] ],  # 19
+
+    [ -1, 1, Conv, [ 256, 1, 1 ] ],
+    [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
+    [ [ -1, 4 ], 1, Concat, [ 1 ] ],  # cat backbone P3
+    [ -1, 3, C3, [ 256, False ] ],  # 23 (P3/8-small)
+
+    [ -1, 1, Conv, [ 256, 3, 2 ] ],
+    [ [ -1, 20 ], 1, Concat, [ 1 ] ],  # cat head P4
+    [ -1, 3, C3, [ 512, False ] ],  # 26 (P4/16-medium)
+
+    [ -1, 1, Conv, [ 512, 3, 2 ] ],
+    [ [ -1, 16 ], 1, Concat, [ 1 ] ],  # cat head P5
+    [ -1, 3, C3, [ 768, False ] ],  # 29 (P5/32-large)
+
+    [ -1, 1, Conv, [ 768, 3, 2 ] ],
+    [ [ -1, 12 ], 1, Concat, [ 1 ] ],  # cat head P6
+    [ -1, 3, C3, [ 1024, False ] ],  # 32 (P6/64-xlarge)
+
+    [ [ 23, 26, 29, 32 ], 1, Detect, [ nc, anchors ] ],  # Detect(P3, P4, P5, P6)
+  ]

Y5/models/yolo.py

@@ -0,0 +1,275 @@
+import argparse
+import logging
+import sys
+from copy import deepcopy
+
+sys.path.append('./')  # to run '$ python *.py' files in subdirectories
+logger = logging.getLogger(__name__)
+
+from Y5.models.common import *
+from Y5.models.experimental import *
+from Y5.utils.autoanchor import check_anchor_order
+from Y5.utils.general import make_divisible, check_file, set_logging
+from Y5.utils.torch_utils import time_synchronized, fuse_conv_and_bn, model_info, scale_img, initialize_weights, \
+    select_device, copy_attr
+
+try:
+    import thop  # for FLOPS computation
+except ImportError:
+    thop = None
+
+
+class Detect(nn.Module):
+    stride = None  # strides computed during build
+    export = False  # onnx export
+
+    def __init__(self, nc=80, anchors=(), ch=()):  # detection layer
+        super(Detect, self).__init__()
+        self.nc = nc  # number of classes
+        self.no = nc + 5  # number of outputs per anchor
+        self.nl = len(anchors)  # number of detection layers
+        self.na = len(anchors[0]) // 2  # number of anchors
+        self.grid = [torch.zeros(1)] * self.nl  # init grid
+        a = torch.tensor(anchors).float().view(self.nl, -1, 2)
+        self.register_buffer('anchors', a)  # shape(nl,na,2)
+        self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2))  # shape(nl,1,na,1,1,2)
+        self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch)  # output conv
+
+    def forward(self, x):
+        # x = x.copy()  # for profiling
+        z = []  # inference output
+        self.training |= self.export
+        for i in range(self.nl):
+            x[i] = self.m[i](x[i])  # conv
+            bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
+            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
+
+            if not self.training:  # inference
+                if self.grid[i].shape[2:4] != x[i].shape[2:4]:
+                    self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
+
+                y = x[i].sigmoid()
+                y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
+                y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
+                z.append(y.view(bs, -1, self.no))
+
+        return x if self.training else (torch.cat(z, 1), x)
+
+    @staticmethod
+    def _make_grid(nx=20, ny=20):
+        yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
+        return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()
+
+
+class Model(nn.Module):
+    def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None):  # model, input channels, number of classes
+        super(Model, self).__init__()
+        if isinstance(cfg, dict):
+            self.yaml = cfg  # model dict
+        else:  # is *.yaml
+            import yaml  # for torch hub
+            self.yaml_file = Path(cfg).name
+            with open(cfg) as f:
+                self.yaml = yaml.load(f, Loader=yaml.SafeLoader)  # model dict
+
+        # Define model
+        ch = self.yaml['ch'] = self.yaml.get('ch', ch)  # input channels
+        if nc and nc != self.yaml['nc']:
+            logger.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
+            self.yaml['nc'] = nc  # override yaml value
+        if anchors:
+            logger.info(f'Overriding model.yaml anchors with anchors={anchors}')
+            self.yaml['anchors'] = round(anchors)  # override yaml value
+        self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch])  # model, savelist
+        self.names = [str(i) for i in range(self.yaml['nc'])]  # default names
+        # print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))])
+
+        # Build strides, anchors
+        m = self.model[-1]  # Detect()
+        if isinstance(m, Detect):
+            s = 256  # 2x min stride
+            m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))])  # forward
+            m.anchors /= m.stride.view(-1, 1, 1)
+            check_anchor_order(m)
+            self.stride = m.stride
+            self._initialize_biases()  # only run once
+            # print('Strides: %s' % m.stride.tolist())
+
+        # Init weights, biases
+        initialize_weights(self)
+        self.info()
+        logger.info('')
+
+    def forward(self, x, augment=False, profile=False):
+        if augment:
+            img_size = x.shape[-2:]  # height, width
+            s = [1, 0.83, 0.67]  # scales
+            f = [None, 3, None]  # flips (2-ud, 3-lr)
+            y = []  # outputs
+            for si, fi in zip(s, f):
+                xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))
+                yi = self.forward_once(xi)[0]  # forward
+                # cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1])  # save
+                yi[..., :4] /= si  # de-scale
+                if fi == 2:
+                    yi[..., 1] = img_size[0] - yi[..., 1]  # de-flip ud
+                elif fi == 3:
+                    yi[..., 0] = img_size[1] - yi[..., 0]  # de-flip lr
+                y.append(yi)
+            return torch.cat(y, 1), None  # augmented inference, train
+        else:
+            return self.forward_once(x, profile)  # single-scale inference, train
+
+    def forward_once(self, x, profile=False):
+        y, dt = [], []  # outputs
+        for m in self.model:
+            if m.f != -1:  # if not from previous layer
+                x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]  # from earlier layers
+
+            if profile:
+                o = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 if thop else 0  # FLOPS
+                t = time_synchronized()
+                for _ in range(10):
+                    _ = m(x)
+                dt.append((time_synchronized() - t) * 100)
+                print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type))
+
+            x = m(x)  # run
+            y.append(x if m.i in self.save else None)  # save output
+
+        if profile:
+            print('%.1fms total' % sum(dt))
+        return x
+
+    def _initialize_biases(self, cf=None):  # initialize biases into Detect(), cf is class frequency
+        # https://arxiv.org/abs/1708.02002 section 3.3
+        # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1.
+        m = self.model[-1]  # Detect() module
+        for mi, s in zip(m.m, m.stride):  # from
+            b = mi.bias.view(m.na, -1)  # conv.bias(255) to (3,85)
+            b.data[:, 4] += math.log(8 / (640 / s) ** 2)  # obj (8 objects per 640 image)
+            b.data[:, 5:] += math.log(0.6 / (m.nc - 0.99)) if cf is None else torch.log(cf / cf.sum())  # cls
+            mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
+
+    def _print_biases(self):
+        m = self.model[-1]  # Detect() module
+        for mi in m.m:  # from
+            b = mi.bias.detach().view(m.na, -1).T  # conv.bias(255) to (3,85)
+            print(('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean()))
+
+    # def _print_weights(self):
+    #     for m in self.model.modules():
+    #         if type(m) is Bottleneck:
+    #             print('%10.3g' % (m.w.detach().sigmoid() * 2))  # shortcut weights
+
+    def fuse(self):  # fuse model Conv2d() + BatchNorm2d() layers
+        #print('Fusing layers... ')
+        for m in self.model.modules():
+            if type(m) is Conv and hasattr(m, 'bn'):
+                m.conv = fuse_conv_and_bn(m.conv, m.bn)  # update conv
+                delattr(m, 'bn')  # remove batchnorm
+                m.forward = m.fuseforward  # update forward
+        self.info()
+        return self
+
+    def nms(self, mode=True):  # add or remove NMS module
+        present = type(self.model[-1]) is NMS  # last layer is NMS
+        if mode and not present:
+            print('Adding NMS... ')
+            m = NMS()  # module
+            m.f = -1  # from
+            m.i = self.model[-1].i + 1  # index
+            self.model.add_module(name='%s' % m.i, module=m)  # add
+            self.eval()
+        elif not mode and present:
+            print('Removing NMS... ')
+            self.model = self.model[:-1]  # remove
+        return self
+
+    def autoshape(self):  # add autoShape module
+        print('Adding autoShape... ')
+        m = autoShape(self)  # wrap model
+        copy_attr(m, self, include=('yaml', 'nc', 'hyp', 'names', 'stride'), exclude=())  # copy attributes
+        return m
+
+    def info(self, verbose=False, img_size=640):  # print model information
+        model_info(self, verbose, img_size)
+
+
+def parse_model(d, ch):  # model_dict, input_channels(3)
+    logger.info('\n%3s%18s%3s%10s  %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments'))
+    anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple']
+    na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors  # number of anchors
+    no = na * (nc + 5)  # number of outputs = anchors * (classes + 5)
+
+    layers, save, c2 = [], [], ch[-1]  # layers, savelist, ch out
+    for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']):  # from, number, module, args
+        m = eval(m) if isinstance(m, str) else m  # eval strings
+        for j, a in enumerate(args):
+            try:
+                args[j] = eval(a) if isinstance(a, str) else a  # eval strings
+            except:
+                pass
+
+        n = max(round(n * gd), 1) if n > 1 else n  # depth gain
+        if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP,
+                 C3]:
+            c1, c2 = ch[f], args[0]
+            if c2 != no:  # if not output
+                c2 = make_divisible(c2 * gw, 8)
+
+            args = [c1, c2, *args[1:]]
+            if m in [BottleneckCSP, C3]:
+                args.insert(2, n)  # number of repeats
+                n = 1
+        elif m is nn.BatchNorm2d:
+            args = [ch[f]]
+        elif m is Concat:
+            c2 = sum([ch[x] for x in f])
+        elif m is Detect:
+            args.append([ch[x] for x in f])
+            if isinstance(args[1], int):  # number of anchors
+                args[1] = [list(range(args[1] * 2))] * len(f)
+        elif m is Contract:
+            c2 = ch[f] * args[0] ** 2
+        elif m is Expand:
+            c2 = ch[f] // args[0] ** 2
+        else:
+            c2 = ch[f]
+
+        m_ = nn.Sequential(*[m(*args) for _ in range(n)]) if n > 1 else m(*args)  # module
+        t = str(m)[8:-2].replace('__main__.', '')  # module type
+        np = sum([x.numel() for x in m_.parameters()])  # number params
+        m_.i, m_.f, m_.type, m_.np = i, f, t, np  # attach index, 'from' index, type, number params
+        logger.info('%3s%18s%3s%10.0f  %-40s%-30s' % (i, f, n, np, t, args))  # print
+        save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelist
+        layers.append(m_)
+        if i == 0:
+            ch = []
+        ch.append(c2)
+    return nn.Sequential(*layers), sorted(save)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml')
+    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    opt = parser.parse_args()
+    opt.cfg = check_file(opt.cfg)  # check file
+    set_logging()
+    device = select_device(opt.device)
+
+    # Create model
+    model = Model(opt.cfg).to(device)
+    model.train()
+
+    # Profile
+    # img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
+    # y = model(img, profile=True)
+
+    # Tensorboard
+    # from torch.utils.tensorboard import SummaryWriter
+    # tb_writer = SummaryWriter()
+    # print("Run 'tensorboard --logdir=models/runs' to view tensorboard at http://localhost:6006/")
+    # tb_writer.add_graph(model.model, img)  # add model to tensorboard
+    # tb_writer.add_image('test', img[0], dataformats='CWH')  # add model to tensorboard

Y5/models/yolov5l.yaml

@@ -0,0 +1,48 @@
+# parameters
+nc: 80  # number of classes
+depth_multiple: 1.0  # model depth multiple
+width_multiple: 1.0  # layer channel multiple
+
+# anchors
+anchors:
+  - [10,13, 16,30, 33,23]  # P3/8
+  - [30,61, 62,45, 59,119]  # P4/16
+  - [116,90, 156,198, 373,326]  # P5/32
+
+# YOLOv5 backbone
+backbone:
+  # [from, number, module, args]
+  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
+   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
+   [-1, 3, C3, [128]],
+   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
+   [-1, 9, C3, [256]],
+   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
+   [-1, 9, C3, [512]],
+   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
+   [-1, 1, SPP, [1024, [5, 9, 13]]],
+   [-1, 3, C3, [1024, False]],  # 9
+  ]
+
+# YOLOv5 head
+head:
+  [[-1, 1, Conv, [512, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
+   [-1, 3, C3, [512, False]],  # 13
+
+   [-1, 1, Conv, [256, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
+   [-1, 3, C3, [256, False]],  # 17 (P3/8-small)
+
+   [-1, 1, Conv, [256, 3, 2]],
+   [[-1, 14], 1, Concat, [1]],  # cat head P4
+   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)
+
+   [-1, 1, Conv, [512, 3, 2]],
+   [[-1, 10], 1, Concat, [1]],  # cat head P5
+   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)
+
+   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
+  ]

Y5/models/yolov5m.yaml

@@ -0,0 +1,48 @@
+# parameters
+nc: 80  # number of classes
+depth_multiple: 0.67  # model depth multiple
+width_multiple: 0.75  # layer channel multiple
+
+# anchors
+anchors:
+  - [10,13, 16,30, 33,23]  # P3/8
+  - [30,61, 62,45, 59,119]  # P4/16
+  - [116,90, 156,198, 373,326]  # P5/32
+
+# YOLOv5 backbone
+backbone:
+  # [from, number, module, args]
+  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
+   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
+   [-1, 3, C3, [128]],
+   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
+   [-1, 9, C3, [256]],
+   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
+   [-1, 9, C3, [512]],
+   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
+   [-1, 1, SPP, [1024, [5, 9, 13]]],
+   [-1, 3, C3, [1024, False]],  # 9
+  ]
+
+# YOLOv5 head
+head:
+  [[-1, 1, Conv, [512, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
+   [-1, 3, C3, [512, False]],  # 13
+
+   [-1, 1, Conv, [256, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
+   [-1, 3, C3, [256, False]],  # 17 (P3/8-small)
+
+   [-1, 1, Conv, [256, 3, 2]],
+   [[-1, 14], 1, Concat, [1]],  # cat head P4
+   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)
+
+   [-1, 1, Conv, [512, 3, 2]],
+   [[-1, 10], 1, Concat, [1]],  # cat head P5
+   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)
+
+   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
+  ]

Y5/models/yolov5s.yaml

@@ -0,0 +1,48 @@
+# parameters
+nc: 80  # number of classes
+depth_multiple: 0.33  # model depth multiple
+width_multiple: 0.50  # layer channel multiple
+
+# anchors
+anchors:
+  - [10,13, 16,30, 33,23]  # P3/8
+  - [30,61, 62,45, 59,119]  # P4/16
+  - [116,90, 156,198, 373,326]  # P5/32
+
+# YOLOv5 backbone
+backbone:
+  # [from, number, module, args]
+  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
+   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
+   [-1, 3, C3, [128]],
+   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
+   [-1, 9, C3, [256]],
+   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
+   [-1, 9, C3, [512]],
+   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
+   [-1, 1, SPP, [1024, [5, 9, 13]]],
+   [-1, 3, C3, [1024, False]],  # 9
+  ]
+
+# YOLOv5 head
+head:
+  [[-1, 1, Conv, [512, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
+   [-1, 3, C3, [512, False]],  # 13
+
+   [-1, 1, Conv, [256, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
+   [-1, 3, C3, [256, False]],  # 17 (P3/8-small)
+
+   [-1, 1, Conv, [256, 3, 2]],
+   [[-1, 14], 1, Concat, [1]],  # cat head P4
+   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)
+
+   [-1, 1, Conv, [512, 3, 2]],
+   [[-1, 10], 1, Concat, [1]],  # cat head P5
+   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)
+
+   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
+  ]

Y5/models/yolov5x.yaml

@@ -0,0 +1,48 @@
+# parameters
+nc: 80  # number of classes
+depth_multiple: 1.33  # model depth multiple
+width_multiple: 1.25  # layer channel multiple
+
+# anchors
+anchors:
+  - [10,13, 16,30, 33,23]  # P3/8
+  - [30,61, 62,45, 59,119]  # P4/16
+  - [116,90, 156,198, 373,326]  # P5/32
+
+# YOLOv5 backbone
+backbone:
+  # [from, number, module, args]
+  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
+   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
+   [-1, 3, C3, [128]],
+   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
+   [-1, 9, C3, [256]],
+   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
+   [-1, 9, C3, [512]],
+   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
+   [-1, 1, SPP, [1024, [5, 9, 13]]],
+   [-1, 3, C3, [1024, False]],  # 9
+  ]
+
+# YOLOv5 head
+head:
+  [[-1, 1, Conv, [512, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
+   [-1, 3, C3, [512, False]],  # 13
+
+   [-1, 1, Conv, [256, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
+   [-1, 3, C3, [256, False]],  # 17 (P3/8-small)
+
+   [-1, 1, Conv, [256, 3, 2]],
+   [[-1, 14], 1, Concat, [1]],  # cat head P4
+   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)
+
+   [-1, 1, Conv, [512, 3, 2]],
+   [[-1, 10], 1, Concat, [1]],  # cat head P5
+   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)
+
+   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
+  ]

Y5/utils/activations.py

@@ -0,0 +1,72 @@
+# Activation functions
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# SiLU https://arxiv.org/pdf/1606.08415.pdf ----------------------------------------------------------------------------
+class SiLU(nn.Module):  # export-friendly version of nn.SiLU()
+    @staticmethod
+    def forward(x):
+        return x * torch.sigmoid(x)
+
+
+class Hardswish(nn.Module):  # export-friendly version of nn.Hardswish()
+    @staticmethod
+    def forward(x):
+        # return x * F.hardsigmoid(x)  # for torchscript and CoreML
+        return x * F.hardtanh(x + 3, 0., 6.) / 6.  # for torchscript, CoreML and ONNX
+
+
+class MemoryEfficientSwish(nn.Module):
+    class F(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, x):
+            ctx.save_for_backward(x)
+            return x * torch.sigmoid(x)
+
+        @staticmethod
+        def backward(ctx, grad_output):
+            x = ctx.saved_tensors[0]
+            sx = torch.sigmoid(x)
+            return grad_output * (sx * (1 + x * (1 - sx)))
+
+    def forward(self, x):
+        return self.F.apply(x)
+
+
+# Mish https://github.com/digantamisra98/Mish --------------------------------------------------------------------------
+class Mish(nn.Module):
+    @staticmethod
+    def forward(x):
+        return x * F.softplus(x).tanh()
+
+
+class MemoryEfficientMish(nn.Module):
+    class F(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, x):
+            ctx.save_for_backward(x)
+            return x.mul(torch.tanh(F.softplus(x)))  # x * tanh(ln(1 + exp(x)))
+
+        @staticmethod
+        def backward(ctx, grad_output):
+            x = ctx.saved_tensors[0]
+            sx = torch.sigmoid(x)
+            fx = F.softplus(x).tanh()
+            return grad_output * (fx + x * sx * (1 - fx * fx))
+
+    def forward(self, x):
+        return self.F.apply(x)
+
+
+# FReLU https://arxiv.org/abs/2007.11824 -------------------------------------------------------------------------------
+class FReLU(nn.Module):
+    def __init__(self, c1, k=3):  # ch_in, kernel
+        super().__init__()
+        self.conv = nn.Conv2d(c1, c1, k, 1, 1, groups=c1, bias=False)
+        self.bn = nn.BatchNorm2d(c1)
+
+    def forward(self, x):
+        return torch.max(x, self.bn(self.conv(x)))

Y5/utils/autoanchor.py

@@ -0,0 +1,155 @@
+# Auto-anchor utils
+
+import numpy as np
+import torch
+import yaml
+from scipy.cluster.vq import kmeans
+from tqdm import tqdm
+
+from Y5.utils.general import colorstr
+
+
+def check_anchor_order(m):
+    # Check anchor order against stride order for YOLOv5 Detect() module m, and correct if necessary
+    a = m.anchor_grid.prod(-1).view(-1)  # anchor area
+    da = a[-1] - a[0]  # delta a
+    ds = m.stride[-1] - m.stride[0]  # delta s
+    if da.sign() != ds.sign():  # same order
+        print('Reversing anchor order')
+        m.anchors[:] = m.anchors.flip(0)
+        m.anchor_grid[:] = m.anchor_grid.flip(0)
+
+
+def check_anchors(dataset, model, thr=4.0, imgsz=640):
+    # Check anchor fit to data, recompute if necessary
+    prefix = colorstr('autoanchor: ')
+    print(f'\n{prefix}Analyzing anchors... ', end='')
+    m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1]  # Detect()
+    shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True)
+    scale = np.random.uniform(0.9, 1.1, size=(shapes.shape[0], 1))  # augment scale
+    wh = torch.tensor(np.concatenate([l[:, 3:5] * s for s, l in zip(shapes * scale, dataset.labels)])).float()  # wh
+
+    def metric(k):  # compute metric
+        r = wh[:, None] / k[None]
+        x = torch.min(r, 1. / r).min(2)[0]  # ratio metric
+        best = x.max(1)[0]  # best_x
+        aat = (x > 1. / thr).float().sum(1).mean()  # anchors above threshold
+        bpr = (best > 1. / thr).float().mean()  # best possible recall
+        return bpr, aat
+
+    bpr, aat = metric(m.anchor_grid.clone().cpu().view(-1, 2))
+    print(f'anchors/target = {aat:.2f}, Best Possible Recall (BPR) = {bpr:.4f}', end='')
+    if bpr < 0.98:  # threshold to recompute
+        print('. Attempting to improve anchors, please wait...')
+        na = m.anchor_grid.numel() // 2  # number of anchors
+        new_anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False)
+        new_bpr = metric(new_anchors.reshape(-1, 2))[0]
+        if new_bpr > bpr:  # replace anchors
+            new_anchors = torch.tensor(new_anchors, device=m.anchors.device).type_as(m.anchors)
+            m.anchor_grid[:] = new_anchors.clone().view_as(m.anchor_grid)  # for inference
+            m.anchors[:] = new_anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1)  # loss
+            check_anchor_order(m)
+            print(f'{prefix}New anchors saved to model. Update model *.yaml to use these anchors in the future.')
+        else:
+            print(f'{prefix}Original anchors better than new anchors. Proceeding with original anchors.')
+    print('')  # newline
+
+
+def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True):
+    """ Creates kmeans-evolved anchors from training dataset
+
+        Arguments:
+            path: path to dataset *.yaml, or a loaded dataset
+            n: number of anchors
+            img_size: image size used for training
+            thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0
+            gen: generations to evolve anchors using genetic algorithm
+            verbose: print all results
+
+        Return:
+            k: kmeans evolved anchors
+
+        Usage:
+            from utils.autoanchor import *; _ = kmean_anchors()
+    """
+    thr = 1. / thr
+    prefix = colorstr('autoanchor: ')
+
+    def metric(k, wh):  # compute metrics
+        r = wh[:, None] / k[None]
+        x = torch.min(r, 1. / r).min(2)[0]  # ratio metric
+        # x = wh_iou(wh, torch.tensor(k))  # iou metric
+        return x, x.max(1)[0]  # x, best_x
+
+    def anchor_fitness(k):  # mutation fitness
+        _, best = metric(torch.tensor(k, dtype=torch.float32), wh)
+        return (best * (best > thr).float()).mean()  # fitness
+
+    def print_results(k):
+        k = k[np.argsort(k.prod(1))]  # sort small to large
+        x, best = metric(k, wh0)
+        bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n  # best possible recall, anch > thr
+        print(f'{prefix}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr')
+        print(f'{prefix}n={n}, img_size={img_size}, metric_all={x.mean():.3f}/{best.mean():.3f}-mean/best, '
+              f'past_thr={x[x > thr].mean():.3f}-mean: ', end='')
+        for i, x in enumerate(k):
+            print('%i,%i' % (round(x[0]), round(x[1])), end=',  ' if i < len(k) - 1 else '\n')  # use in *.cfg
+        return k
+
+    if isinstance(path, str):  # *.yaml file
+        with open(path) as f:
+            data_dict = yaml.load(f, Loader=yaml.SafeLoader)  # model dict
+        from utils.datasets import LoadImagesAndLabels
+        dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True)
+    else:
+        dataset = path  # dataset
+
+    # Get label wh
+    shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True)
+    wh0 = np.concatenate([l[:, 3:5] * s for s, l in zip(shapes, dataset.labels)])  # wh
+
+    # Filter
+    i = (wh0 < 3.0).any(1).sum()
+    if i:
+        print(f'{prefix}WARNING: Extremely small objects found. {i} of {len(wh0)} labels are < 3 pixels in size.')
+    wh = wh0[(wh0 >= 2.0).any(1)]  # filter > 2 pixels
+    # wh = wh * (np.random.rand(wh.shape[0], 1) * 0.9 + 0.1)  # multiply by random scale 0-1
+
+    # Kmeans calculation
+    print(f'{prefix}Running kmeans for {n} anchors on {len(wh)} points...')
+    s = wh.std(0)  # sigmas for whitening
+    k, dist = kmeans(wh / s, n, iter=30)  # points, mean distance
+    k *= s
+    wh = torch.tensor(wh, dtype=torch.float32)  # filtered
+    wh0 = torch.tensor(wh0, dtype=torch.float32)  # unfiltered
+    k = print_results(k)
+
+    # Plot
+    # k, d = [None] * 20, [None] * 20
+    # for i in tqdm(range(1, 21)):
+    #     k[i-1], d[i-1] = kmeans(wh / s, i)  # points, mean distance
+    # fig, ax = plt.subplots(1, 2, figsize=(14, 7), tight_layout=True)
+    # ax = ax.ravel()
+    # ax[0].plot(np.arange(1, 21), np.array(d) ** 2, marker='.')
+    # fig, ax = plt.subplots(1, 2, figsize=(14, 7))  # plot wh
+    # ax[0].hist(wh[wh[:, 0]<100, 0],400)
+    # ax[1].hist(wh[wh[:, 1]<100, 1],400)
+    # fig.savefig('wh.png', dpi=200)
+
+    # Evolve
+    npr = np.random
+    f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1  # fitness, generations, mutation prob, sigma
+    pbar = tqdm(range(gen), desc=f'{prefix}Evolving anchors with Genetic Algorithm:')  # progress bar
+    for _ in pbar:
+        v = np.ones(sh)
+        while (v == 1).all():  # mutate until a change occurs (prevent duplicates)
+            v = ((npr.random(sh) < mp) * npr.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0)
+        kg = (k.copy() * v).clip(min=2.0)
+        fg = anchor_fitness(kg)
+        if fg > f:
+            f, k = fg, kg.copy()
+            pbar.desc = f'{prefix}Evolving anchors with Genetic Algorithm: fitness = {f:.4f}'
+            if verbose:
+                print_results(k)
+
+    return print_results(k)

Y5/utils/aws/mime.sh

@@ -0,0 +1,26 @@
+# AWS EC2 instance startup 'MIME' script https://aws.amazon.com/premiumsupport/knowledge-center/execute-user-data-ec2/
+# This script will run on every instance restart, not only on first start
+# --- DO NOT COPY ABOVE COMMENTS WHEN PASTING INTO USERDATA ---
+
+Content-Type: multipart/mixed; boundary="//"
+MIME-Version: 1.0
+
+--//
+Content-Type: text/cloud-config; charset="us-ascii"
+MIME-Version: 1.0
+Content-Transfer-Encoding: 7bit
+Content-Disposition: attachment; filename="cloud-config.txt"
+
+#cloud-config
+cloud_final_modules:
+- [scripts-user, always]
+
+--//
+Content-Type: text/x-shellscript; charset="us-ascii"
+MIME-Version: 1.0
+Content-Transfer-Encoding: 7bit
+Content-Disposition: attachment; filename="userdata.txt"
+
+#!/bin/bash
+# --- paste contents of userdata.sh here ---
+--//

Y5/utils/aws/resume.py

@@ -0,0 +1,37 @@
+# Resume all interrupted trainings in yolov5/ dir including DPP trainings
+# Usage: $ python utils/aws/resume.py
+
+import os
+import sys
+from pathlib import Path
+
+import torch
+import yaml
+
+sys.path.append('./')  # to run '$ python *.py' files in subdirectories
+
+port = 0  # --master_port
+path = Path('').resolve()
+for last in path.rglob('*/**/last.pt'):
+    ckpt = torch.load(last)
+    if ckpt['optimizer'] is None:
+        continue
+
+    # Load opt.yaml
+    with open(last.parent.parent / 'opt.yaml') as f:
+        opt = yaml.load(f, Loader=yaml.SafeLoader)
+
+    # Get device count
+    d = opt['device'].split(',')  # devices
+    nd = len(d)  # number of devices
+    ddp = nd > 1 or (nd == 0 and torch.cuda.device_count() > 1)  # distributed data parallel
+
+    if ddp:  # multi-GPU
+        port += 1
+        cmd = f'python -m torch.distributed.launch --nproc_per_node {nd} --master_port {port} train.py --resume {last}'
+    else:  # single-GPU
+        cmd = f'python train.py --resume {last}'
+
+    cmd += ' > /dev/null 2>&1 &'  # redirect output to dev/null and run in daemon thread
+    print(cmd)
+    os.system(cmd)

Y5/utils/aws/userdata.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+# AWS EC2 instance startup script https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/user-data.html
+# This script will run only once on first instance start (for a re-start script see mime.sh)
+# /home/ubuntu (ubuntu) or /home/ec2-user (amazon-linux) is working dir
+# Use >300 GB SSD
+
+cd home/ubuntu
+if [ ! -d yolov5 ]; then
+  echo "Running first-time script." # install dependencies, download COCO, pull Docker
+  git clone https://github.com/ultralytics/yolov5 && sudo chmod -R 777 yolov5
+  cd yolov5
+  bash data/scripts/get_coco.sh && echo "Data done." &
+  sudo docker pull ultralytics/yolov5:latest && echo "Docker done." &
+  python -m pip install --upgrade pip && pip install -r requirements.txt && python detect.py && echo "Requirements done." &
+  wait && echo "All tasks done." # finish background tasks
+else
+  echo "Running re-start script." # resume interrupted runs
+  i=0
+  list=$(sudo docker ps -qa) # container list i.e. $'one\ntwo\nthree\nfour'
+  while IFS= read -r id; do
+    ((i++))
+    echo "restarting container $i: $id"
+    sudo docker start $id
+    # sudo docker exec -it $id python train.py --resume # single-GPU
+    sudo docker exec -d $id python utils/aws/resume.py # multi-scenario
+  done <<<"$list"
+fi

Y5/utils/general.py

@@ -0,0 +1,576 @@
+# General utils
+
+import glob
+import logging
+import math
+import os
+import platform
+import random
+import re
+import subprocess
+import time
+from pathlib import Path
+
+import cv2
+import numpy as np
+import torch
+import torchvision
+import yaml
+
+from Y5.utils.google_utils import gsutil_getsize
+from Y5.utils.metrics import fitness
+from Y5.utils.torch_utils import init_torch_seeds
+
+# Settings
+torch.set_printoptions(linewidth=320, precision=5, profile='long')
+np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format})  # format short g, %precision=5
+cv2.setNumThreads(0)  # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader)
+os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8))  # NumExpr max threads
+
+
+def set_logging(rank=-1):
+    logging.basicConfig(
+        format="%(message)s",
+        level=logging.INFO if rank in [-1, 0] else logging.WARN)
+
+
+def init_seeds(seed=0):
+    # Initialize random number generator (RNG) seeds
+    random.seed(seed)
+    np.random.seed(seed)
+    init_torch_seeds(seed)
+
+
+def get_latest_run(search_dir='.'):
+    # Return path to most recent 'last.pt' in /runs (i.e. to --resume from)
+    last_list = glob.glob(f'{search_dir}/**/last*.pt', recursive=True)
+    return max(last_list, key=os.path.getctime) if last_list else ''
+
+
+def isdocker():
+    # Is environment a Docker container
+    return Path('/workspace').exists()  # or Path('/.dockerenv').exists()
+
+
+def check_online():
+    # Check internet connectivity
+    import socket
+    try:
+        socket.create_connection(("1.1.1.1", 443), 5)  # check host accesability
+        return True
+    except OSError:
+        return False
+
+
+def check_git_status():
+    # Recommend 'git pull' if code is out of date
+    print(colorstr('github: '), end='')
+    try:
+        assert Path('.git').exists(), 'skipping check (not a git repository)'
+        assert not isdocker(), 'skipping check (Docker image)'
+        assert check_online(), 'skipping check (offline)'
+
+        cmd = 'git fetch && git config --get remote.origin.url'
+        url = subprocess.check_output(cmd, shell=True).decode().strip().rstrip('.git')  # github repo url
+        branch = subprocess.check_output('git rev-parse --abbrev-ref HEAD', shell=True).decode().strip()  # checked out
+        n = int(subprocess.check_output(f'git rev-list {branch}..origin/master --count', shell=True))  # commits behind
+        if n > 0:
+            s = f"⚠️ WARNING: code is out of date by {n} commit{'s' * (n > 1)}. " \
+                f"Use 'git pull' to update or 'git clone {url}' to download latest."
+        else:
+            s = f'up to date with {url} ✅'
+        print(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s)
+    except Exception as e:
+        print(e)
+
+
+def check_requirements(file='requirements.txt', exclude=()):
+    # Check installed dependencies meet requirements
+    import pkg_resources
+    requirements = [f'{x.name}{x.specifier}' for x in pkg_resources.parse_requirements(Path(file).open())
+                    if x.name not in exclude]
+    pkg_resources.require(requirements)  # DistributionNotFound or VersionConflict exception if requirements not met
+
+
+def check_img_size(img_size, s=32):
+    # Verify img_size is a multiple of stride s
+    new_size = make_divisible(img_size, int(s))  # ceil gs-multiple
+    if new_size != img_size:
+        print('WARNING: --img-size %g must be multiple of max stride %g, updating to %g' % (img_size, s, new_size))
+    return new_size
+
+
+def check_imshow():
+    # Check if environment supports image displays
+    try:
+        assert not isdocker(), 'cv2.imshow() is disabled in Docker environments'
+        cv2.imshow('test', np.zeros((1, 1, 3)))
+        cv2.waitKey(1)
+        cv2.destroyAllWindows()
+        cv2.waitKey(1)
+        return True
+    except Exception as e:
+        print(f'WARNING: Environment does not support cv2.imshow() or PIL Image.show() image displays\n{e}')
+        return False
+
+
+def check_file(file):
+    # Search for file if not found
+    if os.path.isfile(file) or file == '':
+        return file
+    else:
+        files = glob.glob('./**/' + file, recursive=True)  # find file
+        assert len(files), 'File Not Found: %s' % file  # assert file was found
+        assert len(files) == 1, "Multiple files match '%s', specify exact path: %s" % (file, files)  # assert unique
+        return files[0]  # return file
+
+
+def check_dataset(dict):
+    # Download dataset if not found locally
+    val, s = dict.get('val'), dict.get('download')
+    if val and len(val):
+        val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])]  # val path
+        if not all(x.exists() for x in val):
+            print('\nWARNING: Dataset not found, nonexistent paths: %s' % [str(x) for x in val if not x.exists()])
+            if s and len(s):  # download script
+                print('Downloading %s ...' % s)
+                if s.startswith('http') and s.endswith('.zip'):  # URL
+                    f = Path(s).name  # filename
+                    torch.hub.download_url_to_file(s, f)
+                    r = os.system('unzip -q %s -d ../ && rm %s' % (f, f))  # unzip
+                else:  # bash script
+                    r = os.system(s)
+                print('Dataset autodownload %s\n' % ('success' if r == 0 else 'failure'))  # analyze return value
+            else:
+                raise Exception('Dataset not found.')
+
+
+def make_divisible(x, divisor):
+    # Returns x evenly divisible by divisor
+    return math.ceil(x / divisor) * divisor
+
+
+def clean_str(s):
+    # Cleans a string by replacing special characters with underscore _
+    return re.sub(pattern="[|@#!¡·$€%&()=?¿^*;:,¨´><+]", repl="_", string=s)
+
+
+def one_cycle(y1=0.0, y2=1.0, steps=100):
+    # lambda function for sinusoidal ramp from y1 to y2
+    return lambda x: ((1 - math.cos(x * math.pi / steps)) / 2) * (y2 - y1) + y1
+
+
+def colorstr(*input):
+    # Colors a string https://en.wikipedia.org/wiki/ANSI_escape_code, i.e.  colorstr('blue', 'hello world')
+    *args, string = input if len(input) > 1 else ('blue', 'bold', input[0])  # color arguments, string
+    colors = {'black': '\033[30m',  # basic colors
+              'red': '\033[31m',
+              'green': '\033[32m',
+              'yellow': '\033[33m',
+              'blue': '\033[34m',
+              'magenta': '\033[35m',
+              'cyan': '\033[36m',
+              'white': '\033[37m',
+              'bright_black': '\033[90m',  # bright colors
+              'bright_red': '\033[91m',
+              'bright_green': '\033[92m',
+              'bright_yellow': '\033[93m',
+              'bright_blue': '\033[94m',
+              'bright_magenta': '\033[95m',
+              'bright_cyan': '\033[96m',
+              'bright_white': '\033[97m',
+              'end': '\033[0m',  # misc
+              'bold': '\033[1m',
+              'underline': '\033[4m'}
+    return ''.join(colors[x] for x in args) + f'{string}' + colors['end']
+
+
+def labels_to_class_weights(labels, nc=80):
+    # Get class weights (inverse frequency) from training labels
+    if labels[0] is None:  # no labels loaded
+        return torch.Tensor()
+
+    labels = np.concatenate(labels, 0)  # labels.shape = (866643, 5) for COCO
+    classes = labels[:, 0].astype(np.int)  # labels = [class xywh]
+    weights = np.bincount(classes, minlength=nc)  # occurrences per class
+
+    # Prepend gridpoint count (for uCE training)
+    # gpi = ((320 / 32 * np.array([1, 2, 4])) ** 2 * 3).sum()  # gridpoints per image
+    # weights = np.hstack([gpi * len(labels)  - weights.sum() * 9, weights * 9]) ** 0.5  # prepend gridpoints to start
+
+    weights[weights == 0] = 1  # replace empty bins with 1
+    weights = 1 / weights  # number of targets per class
+    weights /= weights.sum()  # normalize
+    return torch.from_numpy(weights)
+
+
+def labels_to_image_weights(labels, nc=80, class_weights=np.ones(80)):
+    # Produces image weights based on class_weights and image contents
+    class_counts = np.array([np.bincount(x[:, 0].astype(np.int), minlength=nc) for x in labels])
+    image_weights = (class_weights.reshape(1, nc) * class_counts).sum(1)
+    # index = random.choices(range(n), weights=image_weights, k=1)  # weight image sample
+    return image_weights
+
+
+def coco80_to_coco91_class():  # converts 80-index (val2014) to 91-index (paper)
+    # https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/
+    # a = np.loadtxt('data/coco.names', dtype='str', delimiter='\n')
+    # b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n')
+    # x1 = [list(a[i] == b).index(True) + 1 for i in range(80)]  # darknet to coco
+    # x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)]  # coco to darknet
+    x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34,
+         35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
+         64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]
+    return x
+
+
+def xyxy2xywh(x):
+    # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[:, 0] = (x[:, 0] + x[:, 2]) / 2  # x center
+    y[:, 1] = (x[:, 1] + x[:, 3]) / 2  # y center
+    y[:, 2] = x[:, 2] - x[:, 0]  # width
+    y[:, 3] = x[:, 3] - x[:, 1]  # height
+    return y
+
+
+def xywh2xyxy(x):
+    # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
+    y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
+    y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
+    y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
+    return y
+
+
+def xywhn2xyxy(x, w=640, h=640, padw=0, padh=0):
+    # Convert nx4 boxes from [x, y, w, h] normalized to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[:, 0] = w * (x[:, 0] - x[:, 2] / 2) + padw  # top left x
+    y[:, 1] = h * (x[:, 1] - x[:, 3] / 2) + padh  # top left y
+    y[:, 2] = w * (x[:, 0] + x[:, 2] / 2) + padw  # bottom right x
+    y[:, 3] = h * (x[:, 1] + x[:, 3] / 2) + padh  # bottom right y
+    return y
+
+
+def xyn2xy(x, w=640, h=640, padw=0, padh=0):
+    # Convert normalized segments into pixel segments, shape (n,2)
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[:, 0] = w * x[:, 0] + padw  # top left x
+    y[:, 1] = h * x[:, 1] + padh  # top left y
+    return y
+
+
+def segment2box(segment, width=640, height=640):
+    # Convert 1 segment label to 1 box label, applying inside-image constraint, i.e. (xy1, xy2, ...) to (xyxy)
+    x, y = segment.T  # segment xy
+    inside = (x >= 0) & (y >= 0) & (x <= width) & (y <= height)
+    x, y, = x[inside], y[inside]
+    return np.array([x.min(), y.min(), x.max(), y.max()]) if any(x) else np.zeros((1, 4))  # cls, xyxy
+
+
+def segments2boxes(segments):
+    # Convert segment labels to box labels, i.e. (cls, xy1, xy2, ...) to (cls, xywh)
+    boxes = []
+    for s in segments:
+        x, y = s.T  # segment xy
+        boxes.append([x.min(), y.min(), x.max(), y.max()])  # cls, xyxy
+    return xyxy2xywh(np.array(boxes))  # cls, xywh
+
+
+def resample_segments(segments, n=1000):
+    # Up-sample an (n,2) segment
+    for i, s in enumerate(segments):
+        x = np.linspace(0, len(s) - 1, n)
+        xp = np.arange(len(s))
+        segments[i] = np.concatenate([np.interp(x, xp, s[:, i]) for i in range(2)]).reshape(2, -1).T  # segment xy
+    return segments
+
+
+def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
+    # Rescale coords (xyxy) from img1_shape to img0_shape
+    if ratio_pad is None:  # calculate from img0_shape
+        gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new
+        pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding
+    else:
+        gain = ratio_pad[0][0]
+        pad = ratio_pad[1]
+
+    coords[:, [0, 2]] -= pad[0]  # x padding
+    coords[:, [1, 3]] -= pad[1]  # y padding
+    coords[:, :4] /= gain
+    clip_coords(coords, img0_shape)
+    return coords
+
+
+def clip_coords(boxes, img_shape):
+    # Clip bounding xyxy bounding boxes to image shape (height, width)
+    boxes[:, 0].clamp_(0, img_shape[1])  # x1
+    boxes[:, 1].clamp_(0, img_shape[0])  # y1
+    boxes[:, 2].clamp_(0, img_shape[1])  # x2
+    boxes[:, 3].clamp_(0, img_shape[0])  # y2
+
+
+def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7):
+    # Returns the IoU of box1 to box2. box1 is 4, box2 is nx4
+    box2 = box2.T
+
+    # Get the coordinates of bounding boxes
+    if x1y1x2y2:  # x1, y1, x2, y2 = box1
+        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]
+    else:  # transform from xywh to xyxy
+        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
+
+    # Intersection area
+    inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
+            (torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)
+
+    # Union Area
+    w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
+    w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
+    union = w1 * h1 + w2 * h2 - inter + eps
+
+    iou = inter / union
+    if GIoU or DIoU or CIoU:
+        cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1)  # convex (smallest enclosing box) width
+        ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1)  # convex height
+        if CIoU or DIoU:  # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
+            c2 = cw ** 2 + ch ** 2 + eps  # convex diagonal squared
+            rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 +
+                    (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4  # center distance squared
+            if DIoU:
+                return iou - rho2 / c2  # DIoU
+            elif CIoU:  # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
+                v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
+                with torch.no_grad():
+                    alpha = v / (v - iou + (1 + eps))
+                return iou - (rho2 / c2 + v * alpha)  # CIoU
+        else:  # GIoU https://arxiv.org/pdf/1902.09630.pdf
+            c_area = cw * ch + eps  # convex area
+            return iou - (c_area - union) / c_area  # GIoU
+    else:
+        return iou  # IoU
+
+
+def box_iou(box1, box2):
+    # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
+    """
+    Return intersection-over-union (Jaccard index) of boxes.
+    Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
+    Arguments:
+        box1 (Tensor[N, 4])
+        box2 (Tensor[M, 4])
+    Returns:
+        iou (Tensor[N, M]): the NxM matrix containing the pairwise
+            IoU values for every element in boxes1 and boxes2
+    """
+
+    def box_area(box):
+        # box = 4xn
+        return (box[2] - box[0]) * (box[3] - box[1])
+
+    area1 = box_area(box1.T)
+    area2 = box_area(box2.T)
+
+    # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
+    inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
+    return inter / (area1[:, None] + area2 - inter)  # iou = inter / (area1 + area2 - inter)
+
+
+def wh_iou(wh1, wh2):
+    # Returns the nxm IoU matrix. wh1 is nx2, wh2 is mx2
+    wh1 = wh1[:, None]  # [N,1,2]
+    wh2 = wh2[None]  # [1,M,2]
+    inter = torch.min(wh1, wh2).prod(2)  # [N,M]
+    return inter / (wh1.prod(2) + wh2.prod(2) - inter)  # iou = inter / (area1 + area2 - inter)
+
+
+def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
+                        labels=()):
+    """Runs Non-Maximum Suppression (NMS) on inference results
+
+    Returns:
+         list of detections, on (n,6) tensor per image [xyxy, conf, cls]
+    """
+
+    nc = prediction.shape[2] - 5  # number of classes
+    xc = prediction[..., 4] > conf_thres  # candidates
+
+    # Settings
+    min_wh, max_wh = 2, 4096  # (pixels) minimum and maximum box width and height
+    max_det = 300  # maximum number of detections per image
+    max_nms = 30000  # maximum number of boxes into torchvision.ops.nms()
+    time_limit = 10.0  # seconds to quit after
+    redundant = True  # require redundant detections
+    multi_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)
+    merge = False  # use merge-NMS
+
+    t = time.time()
+    output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
+    for xi, x in enumerate(prediction):  # image index, image inference
+        # Apply constraints
+        # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0  # width-height
+        x = x[xc[xi]]  # confidence
+
+        # Cat apriori labels if autolabelling
+        if labels and len(labels[xi]):
+            l = labels[xi]
+            v = torch.zeros((len(l), nc + 5), device=x.device)
+            v[:, :4] = l[:, 1:5]  # box
+            v[:, 4] = 1.0  # conf
+            v[range(len(l)), l[:, 0].long() + 5] = 1.0  # cls
+            x = torch.cat((x, v), 0)
+
+        # If none remain process next image
+        if not x.shape[0]:
+            continue
+
+        # Compute conf
+        x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf
+
+        # Box (center x, center y, width, height) to (x1, y1, x2, y2)
+        box = xywh2xyxy(x[:, :4])
+
+        # Detections matrix nx6 (xyxy, conf, cls)
+        if multi_label:
+            i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
+            x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
+        else:  # best class only
+            conf, j = x[:, 5:].max(1, keepdim=True)
+            x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
+
+        # Filter by class
+        if classes is not None:
+            x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
+
+        # Apply finite constraint
+        # if not torch.isfinite(x).all():
+        #     x = x[torch.isfinite(x).all(1)]
+
+        # Check shape
+        n = x.shape[0]  # number of boxes
+        if not n:  # no boxes
+            continue
+        elif n > max_nms:  # excess boxes
+            x = x[x[:, 4].argsort(descending=True)[:max_nms]]  # sort by confidence
+
+        # Batched NMS
+        c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
+        boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
+        i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
+        if i.shape[0] > max_det:  # limit detections
+            i = i[:max_det]
+        if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)
+            # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
+            iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
+            weights = iou * scores[None]  # box weights
+            x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes
+            if redundant:
+                i = i[iou.sum(1) > 1]  # require redundancy
+
+        output[xi] = x[i]
+        if (time.time() - t) > time_limit:
+            print(f'WARNING: NMS time limit {time_limit}s exceeded')
+            break  # time limit exceeded
+
+    return output
+
+
+def strip_optimizer(f='best.pt', s=''):  # from utils.general import *; strip_optimizer()
+    # Strip optimizer from 'f' to finalize training, optionally save as 's'
+    x = torch.load(f, map_location=torch.device('cpu'))
+    if x.get('ema'):
+        x['model'] = x['ema']  # replace model with ema
+    for k in 'optimizer', 'training_results', 'wandb_id', 'ema', 'updates':  # keys
+        x[k] = None
+    x['epoch'] = -1
+    x['model'].half()  # to FP16
+    for p in x['model'].parameters():
+        p.requires_grad = False
+    torch.save(x, s or f)
+    mb = os.path.getsize(s or f) / 1E6  # filesize
+    print(f"Optimizer stripped from {f},{(' saved as %s,' % s) if s else ''} {mb:.1f}MB")
+
+
+def print_mutation(hyp, results, yaml_file='hyp_evolved.yaml', bucket=''):
+    # Print mutation results to evolve.txt (for use with train.py --evolve)
+    a = '%10s' * len(hyp) % tuple(hyp.keys())  # hyperparam keys
+    b = '%10.3g' * len(hyp) % tuple(hyp.values())  # hyperparam values
+    c = '%10.4g' * len(results) % results  # results (P, R, mAP@0.5, mAP@0.5:0.95, val_losses x 3)
+    print('\n%s\n%s\nEvolved fitness: %s\n' % (a, b, c))
+
+    if bucket:
+        url = 'gs://%s/evolve.txt' % bucket
+        if gsutil_getsize(url) > (os.path.getsize('evolve.txt') if os.path.exists('evolve.txt') else 0):
+            os.system('gsutil cp %s .' % url)  # download evolve.txt if larger than local
+
+    with open('evolve.txt', 'a') as f:  # append result
+        f.write(c + b + '\n')
+    x = np.unique(np.loadtxt('evolve.txt', ndmin=2), axis=0)  # load unique rows
+    x = x[np.argsort(-fitness(x))]  # sort
+    np.savetxt('evolve.txt', x, '%10.3g')  # save sort by fitness
+
+    # Save yaml
+    for i, k in enumerate(hyp.keys()):
+        hyp[k] = float(x[0, i + 7])
+    with open(yaml_file, 'w') as f:
+        results = tuple(x[0, :7])
+        c = '%10.4g' * len(results) % results  # results (P, R, mAP@0.5, mAP@0.5:0.95, val_losses x 3)
+        f.write('# Hyperparameter Evolution Results\n# Generations: %g\n# Metrics: ' % len(x) + c + '\n\n')
+        yaml.dump(hyp, f, sort_keys=False)
+
+    if bucket:
+        os.system('gsutil cp evolve.txt %s gs://%s' % (yaml_file, bucket))  # upload
+
+
+def apply_classifier(x, model, img, im0):
+    # applies a second stage classifier to yolo outputs
+    im0 = [im0] if isinstance(im0, np.ndarray) else im0
+    for i, d in enumerate(x):  # per image
+        if d is not None and len(d):
+            d = d.clone()
+
+            # Reshape and pad cutouts
+            b = xyxy2xywh(d[:, :4])  # boxes
+            b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1)  # rectangle to square
+            b[:, 2:] = b[:, 2:] * 1.3 + 30  # pad
+            d[:, :4] = xywh2xyxy(b).long()
+
+            # Rescale boxes from img_size to im0 size
+            scale_coords(img.shape[2:], d[:, :4], im0[i].shape)
+
+            # Classes
+            pred_cls1 = d[:, 5].long()
+            ims = []
+            for j, a in enumerate(d):  # per item
+                cutout = im0[i][int(a[1]):int(a[3]), int(a[0]):int(a[2])]
+                im = cv2.resize(cutout, (224, 224))  # BGR
+                # cv2.imwrite('test%i.jpg' % j, cutout)
+
+                im = im[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
+                im = np.ascontiguousarray(im, dtype=np.float32)  # uint8 to float32
+                im /= 255.0  # 0 - 255 to 0.0 - 1.0
+                ims.append(im)
+
+            pred_cls2 = model(torch.Tensor(ims).to(d.device)).argmax(1)  # classifier prediction
+            x[i] = x[i][pred_cls1 == pred_cls2]  # retain matching class detections
+
+    return x
+
+
+def increment_path(path, exist_ok=True, sep=''):
+    # Increment path, i.e. runs/exp --> runs/exp{sep}0, runs/exp{sep}1 etc.
+    path = Path(path)  # os-agnostic
+    if (path.exists() and exist_ok) or (not path.exists()):
+        return str(path)
+    else:
+        dirs = glob.glob(f"{path}{sep}*")  # similar paths
+        matches = [re.search(rf"%s{sep}(\d+)" % path.stem, d) for d in dirs]
+        i = [int(m.groups()[0]) for m in matches if m]  # indices
+        n = max(i) + 1 if i else 2  # increment number
+        return f"{path}{sep}{n}"  # update path

Y5/utils/google_app_engine/Dockerfile

@@ -0,0 +1,25 @@
+FROM gcr.io/google-appengine/python
+
+# Create a virtualenv for dependencies. This isolates these packages from
+# system-level packages.
+# Use -p python3 or -p python3.7 to select python version. Default is version 2.
+RUN virtualenv /env -p python3
+
+# Setting these environment variables are the same as running
+# source /env/bin/activate.
+ENV VIRTUAL_ENV /env
+ENV PATH /env/bin:$PATH
+
+RUN apt-get update && apt-get install -y python-opencv
+
+# Copy the application's requirements.txt and run pip to install all
+# dependencies into the virtualenv.
+ADD requirements.txt /app/requirements.txt
+RUN pip install -r /app/requirements.txt
+
+# Add the application source code.
+ADD . /app
+
+# Run a WSGI server to serve the application. gunicorn must be declared as
+# a dependency in requirements.txt.
+CMD gunicorn -b :$PORT main:app

Y5/utils/google_app_engine/additional_requirements.txt

@@ -0,0 +1,4 @@
+# add these requirements in your app on top of the existing ones
+pip==18.1
+Flask==1.0.2
+gunicorn==19.9.0

Y5/utils/google_app_engine/app.yaml

@@ -0,0 +1,14 @@
+runtime: custom
+env: flex
+
+service: yolov5app
+
+liveness_check:
+  initial_delay_sec: 600
+
+manual_scaling:
+  instances: 1
+resources:
+  cpu: 1
+  memory_gb: 4
+  disk_size_gb: 20

Y5/utils/google_utils.py

@@ -0,0 +1,122 @@
+# Google utils: https://cloud.google.com/storage/docs/reference/libraries
+
+import os
+import platform
+import subprocess
+import time
+from pathlib import Path
+
+import requests
+import torch
+
+
+def gsutil_getsize(url=''):
+    # gs://bucket/file size https://cloud.google.com/storage/docs/gsutil/commands/du
+    s = subprocess.check_output(f'gsutil du {url}', shell=True).decode('utf-8')
+    return eval(s.split(' ')[0]) if len(s) else 0  # bytes
+
+
+def attempt_download(file, repo='ultralytics/yolov5'):
+    # Attempt file download if does not exist
+    file = Path(str(file).strip().replace("'", ''))
+
+    if not file.exists():
+        try:
+            response = requests.get(f'https://api.github.com/repos/{repo}/releases/latest').json()  # github api
+            assets = [x['name'] for x in response['assets']]  # release assets, i.e. ['yolov5s.pt', 'yolov5m.pt', ...]
+            tag = response['tag_name']  # i.e. 'v1.0'
+        except:  # fallback plan
+            assets = ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']
+            tag = subprocess.check_output('git tag', shell=True).decode().split()[-1]
+
+        name = file.name
+        if name in assets:
+            msg = f'{file} missing, try downloading from https://github.com/{repo}/releases/'
+            redundant = False  # second download option
+            try:  # GitHub
+                url = f'https://github.com/{repo}/releases/download/{tag}/{name}'
+                print(f'Downloading {url} to {file}...')
+                torch.hub.download_url_to_file(url, file)
+                assert file.exists() and file.stat().st_size > 1E6  # check
+            except Exception as e:  # GCP
+                print(f'Download error: {e}')
+                assert redundant, 'No secondary mirror'
+                url = f'https://storage.googleapis.com/{repo}/ckpt/{name}'
+                print(f'Downloading {url} to {file}...')
+                os.system(f'curl -L {url} -o {file}')  # torch.hub.download_url_to_file(url, weights)
+            finally:
+                if not file.exists() or file.stat().st_size < 1E6:  # check
+                    file.unlink(missing_ok=True)  # remove partial downloads
+                    print(f'ERROR: Download failure: {msg}')
+                print('')
+                return
+
+
+def gdrive_download(id='16TiPfZj7htmTyhntwcZyEEAejOUxuT6m', file='tmp.zip'):
+    # Downloads a file from Google Drive. from yolov5.utils.google_utils import *; gdrive_download()
+    t = time.time()
+    file = Path(file)
+    cookie = Path('cookie')  # gdrive cookie
+    print(f'Downloading https://drive.google.com/uc?export=download&id={id} as {file}... ', end='')
+    file.unlink(missing_ok=True)  # remove existing file
+    cookie.unlink(missing_ok=True)  # remove existing cookie
+
+    # Attempt file download
+    out = "NUL" if platform.system() == "Windows" else "/dev/null"
+    os.system(f'curl -c ./cookie -s -L "drive.google.com/uc?export=download&id={id}" > {out}')
+    if os.path.exists('cookie'):  # large file
+        s = f'curl -Lb ./cookie "drive.google.com/uc?export=download&confirm={get_token()}&id={id}" -o {file}'
+    else:  # small file
+        s = f'curl -s -L -o {file} "drive.google.com/uc?export=download&id={id}"'
+    r = os.system(s)  # execute, capture return
+    cookie.unlink(missing_ok=True)  # remove existing cookie
+
+    # Error check
+    if r != 0:
+        file.unlink(missing_ok=True)  # remove partial
+        print('Download error ')  # raise Exception('Download error')
+        return r
+
+    # Unzip if archive
+    if file.suffix == '.zip':
+        print('unzipping... ', end='')
+        os.system(f'unzip -q {file}')  # unzip
+        file.unlink()  # remove zip to free space
+
+    print(f'Done ({time.time() - t:.1f}s)')
+    return r
+
+
+def get_token(cookie="./cookie"):
+    with open(cookie) as f:
+        for line in f:
+            if "download" in line:
+                return line.split()[-1]
+    return ""
+
+# def upload_blob(bucket_name, source_file_name, destination_blob_name):
+#     # Uploads a file to a bucket
+#     # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python
+#
+#     storage_client = storage.Client()
+#     bucket = storage_client.get_bucket(bucket_name)
+#     blob = bucket.blob(destination_blob_name)
+#
+#     blob.upload_from_filename(source_file_name)
+#
+#     print('File {} uploaded to {}.'.format(
+#         source_file_name,
+#         destination_blob_name))
+#
+#
+# def download_blob(bucket_name, source_blob_name, destination_file_name):
+#     # Uploads a blob from a bucket
+#     storage_client = storage.Client()
+#     bucket = storage_client.get_bucket(bucket_name)
+#     blob = bucket.blob(source_blob_name)
+#
+#     blob.download_to_filename(destination_file_name)
+#
+#     print('Blob {} downloaded to {}.'.format(
+#         source_blob_name,
+#         destination_file_name))

Y5/utils/loss.py

@@ -0,0 +1,216 @@
+# Loss functions
+
+import torch
+import torch.nn as nn
+
+from Y5.utils.general import bbox_iou
+from Y5.utils.torch_utils import is_parallel
+
+
+def smooth_BCE(eps=0.1):  # https://github.com/ultralytics/yolov3/issues/238#issuecomment-598028441
+    # return positive, negative label smoothing BCE targets
+    return 1.0 - 0.5 * eps, 0.5 * eps
+
+
+class BCEBlurWithLogitsLoss(nn.Module):
+    # BCEwithLogitLoss() with reduced missing label effects.
+    def __init__(self, alpha=0.05):
+        super(BCEBlurWithLogitsLoss, self).__init__()
+        self.loss_fcn = nn.BCEWithLogitsLoss(reduction='none')  # must be nn.BCEWithLogitsLoss()
+        self.alpha = alpha
+
+    def forward(self, pred, true):
+        loss = self.loss_fcn(pred, true)
+        pred = torch.sigmoid(pred)  # prob from logits
+        dx = pred - true  # reduce only missing label effects
+        # dx = (pred - true).abs()  # reduce missing label and false label effects
+        alpha_factor = 1 - torch.exp((dx - 1) / (self.alpha + 1e-4))
+        loss *= alpha_factor
+        return loss.mean()
+
+
+class FocalLoss(nn.Module):
+    # Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
+    def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
+        super(FocalLoss, self).__init__()
+        self.loss_fcn = loss_fcn  # must be nn.BCEWithLogitsLoss()
+        self.gamma = gamma
+        self.alpha = alpha
+        self.reduction = loss_fcn.reduction
+        self.loss_fcn.reduction = 'none'  # required to apply FL to each element
+
+    def forward(self, pred, true):
+        loss = self.loss_fcn(pred, true)
+        # p_t = torch.exp(-loss)
+        # loss *= self.alpha * (1.000001 - p_t) ** self.gamma  # non-zero power for gradient stability
+
+        # TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py
+        pred_prob = torch.sigmoid(pred)  # prob from logits
+        p_t = true * pred_prob + (1 - true) * (1 - pred_prob)
+        alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
+        modulating_factor = (1.0 - p_t) ** self.gamma
+        loss *= alpha_factor * modulating_factor
+
+        if self.reduction == 'mean':
+            return loss.mean()
+        elif self.reduction == 'sum':
+            return loss.sum()
+        else:  # 'none'
+            return loss
+
+
+class QFocalLoss(nn.Module):
+    # Wraps Quality focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
+    def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
+        super(QFocalLoss, self).__init__()
+        self.loss_fcn = loss_fcn  # must be nn.BCEWithLogitsLoss()
+        self.gamma = gamma
+        self.alpha = alpha
+        self.reduction = loss_fcn.reduction
+        self.loss_fcn.reduction = 'none'  # required to apply FL to each element
+
+    def forward(self, pred, true):
+        loss = self.loss_fcn(pred, true)
+
+        pred_prob = torch.sigmoid(pred)  # prob from logits
+        alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
+        modulating_factor = torch.abs(true - pred_prob) ** self.gamma
+        loss *= alpha_factor * modulating_factor
+
+        if self.reduction == 'mean':
+            return loss.mean()
+        elif self.reduction == 'sum':
+            return loss.sum()
+        else:  # 'none'
+            return loss
+
+
+class ComputeLoss:
+    # Compute losses
+    def __init__(self, model, autobalance=False):
+        super(ComputeLoss, self).__init__()
+        device = next(model.parameters()).device  # get model device
+        h = model.hyp  # hyperparameters
+
+        # Define criteria
+        BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device))
+        BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))
+
+        # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
+        self.cp, self.cn = smooth_BCE(eps=0.0)
+
+        # Focal loss
+        g = h['fl_gamma']  # focal loss gamma
+        if g > 0:
+            BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
+
+        det = model.module.model[-1] if is_parallel(model) else model.model[-1]  # Detect() module
+        self.balance = {3: [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.25, 0.06, .02])  # P3-P7
+        self.ssi = list(det.stride).index(16) if autobalance else 0  # stride 16 index
+        self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, model.gr, h, autobalance
+        for k in 'na', 'nc', 'nl', 'anchors':
+            setattr(self, k, getattr(det, k))
+
+    def __call__(self, p, targets):  # predictions, targets, model
+        device = targets.device
+        lcls, lbox, lobj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device)
+        tcls, tbox, indices, anchors = self.build_targets(p, targets)  # targets
+
+        # Losses
+        for i, pi in enumerate(p):  # layer index, layer predictions
+            b, a, gj, gi = indices[i]  # image, anchor, gridy, gridx
+            tobj = torch.zeros_like(pi[..., 0], device=device)  # target obj
+
+            n = b.shape[0]  # number of targets
+            if n:
+                ps = pi[b, a, gj, gi]  # prediction subset corresponding to targets
+
+                # Regression
+                pxy = ps[:, :2].sigmoid() * 2. - 0.5
+                pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]
+                pbox = torch.cat((pxy, pwh), 1)  # predicted box
+                iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True)  # iou(prediction, target)
+                lbox += (1.0 - iou).mean()  # iou loss
+
+                # Objectness
+                tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * iou.detach().clamp(0).type(tobj.dtype)  # iou ratio
+
+                # Classification
+                if self.nc > 1:  # cls loss (only if multiple classes)
+                    t = torch.full_like(ps[:, 5:], self.cn, device=device)  # targets
+                    t[range(n), tcls[i]] = self.cp
+                    lcls += self.BCEcls(ps[:, 5:], t)  # BCE
+
+                # Append targets to text file
+                # with open('targets.txt', 'a') as file:
+                #     [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]
+
+            obji = self.BCEobj(pi[..., 4], tobj)
+            lobj += obji * self.balance[i]  # obj loss
+            if self.autobalance:
+                self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item()
+
+        if self.autobalance:
+            self.balance = [x / self.balance[self.ssi] for x in self.balance]
+        lbox *= self.hyp['box']
+        lobj *= self.hyp['obj']
+        lcls *= self.hyp['cls']
+        bs = tobj.shape[0]  # batch size
+
+        loss = lbox + lobj + lcls
+        return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach()
+
+    def build_targets(self, p, targets):
+        # Build targets for compute_loss(), input targets(image,class,x,y,w,h)
+        na, nt = self.na, targets.shape[0]  # number of anchors, targets
+        tcls, tbox, indices, anch = [], [], [], []
+        gain = torch.ones(7, device=targets.device)  # normalized to gridspace gain
+        ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt)  # same as .repeat_interleave(nt)
+        targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2)  # append anchor indices
+
+        g = 0.5  # bias
+        off = torch.tensor([[0, 0],
+                            [1, 0], [0, 1], [-1, 0], [0, -1],  # j,k,l,m
+                            # [1, 1], [1, -1], [-1, 1], [-1, -1],  # jk,jm,lk,lm
+                            ], device=targets.device).float() * g  # offsets
+
+        for i in range(self.nl):
+            anchors = self.anchors[i]
+            gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]]  # xyxy gain
+
+            # Match targets to anchors
+            t = targets * gain
+            if nt:
+                # Matches
+                r = t[:, :, 4:6] / anchors[:, None]  # wh ratio
+                j = torch.max(r, 1. / r).max(2)[0] < self.hyp['anchor_t']  # compare
+                # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t']  # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
+                t = t[j]  # filter
+
+                # Offsets
+                gxy = t[:, 2:4]  # grid xy
+                gxi = gain[[2, 3]] - gxy  # inverse
+                j, k = ((gxy % 1. < g) & (gxy > 1.)).T
+                l, m = ((gxi % 1. < g) & (gxi > 1.)).T
+                j = torch.stack((torch.ones_like(j), j, k, l, m))
+                t = t.repeat((5, 1, 1))[j]
+                offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
+            else:
+                t = targets[0]
+                offsets = 0
+
+            # Define
+            b, c = t[:, :2].long().T  # image, class
+            gxy = t[:, 2:4]  # grid xy
+            gwh = t[:, 4:6]  # grid wh
+            gij = (gxy - offsets).long()
+            gi, gj = gij.T  # grid xy indices
+
+            # Append
+            a = t[:, 6].long()  # anchor indices
+            indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1)))  # image, anchor, grid indices
+            tbox.append(torch.cat((gxy - gij, gwh), 1))  # box
+            anch.append(anchors[a])  # anchors
+            tcls.append(c)  # class
+
+        return tcls, tbox, indices, anch

Y5/utils/metrics.py

@@ -0,0 +1,223 @@
+# Model validation metrics
+
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from . import general
+
+
+def fitness(x):
+    # Model fitness as a weighted combination of metrics
+    w = [0.0, 0.0, 0.1, 0.9]  # weights for [P, R, mAP@0.5, mAP@0.5:0.95]
+    return (x[:, :4] * w).sum(1)
+
+
+def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names=()):
+    """ Compute the average precision, given the recall and precision curves.
+    Source: https://github.com/rafaelpadilla/Object-Detection-Metrics.
+    # Arguments
+        tp:  True positives (nparray, nx1 or nx10).
+        conf:  Objectness value from 0-1 (nparray).
+        pred_cls:  Predicted object classes (nparray).
+        target_cls:  True object classes (nparray).
+        plot:  Plot precision-recall curve at mAP@0.5
+        save_dir:  Plot save directory
+    # Returns
+        The average precision as computed in py-faster-rcnn.
+    """
+
+    # Sort by objectness
+    i = np.argsort(-conf)
+    tp, conf, pred_cls = tp[i], conf[i], pred_cls[i]
+
+    # Find unique classes
+    unique_classes = np.unique(target_cls)
+    nc = unique_classes.shape[0]  # number of classes, number of detections
+
+    # Create Precision-Recall curve and compute AP for each class
+    px, py = np.linspace(0, 1, 1000), []  # for plotting
+    ap, p, r = np.zeros((nc, tp.shape[1])), np.zeros((nc, 1000)), np.zeros((nc, 1000))
+    for ci, c in enumerate(unique_classes):
+        i = pred_cls == c
+        n_l = (target_cls == c).sum()  # number of labels
+        n_p = i.sum()  # number of predictions
+
+        if n_p == 0 or n_l == 0:
+            continue
+        else:
+            # Accumulate FPs and TPs
+            fpc = (1 - tp[i]).cumsum(0)
+            tpc = tp[i].cumsum(0)
+
+            # Recall
+            recall = tpc / (n_l + 1e-16)  # recall curve
+            r[ci] = np.interp(-px, -conf[i], recall[:, 0], left=0)  # negative x, xp because xp decreases
+
+            # Precision
+            precision = tpc / (tpc + fpc)  # precision curve
+            p[ci] = np.interp(-px, -conf[i], precision[:, 0], left=1)  # p at pr_score
+
+            # AP from recall-precision curve
+            for j in range(tp.shape[1]):
+                ap[ci, j], mpre, mrec = compute_ap(recall[:, j], precision[:, j])
+                if plot and j == 0:
+                    py.append(np.interp(px, mrec, mpre))  # precision at mAP@0.5
+
+    # Compute F1 (harmonic mean of precision and recall)
+    f1 = 2 * p * r / (p + r + 1e-16)
+    if plot:
+        plot_pr_curve(px, py, ap, Path(save_dir) / 'PR_curve.png', names)
+        plot_mc_curve(px, f1, Path(save_dir) / 'F1_curve.png', names, ylabel='F1')
+        plot_mc_curve(px, p, Path(save_dir) / 'P_curve.png', names, ylabel='Precision')
+        plot_mc_curve(px, r, Path(save_dir) / 'R_curve.png', names, ylabel='Recall')
+
+    i = f1.mean(0).argmax()  # max F1 index
+    return p[:, i], r[:, i], ap, f1[:, i], unique_classes.astype('int32')
+
+
+def compute_ap(recall, precision):
+    """ Compute the average precision, given the recall and precision curves
+    # Arguments
+        recall:    The recall curve (list)
+        precision: The precision curve (list)
+    # Returns
+        Average precision, precision curve, recall curve
+    """
+
+    # Append sentinel values to beginning and end
+    mrec = np.concatenate(([0.], recall, [recall[-1] + 0.01]))
+    mpre = np.concatenate(([1.], precision, [0.]))
+
+    # Compute the precision envelope
+    mpre = np.flip(np.maximum.accumulate(np.flip(mpre)))
+
+    # Integrate area under curve
+    method = 'interp'  # methods: 'continuous', 'interp'
+    if method == 'interp':
+        x = np.linspace(0, 1, 101)  # 101-point interp (COCO)
+        ap = np.trapz(np.interp(x, mrec, mpre), x)  # integrate
+    else:  # 'continuous'
+        i = np.where(mrec[1:] != mrec[:-1])[0]  # points where x axis (recall) changes
+        ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])  # area under curve
+
+    return ap, mpre, mrec
+
+
+class ConfusionMatrix:
+    # Updated version of https://github.com/kaanakan/object_detection_confusion_matrix
+    def __init__(self, nc, conf=0.25, iou_thres=0.45):
+        self.matrix = np.zeros((nc + 1, nc + 1))
+        self.nc = nc  # number of classes
+        self.conf = conf
+        self.iou_thres = iou_thres
+
+    def process_batch(self, detections, labels):
+        """
+        Return intersection-over-union (Jaccard index) of boxes.
+        Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
+        Arguments:
+            detections (Array[N, 6]), x1, y1, x2, y2, conf, class
+            labels (Array[M, 5]), class, x1, y1, x2, y2
+        Returns:
+            None, updates confusion matrix accordingly
+        """
+        detections = detections[detections[:, 4] > self.conf]
+        gt_classes = labels[:, 0].int()
+        detection_classes = detections[:, 5].int()
+        iou = general.box_iou(labels[:, 1:], detections[:, :4])
+
+        x = torch.where(iou > self.iou_thres)
+        if x[0].shape[0]:
+            matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy()
+            if x[0].shape[0] > 1:
+                matches = matches[matches[:, 2].argsort()[::-1]]
+                matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
+                matches = matches[matches[:, 2].argsort()[::-1]]
+                matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
+        else:
+            matches = np.zeros((0, 3))
+
+        n = matches.shape[0] > 0
+        m0, m1, _ = matches.transpose().astype(np.int16)
+        for i, gc in enumerate(gt_classes):
+            j = m0 == i
+            if n and sum(j) == 1:
+                self.matrix[gc, detection_classes[m1[j]]] += 1  # correct
+            else:
+                self.matrix[self.nc, gc] += 1  # background FP
+
+        if n:
+            for i, dc in enumerate(detection_classes):
+                if not any(m1 == i):
+                    self.matrix[dc, self.nc] += 1  # background FN
+
+    def matrix(self):
+        return self.matrix
+
+    def plot(self, save_dir='', names=()):
+        try:
+            import seaborn as sn
+
+            array = self.matrix / (self.matrix.sum(0).reshape(1, self.nc + 1) + 1E-6)  # normalize
+            array[array < 0.005] = np.nan  # don't annotate (would appear as 0.00)
+
+            fig = plt.figure(figsize=(12, 9), tight_layout=True)
+            sn.set(font_scale=1.0 if self.nc < 50 else 0.8)  # for label size
+            labels = (0 < len(names) < 99) and len(names) == self.nc  # apply names to ticklabels
+            sn.heatmap(array, annot=self.nc < 30, annot_kws={"size": 8}, cmap='Blues', fmt='.2f', square=True,
+                       xticklabels=names + ['background FP'] if labels else "auto",
+                       yticklabels=names + ['background FN'] if labels else "auto").set_facecolor((1, 1, 1))
+            fig.axes[0].set_xlabel('True')
+            fig.axes[0].set_ylabel('Predicted')
+            fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250)
+        except Exception as e:
+            pass
+
+    def print(self):
+        for i in range(self.nc + 1):
+            print(' '.join(map(str, self.matrix[i])))
+
+
+# Plots ----------------------------------------------------------------------------------------------------------------
+
+def plot_pr_curve(px, py, ap, save_dir='pr_curve.png', names=()):
+    # Precision-recall curve
+    fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True)
+    py = np.stack(py, axis=1)
+
+    if 0 < len(names) < 21:  # display per-class legend if < 21 classes
+        for i, y in enumerate(py.T):
+            ax.plot(px, y, linewidth=1, label=f'{names[i]} {ap[i, 0]:.3f}')  # plot(recall, precision)
+    else:
+        ax.plot(px, py, linewidth=1, color='grey')  # plot(recall, precision)
+
+    ax.plot(px, py.mean(1), linewidth=3, color='blue', label='all classes %.3f mAP@0.5' % ap[:, 0].mean())
+    ax.set_xlabel('Recall')
+    ax.set_ylabel('Precision')
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
+    fig.savefig(Path(save_dir), dpi=250)
+
+
+def plot_mc_curve(px, py, save_dir='mc_curve.png', names=(), xlabel='Confidence', ylabel='Metric'):
+    # Metric-confidence curve
+    fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True)
+
+    if 0 < len(names) < 21:  # display per-class legend if < 21 classes
+        for i, y in enumerate(py):
+            ax.plot(px, y, linewidth=1, label=f'{names[i]}')  # plot(confidence, metric)
+    else:
+        ax.plot(px, py.T, linewidth=1, color='grey')  # plot(confidence, metric)
+
+    y = py.mean(0)
+    ax.plot(px, y, linewidth=3, color='blue', label=f'all classes {y.max():.2f} at {px[y.argmax()]:.3f}')
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel(ylabel)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
+    fig.savefig(Path(save_dir), dpi=250)

Y5/utils/plots.py

@@ -0,0 +1,429 @@
+# Plotting utils
+# !usr/bin/env python
+# -*- coding:utf-8 _*-
+import glob
+import math
+import os
+import random
+from copy import copy
+from pathlib import Path
+
+import cv2
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import seaborn as sns
+import torch
+import yaml
+from PIL import Image, ImageDraw, ImageFont
+from scipy.signal import butter, filtfilt
+
+from Y5.utils.general import xywh2xyxy, xyxy2xywh
+from Y5.utils.metrics import fitness
+
+# Settings
+matplotlib.rc('font', **{'size': 11})
+matplotlib.use('Agg')  # for writing to files only
+
+
+def color_list():
+    # Return first 10 plt colors as (r,g,b) https://stackoverflow.com/questions/51350872/python-from-color-name-to-rgb
+    def hex2rgb(h):
+        return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))
+
+    return [hex2rgb(h) for h in matplotlib.colors.TABLEAU_COLORS.values()]  # or BASE_ (8), CSS4_ (148), XKCD_ (949)
+
+
+def hist2d(x, y, n=100):
+    # 2d histogram used in labels.png and evolve.png
+    xedges, yedges = np.linspace(x.min(), x.max(), n), np.linspace(y.min(), y.max(), n)
+    hist, xedges, yedges = np.histogram2d(x, y, (xedges, yedges))
+    xidx = np.clip(np.digitize(x, xedges) - 1, 0, hist.shape[0] - 1)
+    yidx = np.clip(np.digitize(y, yedges) - 1, 0, hist.shape[1] - 1)
+    return np.log(hist[xidx, yidx])
+
+
+def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5):
+    # https://stackoverflow.com/questions/28536191/how-to-filter-smooth-with-scipy-numpy
+    def butter_lowpass(cutoff, fs, order):
+        nyq = 0.5 * fs
+        normal_cutoff = cutoff / nyq
+        return butter(order, normal_cutoff, btype='low', analog=False)
+
+    b, a = butter_lowpass(cutoff, fs, order=order)
+    return filtfilt(b, a, data)  # forward-backward filter
+
+
+def plot_one_box(x, img, color=None, label=None, line_thickness=3):
+    # Plots one bounding box on image img
+    tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1  # line/font thickness
+    color = color or [random.randint(0, 255) for _ in range(3)]
+    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
+    cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
+    if label:
+        tf = max(tl - 1, 1)  # font thickness
+        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
+        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
+        cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filled
+        cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
+
+
+def plot_one_box_PIL(box, img, color=None, label=None, line_thickness=None):
+    img = Image.fromarray(img)
+    draw = ImageDraw.Draw(img)
+    line_thickness = line_thickness or max(int(min(img.size) / 200), 2)
+    draw.rectangle(box, width=line_thickness, outline=tuple(color))  # plot
+    if label:
+        fontsize = max(round(max(img.size) / 40), 12)
+        font = ImageFont.truetype("Arial.ttf", fontsize)
+        txt_width, txt_height = font.getsize(label)
+        draw.rectangle([box[0], box[1] - txt_height + 4, box[0] + txt_width, box[1]], fill=tuple(color))
+        draw.text((box[0], box[1] - txt_height + 1), label, fill=(255, 255, 255), font=font)
+    return np.asarray(img)
+
+
+def plot_wh_methods():  # from utils.plots import *; plot_wh_methods()
+    # Compares the two methods for width-height anchor multiplication
+    # https://github.com/ultralytics/yolov3/issues/168
+    x = np.arange(-4.0, 4.0, .1)
+    ya = np.exp(x)
+    yb = torch.sigmoid(torch.from_numpy(x)).numpy() * 2
+
+    fig = plt.figure(figsize=(6, 3), tight_layout=True)
+    plt.plot(x, ya, '.-', label='YOLOv3')
+    plt.plot(x, yb ** 2, '.-', label='YOLOv5 ^2')
+    plt.plot(x, yb ** 1.6, '.-', label='YOLOv5 ^1.6')
+    plt.xlim(left=-4, right=4)
+    plt.ylim(bottom=0, top=6)
+    plt.xlabel('input')
+    plt.ylabel('output')
+    plt.grid()
+    plt.legend()
+    fig.savefig('comparison.png', dpi=200)
+
+
+def output_to_target(output):
+    # Convert model output to target format [batch_id, class_id, x, y, w, h, conf]
+    targets = []
+    for i, o in enumerate(output):
+        for *box, conf, cls in o.cpu().numpy():
+            targets.append([i, cls, *list(*xyxy2xywh(np.array(box)[None])), conf])
+    return np.array(targets)
+
+
+def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max_size=640, max_subplots=16):
+    # Plot image grid with labels
+
+    if isinstance(images, torch.Tensor):
+        images = images.cpu().float().numpy()
+    if isinstance(targets, torch.Tensor):
+        targets = targets.cpu().numpy()
+
+    # un-normalise
+    if np.max(images[0]) <= 1:
+        images *= 255
+
+    tl = 3  # line thickness
+    tf = max(tl - 1, 1)  # font thickness
+    bs, _, h, w = images.shape  # batch size, _, height, width
+    bs = min(bs, max_subplots)  # limit plot images
+    ns = np.ceil(bs ** 0.5)  # number of subplots (square)
+
+    # Check if we should resize
+    scale_factor = max_size / max(h, w)
+    if scale_factor < 1:
+        h = math.ceil(scale_factor * h)
+        w = math.ceil(scale_factor * w)
+
+    colors = color_list()  # list of colors
+    mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8)  # init
+    for i, img in enumerate(images):
+        if i == max_subplots:  # if last batch has fewer images than we expect
+            break
+
+        block_x = int(w * (i // ns))
+        block_y = int(h * (i % ns))
+
+        img = img.transpose(1, 2, 0)
+        if scale_factor < 1:
+            img = cv2.resize(img, (w, h))
+
+        mosaic[block_y:block_y + h, block_x:block_x + w, :] = img
+        if len(targets) > 0:
+            image_targets = targets[targets[:, 0] == i]
+            boxes = xywh2xyxy(image_targets[:, 2:6]).T
+            classes = image_targets[:, 1].astype('int')
+            labels = image_targets.shape[1] == 6  # labels if no conf column
+            conf = None if labels else image_targets[:, 6]  # check for confidence presence (label vs pred)
+
+            if boxes.shape[1]:
+                if boxes.max() <= 1.01:  # if normalized with tolerance 0.01
+                    boxes[[0, 2]] *= w  # scale to pixels
+                    boxes[[1, 3]] *= h
+                elif scale_factor < 1:  # absolute coords need scale if image scales
+                    boxes *= scale_factor
+            boxes[[0, 2]] += block_x
+            boxes[[1, 3]] += block_y
+            for j, box in enumerate(boxes.T):
+                cls = int(classes[j])
+                color = colors[cls % len(colors)]
+                cls = names[cls] if names else cls
+                if labels or conf[j] > 0.25:  # 0.25 conf thresh
+                    label = '%s' % cls if labels else '%s %.1f' % (cls, conf[j])
+                    plot_one_box(box, mosaic, label=label, color=color, line_thickness=tl)
+
+        # Draw image filename labels
+        if paths:
+            label = Path(paths[i]).name[:40]  # trim to 40 char
+            t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
+            cv2.putText(mosaic, label, (block_x + 5, block_y + t_size[1] + 5), 0, tl / 3, [220, 220, 220], thickness=tf,
+                        lineType=cv2.LINE_AA)
+
+        # Image border
+        cv2.rectangle(mosaic, (block_x, block_y), (block_x + w, block_y + h), (255, 255, 255), thickness=3)
+
+    if fname:
+        r = min(1280. / max(h, w) / ns, 1.0)  # ratio to limit image size
+        mosaic = cv2.resize(mosaic, (int(ns * w * r), int(ns * h * r)), interpolation=cv2.INTER_AREA)
+        # cv2.imwrite(fname, cv2.cvtColor(mosaic, cv2.COLOR_BGR2RGB))  # cv2 save
+        Image.fromarray(mosaic).save(fname)  # PIL save
+    return mosaic
+
+
+def plot_lr_scheduler(optimizer, scheduler, epochs=300, save_dir=''):
+    # Plot LR simulating training for full epochs
+    optimizer, scheduler = copy(optimizer), copy(scheduler)  # do not modify originals
+    y = []
+    for _ in range(epochs):
+        scheduler.step()
+        y.append(optimizer.param_groups[0]['lr'])
+    plt.plot(y, '.-', label='LR')
+    plt.xlabel('epoch')
+    plt.ylabel('LR')
+    plt.grid()
+    plt.xlim(0, epochs)
+    plt.ylim(0)
+    plt.savefig(Path(save_dir) / 'LR.png', dpi=200)
+    plt.close()
+
+
+def plot_test_txt():  # from utils.plots import *; plot_test()
+    # Plot test.txt histograms
+    x = np.loadtxt('test.txt', dtype=np.float32)
+    box = xyxy2xywh(x[:, :4])
+    cx, cy = box[:, 0], box[:, 1]
+
+    fig, ax = plt.subplots(1, 1, figsize=(6, 6), tight_layout=True)
+    ax.hist2d(cx, cy, bins=600, cmax=10, cmin=0)
+    ax.set_aspect('equal')
+    plt.savefig('hist2d.png', dpi=300)
+
+    fig, ax = plt.subplots(1, 2, figsize=(12, 6), tight_layout=True)
+    ax[0].hist(cx, bins=600)
+    ax[1].hist(cy, bins=600)
+    plt.savefig('hist1d.png', dpi=200)
+
+
+def plot_targets_txt():  # from utils.plots import *; plot_targets_txt()
+    # Plot targets.txt histograms
+    x = np.loadtxt('targets.txt', dtype=np.float32).T
+    s = ['x targets', 'y targets', 'width targets', 'height targets']
+    fig, ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)
+    ax = ax.ravel()
+    for i in range(4):
+        ax[i].hist(x[i], bins=100, label='%.3g +/- %.3g' % (x[i].mean(), x[i].std()))
+        ax[i].legend()
+        ax[i].set_title(s[i])
+    plt.savefig('targets.jpg', dpi=200)
+
+
+def plot_study_txt(path='', x=None):  # from utils.plots import *; plot_study_txt()
+    # Plot study.txt generated by test.py
+    fig, ax = plt.subplots(2, 4, figsize=(10, 6), tight_layout=True)
+    # ax = ax.ravel()
+
+    fig2, ax2 = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True)
+    # for f in [Path(path) / f'study_coco_{x}.txt' for x in ['yolov5s', 'yolov5m', 'yolov5l', 'yolov5x']]:
+    for f in sorted(Path(path).glob('study*.txt')):
+        y = np.loadtxt(f, dtype=np.float32, usecols=[0, 1, 2, 3, 7, 8, 9], ndmin=2).T
+        x = np.arange(y.shape[1]) if x is None else np.array(x)
+        s = ['P', 'R', 'mAP@.5', 'mAP@.5:.95', 't_inference (ms/img)', 't_NMS (ms/img)', 't_total (ms/img)']
+        # for i in range(7):
+        #     ax[i].plot(x, y[i], '.-', linewidth=2, markersize=8)
+        #     ax[i].set_title(s[i])
+
+        j = y[3].argmax() + 1
+        ax2.plot(y[6, :j], y[3, :j] * 1E2, '.-', linewidth=2, markersize=8,
+                 label=f.stem.replace('study_coco_', '').replace('yolo', 'YOLO'))
+
+    ax2.plot(1E3 / np.array([209, 140, 97, 58, 35, 18]), [34.6, 40.5, 43.0, 47.5, 49.7, 51.5],
+             'k.-', linewidth=2, markersize=8, alpha=.25, label='EfficientDet')
+
+    ax2.grid(alpha=0.2)
+    ax2.set_yticks(np.arange(20, 60, 5))
+    ax2.set_xlim(0, 30)
+    ax2.set_ylim(30, 55)
+    ax2.set_xlabel('GPU Speed (ms/img)')
+    ax2.set_ylabel('COCO AP val')
+    ax2.legend(loc='lower right')
+    plt.savefig(str(Path(path).name) + '.png', dpi=300)
+
+
+def plot_labels(labels, save_dir=Path(''), loggers=None):
+    # plot dataset labels
+    print('Plotting labels... ')
+    c, b = labels[:, 0], labels[:, 1:].transpose()  # classes, boxes
+    nc = int(c.max() + 1)  # number of classes
+    colors = color_list()
+    x = pd.DataFrame(b.transpose(), columns=['x', 'y', 'width', 'height'])
+
+    # seaborn correlogram
+    sns.pairplot(x, corner=True, diag_kind='auto', kind='hist', diag_kws=dict(bins=50), plot_kws=dict(pmax=0.9))
+    plt.savefig(save_dir / 'labels_correlogram.jpg', dpi=200)
+    plt.close()
+
+    # matplotlib labels
+    matplotlib.use('svg')  # faster
+    ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)[1].ravel()
+    ax[0].hist(c, bins=np.linspace(0, nc, nc + 1) - 0.5, rwidth=0.8)
+    ax[0].set_xlabel('classes')
+    sns.histplot(x, x='x', y='y', ax=ax[2], bins=50, pmax=0.9)
+    sns.histplot(x, x='width', y='height', ax=ax[3], bins=50, pmax=0.9)
+
+    # rectangles
+    labels[:, 1:3] = 0.5  # center
+    labels[:, 1:] = xywh2xyxy(labels[:, 1:]) * 2000
+    img = Image.fromarray(np.ones((2000, 2000, 3), dtype=np.uint8) * 255)
+    for cls, *box in labels[:1000]:
+        ImageDraw.Draw(img).rectangle(box, width=1, outline=colors[int(cls) % 10])  # plot
+    ax[1].imshow(img)
+    ax[1].axis('off')
+
+    for a in [0, 1, 2, 3]:
+        for s in ['top', 'right', 'left', 'bottom']:
+            ax[a].spines[s].set_visible(False)
+
+    plt.savefig(save_dir / 'labels.jpg', dpi=200)
+    matplotlib.use('Agg')
+    plt.close()
+
+    # loggers
+    for k, v in loggers.items() or {}:
+        if k == 'wandb' and v:
+            v.log({"Labels": [v.Image(str(x), caption=x.name) for x in save_dir.glob('*labels*.jpg')]}, commit=False)
+
+
+def plot_evolution(yaml_file='data/hyp.finetune.yaml'):  # from utils.plots import *; plot_evolution()
+    # Plot hyperparameter evolution results in evolve.txt
+    with open(yaml_file) as f:
+        hyp = yaml.load(f, Loader=yaml.SafeLoader)
+    x = np.loadtxt('evolve.txt', ndmin=2)
+    f = fitness(x)
+    # weights = (f - f.min()) ** 2  # for weighted results
+    plt.figure(figsize=(10, 12), tight_layout=True)
+    matplotlib.rc('font', **{'size': 8})
+    for i, (k, v) in enumerate(hyp.items()):
+        y = x[:, i + 7]
+        # mu = (y * weights).sum() / weights.sum()  # best weighted result
+        mu = y[f.argmax()]  # best single result
+        plt.subplot(6, 5, i + 1)
+        plt.scatter(y, f, c=hist2d(y, f, 20), cmap='viridis', alpha=.8, edgecolors='none')
+        plt.plot(mu, f.max(), 'k+', markersize=15)
+        plt.title('%s = %.3g' % (k, mu), fontdict={'size': 9})  # limit to 40 characters
+        if i % 5 != 0:
+            plt.yticks([])
+        print('%15s: %.3g' % (k, mu))
+    plt.savefig('evolve.png', dpi=200)
+    print('\nPlot saved as evolve.png')
+
+
+def profile_idetection(start=0, stop=0, labels=(), save_dir=''):
+    # Plot iDetection '*.txt' per-image logs. from utils.plots import *; profile_idetection()
+    ax = plt.subplots(2, 4, figsize=(12, 6), tight_layout=True)[1].ravel()
+    s = ['Images', 'Free Storage (GB)', 'RAM Usage (GB)', 'Battery', 'dt_raw (ms)', 'dt_smooth (ms)', 'real-world FPS']
+    files = list(Path(save_dir).glob('frames*.txt'))
+    for fi, f in enumerate(files):
+        try:
+            results = np.loadtxt(f, ndmin=2).T[:, 90:-30]  # clip first and last rows
+            n = results.shape[1]  # number of rows
+            x = np.arange(start, min(stop, n) if stop else n)
+            results = results[:, x]
+            t = (results[0] - results[0].min())  # set t0=0s
+            results[0] = x
+            for i, a in enumerate(ax):
+                if i < len(results):
+                    label = labels[fi] if len(labels) else f.stem.replace('frames_', '')
+                    a.plot(t, results[i], marker='.', label=label, linewidth=1, markersize=5)
+                    a.set_title(s[i])
+                    a.set_xlabel('time (s)')
+                    # if fi == len(files) - 1:
+                    #     a.set_ylim(bottom=0)
+                    for side in ['top', 'right']:
+                        a.spines[side].set_visible(False)
+                else:
+                    a.remove()
+        except Exception as e:
+            print('Warning: Plotting error for %s; %s' % (f, e))
+
+    ax[1].legend()
+    plt.savefig(Path(save_dir) / 'idetection_profile.png', dpi=200)
+
+
+def plot_results_overlay(start=0, stop=0):  # from utils.plots import *; plot_results_overlay()
+    # Plot training 'results*.txt', overlaying train and val losses
+    s = ['train', 'train', 'train', 'Precision', 'mAP@0.5', 'val', 'val', 'val', 'Recall', 'mAP@0.5:0.95']  # legends
+    t = ['Box', 'Objectness', 'Classification', 'P-R', 'mAP-F1']  # titles
+    for f in sorted(glob.glob('results*.txt') + glob.glob('../../Downloads/results*.txt')):
+        results = np.loadtxt(f, usecols=[2, 3, 4, 8, 9, 12, 13, 14, 10, 11], ndmin=2).T
+        n = results.shape[1]  # number of rows
+        x = range(start, min(stop, n) if stop else n)
+        fig, ax = plt.subplots(1, 5, figsize=(14, 3.5), tight_layout=True)
+        ax = ax.ravel()
+        for i in range(5):
+            for j in [i, i + 5]:
+                y = results[j, x]
+                ax[i].plot(x, y, marker='.', label=s[j])
+                # y_smooth = butter_lowpass_filtfilt(y)
+                # ax[i].plot(x, np.gradient(y_smooth), marker='.', label=s[j])
+
+            ax[i].set_title(t[i])
+            ax[i].legend()
+            ax[i].set_ylabel(f) if i == 0 else None  # add filename
+        fig.savefig(f.replace('.txt', '.png'), dpi=200)
+
+
+def plot_results(start=0, stop=0, bucket='', id=(), labels=(), save_dir=''):
+    # Plot training 'results*.txt'. from utils.plots import *; plot_results(save_dir='runs/train/exp')
+    fig, ax = plt.subplots(2, 5, figsize=(12, 6), tight_layout=True)
+    ax = ax.ravel()
+    s = ['Box', 'Objectness', 'Classification', 'Precision', 'Recall',
+         'val Box', 'val Objectness', 'val Classification', 'mAP@0.5', 'mAP@0.5:0.95']
+    if bucket:
+        # files = ['https://storage.googleapis.com/%s/results%g.txt' % (bucket, x) for x in id]
+        files = ['results%g.txt' % x for x in id]
+        c = ('gsutil cp ' + '%s ' * len(files) + '.') % tuple('gs://%s/results%g.txt' % (bucket, x) for x in id)
+        os.system(c)
+    else:
+        files = list(Path(save_dir).glob('results*.txt'))
+    assert len(files), 'No results.txt files found in %s, nothing to plot.' % os.path.abspath(save_dir)
+    for fi, f in enumerate(files):
+        try:
+            results = np.loadtxt(f, usecols=[2, 3, 4, 8, 9, 12, 13, 14, 10, 11], ndmin=2).T
+            n = results.shape[1]  # number of rows
+            x = range(start, min(stop, n) if stop else n)
+            for i in range(10):
+                y = results[i, x]
+                if i in [0, 1, 2, 5, 6, 7]:
+                    y[y == 0] = np.nan  # don't show zero loss values
+                    # y /= y[0]  # normalize
+                label = labels[fi] if len(labels) else f.stem
+                ax[i].plot(x, y, marker='.', label=label, linewidth=2, markersize=8)
+                ax[i].set_title(s[i])
+                # if i in [5, 6, 7]:  # share train and val loss y axes
+                #     ax[i].get_shared_y_axes().join(ax[i], ax[i - 5])
+        except Exception as e:
+            print('Warning: Plotting error for %s; %s' % (f, e))
+
+    ax[1].legend()
+    fig.savefig(Path(save_dir) / 'results.png', dpi=200)

Y5/utils/torch_utils.py

@@ -0,0 +1,294 @@
+# PyTorch utils
+
+import logging
+import math
+import os
+import subprocess
+import time
+from contextlib import contextmanager
+from copy import deepcopy
+from pathlib import Path
+
+import torch
+import torch.backends.cudnn as cudnn
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+
+try:
+    import thop  # for FLOPS computation
+except ImportError:
+    thop = None
+logger = logging.getLogger(__name__)
+
+
+@contextmanager
+def torch_distributed_zero_first(local_rank: int):
+    """
+    Decorator to make all processes in distributed training wait for each local_master to do something.
+    """
+    if local_rank not in [-1, 0]:
+        torch.distributed.barrier()
+    yield
+    if local_rank == 0:
+        torch.distributed.barrier()
+
+
+def init_torch_seeds(seed=0):
+    # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html
+    torch.manual_seed(seed)
+    if seed == 0:  # slower, more reproducible
+        cudnn.benchmark, cudnn.deterministic = False, True
+    else:  # faster, less reproducible
+        cudnn.benchmark, cudnn.deterministic = True, False
+
+
+def git_describe():
+    # return human-readable git description, i.e. v5.0-5-g3e25f1e https://git-scm.com/docs/git-describe
+    if Path('.git').exists():
+        return subprocess.check_output('git describe --tags --long --always', shell=True).decode('utf-8')[:-1]
+    else:
+        return ''
+
+
+def select_device(device='', batch_size=None):
+    # device = 'cpu' or '0' or '0,1,2,3'
+    s = f'YOLOv5 {git_describe()} torch {torch.__version__} '  # string
+    cpu = device.lower() == 'cpu'
+    if cpu:
+        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'  # force torch.cuda.is_available() = False
+    elif device:  # non-cpu device requested
+        os.environ['CUDA_VISIBLE_DEVICES'] = device  # set environment variable
+        assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested'  # check availability
+
+    cuda = not cpu and torch.cuda.is_available()
+    if cuda:
+        n = torch.cuda.device_count()
+        if n > 1 and batch_size:  # check that batch_size is compatible with device_count
+            assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}'
+        space = ' ' * len(s)
+        for i, d in enumerate(device.split(',') if device else range(n)):
+            p = torch.cuda.get_device_properties(i)
+            s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / 1024 ** 2}MB)\n"  # bytes to MB
+    else:
+        s += 'CPU\n'
+
+    logger.info(s)  # skip a line
+    return torch.device('cuda:0' if cuda else 'cpu')
+
+
+def time_synchronized():
+    # pytorch-accurate time
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+    return time.time()
+
+
+def profile(x, ops, n=100, device=None):
+    # profile a pytorch module or list of modules. Example usage:
+    #     x = torch.randn(16, 3, 640, 640)  # input
+    #     m1 = lambda x: x * torch.sigmoid(x)
+    #     m2 = nn.SiLU()
+    #     profile(x, [m1, m2], n=100)  # profile speed over 100 iterations
+
+    device = device or torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    x = x.to(device)
+    x.requires_grad = True
+    print(torch.__version__, device.type, torch.cuda.get_device_properties(0) if device.type == 'cuda' else '')
+    print(f"\n{'Params':>12s}{'GFLOPS':>12s}{'forward (ms)':>16s}{'backward (ms)':>16s}{'input':>24s}{'output':>24s}")
+    for m in ops if isinstance(ops, list) else [ops]:
+        m = m.to(device) if hasattr(m, 'to') else m  # device
+        m = m.half() if hasattr(m, 'half') and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m  # type
+        dtf, dtb, t = 0., 0., [0., 0., 0.]  # dt forward, backward
+        try:
+            flops = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2  # GFLOPS
+        except:
+            flops = 0
+
+        for _ in range(n):
+            t[0] = time_synchronized()
+            y = m(x)
+            t[1] = time_synchronized()
+            try:
+                _ = y.sum().backward()
+                t[2] = time_synchronized()
+            except:  # no backward method
+                t[2] = float('nan')
+            dtf += (t[1] - t[0]) * 1000 / n  # ms per op forward
+            dtb += (t[2] - t[1]) * 1000 / n  # ms per op backward
+
+        s_in = tuple(x.shape) if isinstance(x, torch.Tensor) else 'list'
+        s_out = tuple(y.shape) if isinstance(y, torch.Tensor) else 'list'
+        p = sum(list(x.numel() for x in m.parameters())) if isinstance(m, nn.Module) else 0  # parameters
+        print(f'{p:12.4g}{flops:12.4g}{dtf:16.4g}{dtb:16.4g}{str(s_in):>24s}{str(s_out):>24s}')
+
+
+def is_parallel(model):
+    return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)
+
+
+def intersect_dicts(da, db, exclude=()):
+    # Dictionary intersection of matching keys and shapes, omitting 'exclude' keys, using da values
+    return {k: v for k, v in da.items() if k in db and not any(x in k for x in exclude) and v.shape == db[k].shape}
+
+
+def initialize_weights(model):
+    for m in model.modules():
+        t = type(m)
+        if t is nn.Conv2d:
+            pass  # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+        elif t is nn.BatchNorm2d:
+            m.eps = 1e-3
+            m.momentum = 0.03
+        elif t in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6]:
+            m.inplace = True
+
+
+def find_modules(model, mclass=nn.Conv2d):
+    # Finds layer indices matching module class 'mclass'
+    return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)]
+
+
+def sparsity(model):
+    # Return global model sparsity
+    a, b = 0., 0.
+    for p in model.parameters():
+        a += p.numel()
+        b += (p == 0).sum()
+    return b / a
+
+
+def prune(model, amount=0.3):
+    # Prune model to requested global sparsity
+    import torch.nn.utils.prune as prune
+    print('Pruning model... ', end='')
+    for name, m in model.named_modules():
+        if isinstance(m, nn.Conv2d):
+            prune.l1_unstructured(m, name='weight', amount=amount)  # prune
+            prune.remove(m, 'weight')  # make permanent
+    print(' %.3g global sparsity' % sparsity(model))
+
+
+def fuse_conv_and_bn(conv, bn):
+    # Fuse convolution and batchnorm layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/
+    fusedconv = nn.Conv2d(conv.in_channels,
+                          conv.out_channels,
+                          kernel_size=conv.kernel_size,
+                          stride=conv.stride,
+                          padding=conv.padding,
+                          groups=conv.groups,
+                          bias=True).requires_grad_(False).to(conv.weight.device)
+
+    # prepare filters
+    w_conv = conv.weight.clone().view(conv.out_channels, -1)
+    w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
+    fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size()))
+
+    # prepare spatial bias
+    b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias
+    b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
+    fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)
+
+    return fusedconv
+
+
+def model_info(model, verbose=False, img_size=640):
+    # Model information. img_size may be int or list, i.e. img_size=640 or img_size=[640, 320]
+    n_p = sum(x.numel() for x in model.parameters())  # number parameters
+    n_g = sum(x.numel() for x in model.parameters() if x.requires_grad)  # number gradients
+    if verbose:
+        print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma'))
+        for i, (name, p) in enumerate(model.named_parameters()):
+            name = name.replace('module_list.', '')
+            print('%5g %40s %9s %12g %20s %10.3g %10.3g' %
+                  (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
+
+    try:  # FLOPS
+        from thop import profile
+        stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32
+        img = torch.zeros((1, model.yaml.get('ch', 3), stride, stride), device=next(model.parameters()).device)  # input
+        flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2  # stride GFLOPS
+        img_size = img_size if isinstance(img_size, list) else [img_size, img_size]  # expand if int/float
+        fs = ', %.1f GFLOPS' % (flops * img_size[0] / stride * img_size[1] / stride)  # 640x640 GFLOPS
+    except (ImportError, Exception):
+        fs = ''
+
+    logger.info(f"Model Summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}")
+
+
+def load_classifier(name='resnet101', n=2):
+    # Loads a pretrained model reshaped to n-class output
+    model = torchvision.models.__dict__[name](pretrained=True)
+
+    # ResNet model properties
+    # input_size = [3, 224, 224]
+    # input_space = 'RGB'
+    # input_range = [0, 1]
+    # mean = [0.485, 0.456, 0.406]
+    # std = [0.229, 0.224, 0.225]
+
+    # Reshape output to n classes
+    filters = model.fc.weight.shape[1]
+    model.fc.bias = nn.Parameter(torch.zeros(n), requires_grad=True)
+    model.fc.weight = nn.Parameter(torch.zeros(n, filters), requires_grad=True)
+    model.fc.out_features = n
+    return model
+
+
+def scale_img(img, ratio=1.0, same_shape=False, gs=32):  # img(16,3,256,416)
+    # scales img(bs,3,y,x) by ratio constrained to gs-multiple
+    if ratio == 1.0:
+        return img
+    else:
+        h, w = img.shape[2:]
+        s = (int(h * ratio), int(w * ratio))  # new size
+        img = F.interpolate(img, size=s, mode='bilinear', align_corners=False)  # resize
+        if not same_shape:  # pad/crop img
+            h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)]
+        return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447)  # value = imagenet mean
+
+
+def copy_attr(a, b, include=(), exclude=()):
+    # Copy attributes from b to a, options to only include [...] and to exclude [...]
+    for k, v in b.__dict__.items():
+        if (len(include) and k not in include) or k.startswith('_') or k in exclude:
+            continue
+        else:
+            setattr(a, k, v)
+
+
+class ModelEMA:
+    """ Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models
+    Keep a moving average of everything in the model state_dict (parameters and buffers).
+    This is intended to allow functionality like
+    https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
+    A smoothed version of the weights is necessary for some training schemes to perform well.
+    This class is sensitive where it is initialized in the sequence of model init,
+    GPU assignment and distributed training wrappers.
+    """
+
+    def __init__(self, model, decay=0.9999, updates=0):
+        # Create EMA
+        self.ema = deepcopy(model.module if is_parallel(model) else model).eval()  # FP32 EMA
+        # if next(model.parameters()).device.type != 'cpu':
+        #     self.ema.half()  # FP16 EMA
+        self.updates = updates  # number of EMA updates
+        self.decay = lambda x: decay * (1 - math.exp(-x / 2000))  # decay exponential ramp (to help early epochs)
+        for p in self.ema.parameters():
+            p.requires_grad_(False)
+
+    def update(self, model):
+        # Update EMA parameters
+        with torch.no_grad():
+            self.updates += 1
+            d = self.decay(self.updates)
+
+            msd = model.module.state_dict() if is_parallel(model) else model.state_dict()  # model state_dict
+            for k, v in self.ema.state_dict().items():
+                if v.dtype.is_floating_point:
+                    v *= d
+                    v += (1. - d) * msd[k].detach()
+
+    def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):
+        # Update EMA attributes
+        copy_attr(self.ema, model, include, exclude)

Y5/utils/wandb_logging/log_dataset.py

@@ -0,0 +1,39 @@
+import argparse
+from pathlib import Path
+
+import yaml
+
+from wandb_utils import WandbLogger
+from Y5.utils.datasets import LoadImagesAndLabels
+
+WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
+
+
+def create_dataset_artifact(opt):
+    with open(opt.data) as f:
+        data = yaml.load(f, Loader=yaml.SafeLoader)  # data dict
+    logger = WandbLogger(opt, '', None, data, job_type='create_dataset')
+    nc, names = (1, ['item']) if opt.single_cls else (int(data['nc']), data['names'])
+    names = {k: v for k, v in enumerate(names)}  # to index dictionary
+    logger.log_dataset_artifact(LoadImagesAndLabels(data['train']), names, name='train')  # trainset
+    logger.log_dataset_artifact(LoadImagesAndLabels(data['val']), names, name='val')  # valset
+
+    # Update data.yaml with artifact links
+    data['train'] = WANDB_ARTIFACT_PREFIX + str(Path(opt.project) / 'train')
+    data['val'] = WANDB_ARTIFACT_PREFIX + str(Path(opt.project) / 'val')
+    path = opt.data if opt.overwrite_config else opt.data.replace('.', '_wandb.')  # updated data.yaml path
+    data.pop('download', None)  # download via artifact instead of predefined field 'download:'
+    with open(path, 'w') as f:
+        yaml.dump(data, f)
+    print("New Config file => ", path)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
+    parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
+    parser.add_argument('--project', type=str, default='YOLOv5', help='name of W&B Project')
+    parser.add_argument('--overwrite_config', action='store_true', help='overwrite data.yaml')
+    opt = parser.parse_args()
+
+    create_dataset_artifact(opt)

Y5/utils/wandb_logging/wandb_utils.py

@@ -0,0 +1,145 @@
+import json
+import shutil
+import sys
+from datetime import datetime
+from pathlib import Path
+
+import torch
+
+sys.path.append(str(Path(__file__).parent.parent.parent))  # add utils/ to path
+from Y5.utils.general import colorstr, xywh2xyxy
+
+try:
+    import wandb
+except ImportError:
+    wandb = None
+    print(f"{colorstr('wandb: ')}Install Weights & Biases for YOLOv5 logging with 'pip install wandb' (recommended)")
+
+WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
+
+
+def remove_prefix(from_string, prefix):
+    return from_string[len(prefix):]
+
+
+class WandbLogger():
+    def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
+        self.wandb = wandb
+        self.wandb_run = wandb.init(config=opt, resume="allow",
+                                    project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
+                                    name=name,
+                                    job_type=job_type,
+                                    id=run_id) if self.wandb else None
+
+        if job_type == 'Training':
+            self.setup_training(opt, data_dict)
+            if opt.bbox_interval == -1:
+                opt.bbox_interval = (opt.epochs // 10) if opt.epochs > 10 else opt.epochs
+            if opt.save_period == -1:
+                opt.save_period = (opt.epochs // 10) if opt.epochs > 10 else opt.epochs
+
+    def setup_training(self, opt, data_dict):
+        self.log_dict = {}
+        self.train_artifact_path, self.trainset_artifact = \
+            self.download_dataset_artifact(data_dict['train'], opt.artifact_alias)
+        self.test_artifact_path, self.testset_artifact = \
+            self.download_dataset_artifact(data_dict['val'], opt.artifact_alias)
+        self.result_artifact, self.result_table, self.weights = None, None, None
+        if self.train_artifact_path is not None:
+            train_path = Path(self.train_artifact_path) / 'data/images/'
+            data_dict['train'] = str(train_path)
+        if self.test_artifact_path is not None:
+            test_path = Path(self.test_artifact_path) / 'data/images/'
+            data_dict['val'] = str(test_path)
+            self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
+            self.result_table = wandb.Table(["epoch", "id", "prediction", "avg_confidence"])
+        if opt.resume_from_artifact:
+            modeldir, _ = self.download_model_artifact(opt.resume_from_artifact)
+            if modeldir:
+                self.weights = Path(modeldir) / "best.pt"
+                opt.weights = self.weights
+
+    def download_dataset_artifact(self, path, alias):
+        if path.startswith(WANDB_ARTIFACT_PREFIX):
+            dataset_artifact = wandb.use_artifact(remove_prefix(path, WANDB_ARTIFACT_PREFIX) + ":" + alias)
+            assert dataset_artifact is not None, "'Error: W&B dataset artifact doesn\'t exist'"
+            datadir = dataset_artifact.download()
+            labels_zip = Path(datadir) / "data/labels.zip"
+            shutil.unpack_archive(labels_zip, Path(datadir) / 'data/labels', 'zip')
+            print("Downloaded dataset to : ", datadir)
+            return datadir, dataset_artifact
+        return None, None
+
+    def download_model_artifact(self, name):
+        model_artifact = wandb.use_artifact(name + ":latest")
+        assert model_artifact is not None, 'Error: W&B model artifact doesn\'t exist'
+        modeldir = model_artifact.download()
+        print("Downloaded model to : ", modeldir)
+        return modeldir, model_artifact
+
+    def log_model(self, path, opt, epoch):
+        datetime_suffix = datetime.today().strftime('%Y-%m-%d-%H-%M-%S')
+        model_artifact = wandb.Artifact('run_' + wandb.run.id + '_model', type='model', metadata={
+            'original_url': str(path),
+            'epoch': epoch + 1,
+            'save period': opt.save_period,
+            'project': opt.project,
+            'datetime': datetime_suffix
+        })
+        model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
+        model_artifact.add_file(str(path / 'best.pt'), name='best.pt')
+        wandb.log_artifact(model_artifact)
+        print("Saving model artifact on epoch ", epoch + 1)
+
+    def log_dataset_artifact(self, dataset, class_to_id, name='dataset'):
+        artifact = wandb.Artifact(name=name, type="dataset")
+        image_path = dataset.path
+        artifact.add_dir(image_path, name='data/images')
+        table = wandb.Table(columns=["id", "train_image", "Classes"])
+        class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()])
+        for si, (img, labels, paths, shapes) in enumerate(dataset):
+            height, width = shapes[0]
+            labels[:, 2:] = (xywh2xyxy(labels[:, 2:].view(-1, 4)))
+            labels[:, 2:] *= torch.Tensor([width, height, width, height])
+            box_data = []
+            img_classes = {}
+            for cls, *xyxy in labels[:, 1:].tolist():
+                cls = int(cls)
+                box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
+                                 "class_id": cls,
+                                 "box_caption": "%s" % (class_to_id[cls]),
+                                 "scores": {"acc": 1},
+                                 "domain": "pixel"})
+                img_classes[cls] = class_to_id[cls]
+            boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}}  # inference-space
+            table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), json.dumps(img_classes))
+        artifact.add(table, name)
+        labels_path = 'labels'.join(image_path.rsplit('images', 1))
+        zip_path = Path(labels_path).parent / (name + '_labels.zip')
+        if not zip_path.is_file():  # make_archive won't check if file exists
+            shutil.make_archive(zip_path.with_suffix(''), 'zip', labels_path)
+        artifact.add_file(str(zip_path), name='data/labels.zip')
+        wandb.log_artifact(artifact)
+        print("Saving data to W&B...")
+
+    def log(self, log_dict):
+        if self.wandb_run:
+            for key, value in log_dict.items():
+                self.log_dict[key] = value
+
+    def end_epoch(self):
+        if self.wandb_run and self.log_dict:
+            wandb.log(self.log_dict)
+        self.log_dict = {}
+
+    def finish_run(self):
+        if self.wandb_run:
+            if self.result_artifact:
+                print("Add Training Progress Artifact")
+                self.result_artifact.add(self.result_table, 'result')
+                train_results = wandb.JoinedTable(self.testset_artifact.get("val"), self.result_table, "id")
+                self.result_artifact.add(train_results, 'joined_result')
+                wandb.log_artifact(self.result_artifact)
+            if self.log_dict:
+                wandb.log(self.log_dict)
+            wandb.run.finish()

bigdata.py

@@ -0,0 +1,173 @@
+'''
+author ：  Hatcher W
+Email： hcwang@semi.ac.cn
+use: 使用SSD算法对大数据进行筛选，先使用Hatcher is cool 将文件夹中所有的视频和图像的路径进行递归存放在videoroot.json和imageroot.json文件中
+input:两个json文件
+output：相同路径（改为data）有动物图片的结果存储
+'''
+
+import json
+import glob
+import os
+import time
+
+import torch
+from PIL import Image
+from vizer.draw import draw_boxes
+
+from ssd.config import cfg
+from ssd.data.datasets import COCODataset, VOCDataset
+import argparse
+import numpy as np
+
+from ssd.data.transforms import build_transforms
+from ssd.modeling.detector import build_detection_model
+from ssd.utils import mkdir
+from ssd.utils.checkpoint import CheckPointer
+from vediotoima import  vedio_img
+os.environ['CUDA_VISIBLE_DEVICES'] = '5'
+FPS = 50
+@torch.no_grad()
+def run_demo(cfg, ckpt, score_threshold, images_dir, output_dir, dataset_type):
+    if dataset_type == "voc":
+        class_names = VOCDataset.class_names
+    elif dataset_type == 'coco':
+        class_names = COCODataset.class_names
+    else:
+        raise NotImplementedError('Not implemented now.')
+    device = torch.device(cfg.MODEL.DEVICE)
+
+    model = build_detection_model(cfg)
+    model = model.to(device)
+    checkpointer = CheckPointer(model, save_dir=cfg.OUTPUT_DIR)
+    checkpointer.load(ckpt, use_latest=ckpt is None)
+    weight_file = ckpt if ckpt else checkpointer.get_checkpoint_file()
+    print('Loaded weights from {}'.format(weight_file))
+
+    image_paths = glob.glob(os.path.join(images_dir, '*.jpg'))
+    mkdir(output_dir)
+
+    cpu_device = torch.device("cpu")
+    transforms = build_transforms(cfg, is_train=False)
+    model.eval()
+    for i, image_path in enumerate(image_paths):
+        start = time.time()
+        image_name = os.path.basename(image_path)
+
+        image = np.array(Image.open(image_path).convert("RGB"))
+        height, width = image.shape[:2]
+        images = transforms(image)[0].unsqueeze(0)
+        load_time = time.time() - start
+
+        start = time.time()
+        result = model(images.to(device))[0]
+        inference_time = time.time() - start
+
+        result = result.resize((width, height)).to(cpu_device).numpy()
+        boxes, labels, scores = result['boxes'], result['labels'], result['scores']
+
+        indices = scores > score_threshold
+        boxes = boxes[indices]
+        #print(len(boxes))
+        #exit()
+        labels = labels[indices]
+        scores = scores[indices]
+        meters = ' | '.join(
+            [
+                'objects {:02d}'.format(len(boxes)),
+                'load {:03d}ms'.format(round(load_time * 1000)),
+                'inference {:03d}ms'.format(round(inference_time * 1000)),
+                'FPS {}'.format(round(1.0 / inference_time))
+            ]
+        )
+        print('({:04d}/{:04d}) {}: {}'.format(i + 1, len(image_paths), image_name, meters))
+
+        #drawn_image = draw_boxes(image, boxes, labels, scores, class_names).astype(np.uint8)
+        if(len(boxes)):
+
+            Image.fromarray(image).save(os.path.join(output_dir, image_name))
+            print(image_name+"is saved")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="SSD Demo.")
+    parser.add_argument(
+        "--config-file",
+        default="configs/vgg_ssd512_voc0712.yaml",
+        metavar="FILE",
+        help="path to config file",
+        type=str,
+    )
+    parser.add_argument("--ckpt", type=str, default="../data_MPKFNet/weigth/detection/detection.pth", help="Trained weights.")
+    parser.add_argument("--score_threshold", type=float, default=0.3)
+    parser.add_argument("--images_dir", default='../data_MPKFNet/result_raw/yang/', type=str, help='Specify a image dir to do prediction.')
+    parser.add_argument("--output_dir", default='../data_MPKFNet/result/yangd/', type=str, help='Specify a image dir to save predicted images.')
+    parser.add_argument("--dataset_type", default="voc", type=str, help='Specify dataset type. Currently support voc and coco.')
+
+    parser.add_argument(
+        "opts",
+        help="Modify config options using the command-line",
+        default=None,
+        nargs=argparse.REMAINDER,
+    )
+    args = parser.parse_args()
+    print(args)
+
+    cfg.merge_from_file(args.config_file)
+    cfg.merge_from_list(args.opts)
+    cfg.freeze()
+
+    print("Loaded configuration file {}".format(args.config_file))
+    with open(args.config_file, "r") as cf:
+        config_str = "\n" + cf.read()
+        print(config_str)
+    print("Running with config:\n{}".format(cfg))
+    ########################################
+    with open ("./vedioroot.json",'r+') as fvedio:
+        vediolists = json.load(fvedio)
+        fvedio.close()
+    with open("./loglast.txt", 'r') as f:
+        last_path = f.readline()
+        print("lastpath is "+last_path)
+        #exit()
+        f.close()
+    flag_detectlast = True
+    for vediolist in vediolists:
+        if flag_detectlast:
+            if (vediolist == last_path) or (last_path == ""):
+                flag_detectlast=False
+                print("Detecting finish and begining")
+                continue
+            else:
+                print(vediolist +"has been finished")
+                continue
+        vedio_path = vediolist.replace(".MP4","").replace(".mp4","").replace(".MOV","").replace(".mov","")
+        vedio_rawoutpath = vedio_path.replace("/Hatcher/DATA/","/Hatcher/model/data_MPKFNet/bigdata/rawdata/")
+        #print("rawouput is "+vedio_rawoutpath)
+        #exit()
+        if not os.path.exists(vedio_rawoutpath):
+            os.makedirs(vedio_rawoutpath)
+        vedio_img(vediolist,vedio_rawoutpath,FPS)
+        ######finish vedio to image
+        args.images_dir = vedio_rawoutpath
+        args.output_dir = vedio_path.replace("/Hatcher/DATA/","/Hatcher/model/data_MPKFNet/bigdata/resultdata/")
+        #print(vedio_path)
+        #print(type(vedio_path))
+        #exit()
+        if not os.path.exists(args.output_dir):
+            os.makedirs(args.output_dir)
+
+        run_demo(cfg=cfg,
+                 ckpt=args.ckpt,
+                 score_threshold=args.score_threshold,
+                 images_dir=args.images_dir,
+                 output_dir=args.output_dir,
+                 dataset_type=args.dataset_type)
+        with open("./loglast.txt",'w') as f:
+            f.write(vediolist)
+            f.close()
+
+
+
+if __name__ == '__main__':
+    main()

bigdata_copyall.py

@@ -0,0 +1,32 @@
+'''
+author ：  Hatcher W
+Email： hcwang@semi.ac.cn
+use: 将所有的图片存储到一个文件夹中，便于后续筛选
+input:两个json文件 用于提供所有的原始图片路径
+output：在ALL文件中存储所有的图片
+'''
+import json
+import os
+import shutil
+def bigdata_copy(desname):
+    with open("./vedioroot.json", 'r+') as fvedio:
+        vediolists = json.load(fvedio)
+        fvedio.close()
+    flag_detectlast = True
+    num = 0
+    for vediolist in vediolists:
+        image_num = str(num).zfill(7)
+        source_path = vediolist.replace(".MP4", "").replace(".mp4", "").replace(".MOV", "").replace(".mov", "").replace("/Hatcher/DATA/", "/Hatcher/model/data_MPKFNet/bigdata/resultdata/")
+        despath = "../data_MPKFNet/bigdata/ALL/"+desname
+
+        if not os.path.exists(despath):
+            os.makedirs(despath)
+
+        ######finish vedio to image
+        for image in os.listdir(source_path):
+            shutil.copy(os.path.join(source_path,image),os.path.join(despath,image_num)+".jpg")
+            print(os.path.join(source_path,image))
+            num +=1
+            image_num = str(num).zfill(7)
+if __name__ == "__main__":
+    bigdata_copy("Hat_Raw_Data")

bigdata_image.py

@@ -0,0 +1,193 @@
+'''
+author ：  Hatcher W
+Email： hcwang@semi.ac.cn
+use: 使用SSD算法对大数据进行筛选，先使用Hatcher is cool 将文件夹中所有的视频和图像的路径进行递归存放在videoroot.json和imageroot.json文件中
+input:两个json文件
+output：相同路径（改为data）有动物图片的结果存储
+'''
+
+import json
+import glob
+import os
+import time
+
+import torch
+from PIL import Image
+from vizer.draw import draw_boxes
+
+from ssd.config import cfg
+from ssd.data.datasets import COCODataset, VOCDataset
+import argparse
+import numpy as np
+
+from ssd.data.transforms import build_transforms
+from ssd.modeling.detector import build_detection_model
+from ssd.utils import mkdir
+from ssd.utils.checkpoint import CheckPointer
+from vediotoima import  vedio_img
+os.environ['CUDA_VISIBLE_DEVICES'] = '5'
+FPS = 5
+num_image = 380000
+
+
+@torch.no_grad()
+def run_demo(cfg, ckpt, score_threshold, images_dir, output_dir, dataset_type,model):
+    # if dataset_type == "voc":
+    #     class_names = VOCDataset.class_names
+    # elif dataset_type == 'coco':
+    #     class_names = COCODataset.class_names
+    # else:
+    #     raise NotImplementedError('Not implemented now.')
+    # device = torch.device(cfg.MODEL.DEVICE)
+    #
+    # model = build_detection_model(cfg)
+    # model = model.to(device)
+    # checkpointer = CheckPointer(model, save_dir=cfg.OUTPUT_DIR)
+    # checkpointer.load(ckpt, use_latest=ckpt is None)
+    # weight_file = ckpt if ckpt else checkpointer.get_checkpoint_file()
+    # print('Loaded weights from {}'.format(weight_file))
+    #
+    # #image_paths = glob.glob(os.path.join(images_dir, '*.jpg'))
+    global num_image
+    device = torch.device(cfg.MODEL.DEVICE)
+    image_paths = images_dir
+    #mkdir(output_dir)
+
+    cpu_device = torch.device("cpu")
+    transforms = build_transforms(cfg, is_train=False)
+    # model.eval()
+
+    start = time.time()
+    image_name = os.path.basename(image_paths)
+
+    image = np.array(Image.open(image_paths).convert("RGB"))
+    height, width = image.shape[:2]
+    images = transforms(image)[0].unsqueeze(0)
+    load_time = time.time() - start
+
+    start = time.time()
+    result = model(images.to(device))[0]
+    inference_time = time.time() - start
+
+    result = result.resize((width, height)).to(cpu_device).numpy()
+    boxes, labels, scores = result['boxes'], result['labels'], result['scores']
+
+    indices = scores > score_threshold
+    boxes = boxes[indices]
+    #print(len(boxes))
+    #exit()
+    labels = labels[indices]
+    scores = scores[indices]
+    meters = ' | '.join(
+        [
+            'objects {:02d}'.format(len(boxes)),
+            'load {:03d}ms'.format(round(load_time * 1000)),
+            'inference {:03d}ms'.format(round(inference_time * 1000)),
+            'FPS {}'.format(round(1.0 / inference_time))
+        ]
+    )
+    print('({:04d}/{:04d}) {}: {}'.format(1, len(image_paths), image_name, meters))
+
+    #drawn_image = draw_boxes(image, boxes, labels, scores, class_names).astype(np.uint8)
+    if(len(boxes)) :
+        print(labels)
+        Image.fromarray(image).save(os.path.join(output_dir,"{}.jpg".format(num_image)))
+        num_image +=1
+        print(image_name+"is saved")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="SSD Demo.")
+    parser.add_argument(
+        "--config-file",
+        default="configs/vgg_ssd512_voc0712.yaml",
+        metavar="FILE",
+        help="path to config file",
+        type=str,
+    )
+    parser.add_argument("--ckpt", type=str, default="../data_MPKFNet/weigth/detection/detection.pth", help="Trained weights.")
+    parser.add_argument("--score_threshold", type=float, default=0.3)
+    parser.add_argument("--images_dir", default='../data_MPKFNet/result_raw/yang/', type=str, help='Specify a image dir to do prediction.')
+    parser.add_argument("--output_dir", default='../data_MPKFNet/result/yangd/', type=str, help='Specify a image dir to save predicted images.')
+    parser.add_argument("--dataset_type", default="voc", type=str, help='Specify dataset type. Currently support voc and coco.')
+
+    parser.add_argument(
+        "opts",
+        help="Modify config options using the command-line",
+        default=None,
+        nargs=argparse.REMAINDER,
+    )
+    args = parser.parse_args()
+    print(args)
+
+    cfg.merge_from_file(args.config_file)
+    cfg.merge_from_list(args.opts)
+    cfg.freeze()
+
+    print("Loaded configuration file {}".format(args.config_file))
+    with open(args.config_file, "r") as cf:
+        config_str = "\n" + cf.read()
+        print(config_str)
+    print("Running with config:\n{}".format(cfg))
+    ########################################
+    with open ("./imageroot.json",'r+') as fvedio:
+        vediolists = json.load(fvedio)
+        fvedio.close()
+    with open("./log_imglast.txt", 'r') as f:
+        last_path = f.readline()
+        print("lastpath is "+last_path)
+        #exit()
+        f.close()
+    flag_detectlast = True
+    #######
+    if args.dataset_type == "voc":
+        class_names = VOCDataset.class_names
+    elif args.dataset_type == 'coco':
+        class_names = COCODataset.class_names
+    else:
+        raise NotImplementedError('Not implemented now.')
+    device = torch.device(cfg.MODEL.DEVICE)
+
+    model = build_detection_model(cfg)
+    model = model.to(device)
+    checkpointer = CheckPointer(model, save_dir=cfg.OUTPUT_DIR)
+    checkpointer.load(args.ckpt, use_latest=args.ckpt is None)
+    weight_file = args.ckpt if args.ckpt else checkpointer.get_checkpoint_file()
+    print('Loaded weights from {}'.format(weight_file))
+
+
+    model.eval()
+    #####
+
+    for vediolist in vediolists:
+        if flag_detectlast:
+            if (vediolist == last_path) or (last_path == ""):
+                flag_detectlast=False
+                print("Detecting finish and begining")
+                continue
+            else:
+                print(vediolist +"has been finished")
+                continue
+        vedio_path = vediolist
+        args.images_dir = vedio_path
+        args.output_dir = "/data1//Hatcher/model/data_MPKFNet/bigdata/ALL_image/"
+
+        if not os.path.exists(args.output_dir):
+            os.makedirs(args.output_dir)
+
+        run_demo(cfg=cfg,
+                 ckpt=args.ckpt,
+                 score_threshold=args.score_threshold,
+                 images_dir=args.images_dir,
+                 output_dir=args.output_dir,
+                 dataset_type=args.dataset_type,
+                 model = model
+                 )
+        with open("./log_imglast.txt",'w') as f:
+            f.write(vediolist)
+            f.close()
+
+
+
+if __name__ == '__main__':
+    main()

check_VOC.py

@@ -0,0 +1,54 @@
+'''
+author ：  Hatcher W
+Email： hcwang@semi.ac.cn
+use: 用于检测voc数据集是否正确，
+input:（图像的文件夹路径，图像标签文件夹路径）
+output：check路径，所有的可视化结果都在里面
+'''
+import os
+import cv2
+import  xml.dom.minidom
+from tqdm import tqdm
+
+image_path="/data1/Hatcher/model/20210708/ftp_Result/860548046389083/RAWImage/WVGA__20210604_040328_100_0356/"
+annotation_path="/data1/Hatcher/model/20210708/ftp_Result/860548046389083/VideoResult/WVGA_20210604_040328_100_0356/"
+save_path ="/data1/Hatcher/model/20210708/ftp_Result/860548046389083/VideoResult/check/"
+
+files_name = os.listdir(image_path)
+for filename_ in tqdm(files_name):
+    filename, extension= os.path.splitext(filename_)
+    #filename=filename.split('__')[0]
+    img_path =image_path+filename+'.jpg'
+    xml_path =annotation_path+filename+'.xml'
+    img = cv2.imread(img_path)
+    if img is None:
+        pass
+    try:
+        dom = xml.dom.minidom.parse(xml_path)
+    except:
+        #os.remove(img_path)
+        continue
+    root = dom.documentElement
+    objects=dom.getElementsByTagName("object")
+    for object in objects:
+        bndbox = object.getElementsByTagName('bndbox')[0]
+        xmin = bndbox.getElementsByTagName('xmin')[0]
+        ymin = bndbox.getElementsByTagName('ymin')[0]
+        xmax = bndbox.getElementsByTagName('xmax')[0]
+        ymax = bndbox.getElementsByTagName('ymax')[0]
+        #print(xmin.childNodes[0].data)
+        #print(type(xmin.childNodes[0].data))
+        xmin_data=int(xmin.childNodes[0].data)
+        ymin_data=int(ymin.childNodes[0].data)
+        xmax_data=int(xmax.childNodes[0].data)
+        ymax_data=int(ymax.childNodes[0].data)
+        label_name=object.getElementsByTagName('name')[0].childNodes[0].data
+        cv2.rectangle(img,(xmin_data,ymin_data),(xmax_data,ymax_data),(55,255,155),1)
+        cv2.putText(img,label_name,(xmin_data,ymin_data-3),cv2.FONT_HERSHEY_SIMPLEX,0.5,(0,255,0),2)
+    flag=0
+    flag=cv2.imwrite(save_path+"{}.jpg".format(filename),img)
+    if not (flag):
+        print(filename,"error")
+print("all done ====================================")
+
+#(int((xmin_data+xmax_data/2)),int((ymin_data+ymax_data)/2))

check_info.py

@@ -0,0 +1,192 @@
+'''
+author ：  Hatcher W
+Email： hcwang@semi.ac.cn
+use: 检测通用计算平台的硬件信息
+'''
+from __future__ import print_function
+from collections import OrderedDict
+import subprocess
+def cpuinfo():
+    """
+    Return the information in /proc/cpuinfo
+    as a dictionary in the following format:
+    cpu_info['proc0']={...}
+    cpu_info['proc1']={...}
+    """
+
+    cpuinfo = OrderedDict()
+    procinfo = OrderedDict()
+
+    nprocs = 0
+    with open('/proc/cpuinfo') as f:
+        for line in f:
+            if not line.strip():
+                # end of one processor
+                cpuinfo['proc%s' % nprocs] = procinfo
+                nprocs = nprocs + 1
+                # Reset
+                procinfo = OrderedDict()
+            else:
+                if len(line.split(':')) == 2:
+                    procinfo[line.split(':')[0].strip()] = line.split(':')[1].strip()
+                else:
+                    procinfo[line.split(':')[0].strip()] = ''
+
+    return cpuinfo
+
+
+def meminfo():
+    """
+    Return the information in /proc/meminfo
+    as a dictionary
+    """
+
+    meminfo = OrderedDict()
+
+    with open('/proc/meminfo') as f:
+        for line in f:
+            meminfo[line.split(':')[0]] = line.split(':')[1].strip()
+    return meminfo
+def gpuinfo():
+    """
+    Return the information in /proc/gpuinfo
+    as a dictionary
+    """
+
+    gpuinfo = OrderedDict()
+
+    with open('/proc/gpuinfo') as f:
+        for line in f:
+            meminfo[line.split(':')[0]] = line.split(':')[1].strip()
+    return gpuinfo
+def getosversion():
+    with open('/etc/issue') as fd:
+        for line in fd:
+            osver = line.strip()
+            break
+    return {'osver':osver}
+def getBaseboardSerialnumber ():
+    p = subprocess.Popen(["sudo dmidecode -t 2 | grep Serial"], shell=True, stdin=subprocess.PIPE, stdout=subprocess.PIPE , stderr=subprocess.PIPE)
+    data = p.stdout
+    lines = []
+    while True:
+        line = str(data.readline(), encoding="utf-8")
+        if line == '\n':
+            break
+        if line:
+            d = dict([line.strip().split(': ')])
+            lines.append(d)
+        else:
+            break
+    return lines
+def getCpuId():
+    p = subprocess.Popen(["sudo dmidecode -t 4 | grep ID"], shell=True, stdin=subprocess.PIPE, stdout=subprocess.PIPE , stderr=subprocess.PIPE)
+    data = p.stdout
+    lines = []
+    while True:
+        line = str(data.readline(), encoding="utf-8")
+        if line == '\n':
+            break
+        if line:
+            d = dict([line.strip().split(': ')])
+            lines.append(d)
+        else:
+            break
+    return lines
+def getMachineNum():
+    p = subprocess.Popen(["sudo dmidecode | grep 'Product Name'"], shell=True, stdin=subprocess.PIPE, stdout=subprocess.PIPE , stderr=subprocess.PIPE)
+    data = p.stdout
+    lines = []
+    while True:
+        line = str(data.readline(), encoding="utf-8")
+        if line == '\n':
+            break
+        if line:
+            d = dict([line.strip().split(': ')])
+            lines.append(d)
+        else:
+            break
+    return lines
+def getGPU():
+    p = subprocess.Popen(["sudo lspci -vnn | grep VGA -A 12"], shell=True, stdin=subprocess.PIPE, stdout=subprocess.PIPE , stderr=subprocess.PIPE)
+    data = p.stdout
+    print(data)
+    lines = []
+    while True:
+        line = str(data.readline(), encoding="utf-8")
+        print(type(line))
+        if line == '\n' :
+            break
+        if type(line)!="dict":
+            continue
+        if line:
+            d = dict([line.strip().split(': ')])
+            lines.append(d)
+        else:
+            break
+    return lines
+if __name__ == '__main__':
+    #print(meminfo())
+    #gpuinfo = gpuinfo()
+    #getgpu = getGPU()
+    getmachinenum = getMachineNum()
+    getbaseboardserialnumber =getBaseboardSerialnumber()
+    getcpuid = getCpuId()
+    osver = getosversion()
+    meminfo = meminfo()
+    cpuinfo = cpuinfo()
+    #exit()
+    print(cpuinfo["proc0"]["model name"])
+    if cpuinfo["proc0"]["model name"] == "Intel(R) Xeon(R) Gold 6230 CPU @ 2.10GHz":
+        print(cpuinfo["proc0"]["model name"])
+        print("model name is ok")
+    else:
+        print(cpuinfo["proc0"]["model name"])
+        print("CPU model is wrong")
+    if getcpuid[0]['ID'] == '57 06 05 00 FF FB EB BF':
+        print(getcpuid[0]['ID'])
+        print("getcpuid is ok")
+    else:
+        print(getcpuid[0]['ID'])
+        print("getcpuid is wrong")
+    if cpuinfo["proc0"]["microcode"] == "0x5002f01":
+        print(cpuinfo["proc0"]["microcode"])
+        print("microcode is ok")
+    else:
+        print(cpuinfo["proc0"]["microcode"])
+        print("microcode is wrong")
+    if cpuinfo["proc0"]["vendor_id"] == "GenuineIntel":
+        print(cpuinfo["proc0"]["vendor_id"])
+        print("vendor_id is ok")
+    else:
+        print(cpuinfo["proc0"]["vendor_id"])
+        print("vendor_id is wrong")
+    if meminfo['MemTotal'] == "394851452 kB":
+        print(meminfo['MemTotal'])
+        print("MemTotal is ok")
+    else:
+        print(meminfo['MemTotal'])
+        print("MemTotal is wrong")
+    if osver['osver'] == 'Ubuntu 16.04.6 LTS \\n \\l':
+        print(osver['osver'])
+        print("osver is ok")
+    else:
+        print(osver['osver'])
+        print("osver is wrong")
+    if getbaseboardserialnumber[0]['Serial Number'] == 'VM194S002416':
+        print(getbaseboardserialnumber[0]['Serial Number'])
+        print("getbaseboardserialnumber is ok")
+    else:
+        print(getbaseboardserialnumber[0]['Serial Number'])
+        print("getbaseboardserialnumber is wrong")
+    if getmachinenum[0]['Product Name'] == 'SYS-4029GP-TRT':
+        print(getmachinenum[0]['Product Name'])
+        print("Product Name is ok")
+    else:
+        print(getmachinenum[0]['Product Name'])
+        print("Product Name is wrong")
+    #print(cpuinfo)
+    #for processor in cpuinfo.keys():
+        #print(cpuinfo[processor]['model name'])
+
+

common.py

@@ -0,0 +1,44 @@
+#进入线程先把标志位置为True
+# from System_I import Detection
+# import os,torch
+# os.environ['CUDA_VISIBLE_DEVICES'] = '1'
+# devices = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+# root = "/home/wanghc/output"
+import  os,torch, threading,argparse
+class global_para():
+    parser1 = argparse.ArgumentParser(description="GYD")
+    parser1.add_argument("video_num", default='0001', type=str,
+                        help='Video name.')
+    parser1.add_argument("GPU", default='5', type=str,
+                        help='GPU.')
+    parser1.add_argument("thresh", default=0.6, type=float,
+                        help='GPU.')
+    parser1.add_argument("red", default=0, type=int,
+                        help='red.')
+
+    args = parser1.parse_args()
+    print(args)
+    red_flag = args.red
+    detect_thresh = args.thresh
+    is_ex_bird = False  # 目标检测标志位，每个一秒检测一次目标检测，如果有鸟，则置为True，为True的时候说明， 处理30s的线程还在继续，直到标志位变为了False再开启新的线程
+    Density_thread_flag = False  # 密度估计线程标志位，为True说明开启了密度估计线程，处理完750帧为止，标志位标为False，
+    Detection_thread_flag = False  # 目标检测和细粒度识别线程标志位，为True说明开启了密度估计线程，处理完750帧为止，标志位标为False，
+    Vedio_thread_flag = False  # 存储视频线程标志位，存储30s
+    SaveImage_thread_flag = False  # 保存原图线程标志位，存储30张原图
+    ALL_thread_flag = False  # 所有线程的标志位，单独开设一个线程，监控所有线程是否都结束，如果都结束，则置为False，后续添加线程管理，可以放在这里面
+
+    DebugPrint_flag = True  # 打印标志位
+
+    destory_image = False
+    Camera_device = args.video_num  # 摄像头编号，0001
+    image_time = None  # 当前处理的30视频的编号
+    ImgDealQueue = []  # 750帧视频，缓存空间
+    root = "/data/Hatcher/output2"
+    #os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+    os.environ['CUDA_VISIBLE_DEVICES'] = args.GPU
+    devices = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    #event = threading.Event()
+    eventDensity = threading.Event()
+    eventDetection = threading.Event()
+    eventSaveimage = threading.Event()
+    # Detection2 = Detection(devices,root)

con_time_cut.py

@@ -0,0 +1,16 @@
+import os
+import cv2
+root_path = "/data1/Hatcher/model/data_MPKFNet/con/FFOutput(1)/"
+result_path = "/data1/Hatcher/model/data_MPKFNet/con/result/"
+if not os.path.exists(result_path):
+    os.mkdir(result_path)
+file_list = os.listdir(root_path)
+num = 0
+for file_path in file_list:
+    image_list = os.listdir(os.path.join(root_path,file_path))
+    for image_path in image_list:
+        print(image_path)
+        img = cv2.imread(os.path.join(os.path.join(root_path,file_path),image_path))
+        Img = img.copy()[20:96, 806:903, :]
+        cv2.imwrite(os.path.join(result_path,"{}.jpg".format(str(num).zfill(4))),Img)
+        num+=1

configs/configs.py

@@ -0,0 +1,9 @@
+
+trainA_json = '../data_MPKFNet/data/BirdCounting/BirdCounting2/train_data.json'
+trainB_json = '../data_MPKFNet/data/trainB.json'
+testA_json = '../data_MPKFNet/data/BirdCounting/BirdCounting2/test_data.json'
+testB_json = '../data_MPKFNet/data/testB.json'
+inf_json   ='../data_MPKFNet/data/inf.json'
+epochs =200
+model_head = True
+

creat_xml.py

@@ -0,0 +1,71 @@
+'''
+author ：  Hatcher W
+Email： hcwang@semi.ac.cn
+use: 用于创建存储目标检测结果的xml文件
+input:网络运算输出结果
+output：xml文件
+'''
+from lxml import etree
+
+
+class GEN_Annotations:
+    def __init__(self, filename):
+        self.root = etree.Element("annotation")
+
+        child1 = etree.SubElement(self.root, "folder")
+        child1.text = "VOC2007"
+
+        child2 = etree.SubElement(self.root, "filename")
+        child2.text = filename
+
+        child3 = etree.SubElement(self.root, "source")
+
+        child4 = etree.SubElement(child3, "annotation")
+        child4.text = "PASCAL VOC2007"
+        child5 = etree.SubElement(child3, "database")
+        child5.text = "Unknown"
+
+        child6 = etree.SubElement(child3, "image")
+        child6.text = "flickr"
+        child7 = etree.SubElement(child3, "flickrid")
+        child7.text = "35435"
+
+    def set_size(self, witdh, height, channel):
+        size = etree.SubElement(self.root, "size")
+        widthn = etree.SubElement(size, "width")
+        widthn.text = str(witdh)
+        heightn = etree.SubElement(size, "height")
+        heightn.text = str(height)
+        channeln = etree.SubElement(size, "depth")
+        channeln.text = str(channel)
+
+    def savefile(self, filename):
+        tree = etree.ElementTree(self.root)
+        tree.write(filename, pretty_print=True, xml_declaration=False, encoding='utf-8')
+
+    def add_pic_attr(self, label, xmin, ymin, xmax, ymax):
+        object = etree.SubElement(self.root, "object")
+        namen = etree.SubElement(object, "name")
+        namen.text = label
+        bndbox = etree.SubElement(object, "bndbox")
+        xminn = etree.SubElement(bndbox, "xmin")
+        xminn.text = str(xmin)
+        yminn = etree.SubElement(bndbox, "ymin")
+        yminn.text = str(ymin)
+        xmaxn = etree.SubElement(bndbox, "xmax")
+        xmaxn.text = str(xmax)
+        ymaxn = etree.SubElement(bndbox, "ymax")
+        ymaxn.text = str(ymax)
+
+
+if __name__ == '__main__':
+    filename = "000001.jpg"
+    anno = GEN_Annotations(filename)
+    anno.set_size(1280, 720, 3)
+    for i in range(3):
+        xmin = i + 1
+        ymin = i + 10
+        xmax = i + 100
+        ymax = i + 100
+        anno.add_pic_attr("mouse", xmin, ymin, xmax, ymax)
+    anno.savefile("00001.xml")

creat_xml_fordetection.py

@@ -0,0 +1,119 @@
+'''
+author ：  Hatcher W
+Email： hcwang@semi.ac.cn
+use: 用于创建存储目标检测结果的xml文件
+input:网络运算输出结果
+output：xml文件
+'''
+import numpy as np
+from lxml import etree
+
+
+class GEN_Annotations:
+    def __init__(self, filename):
+        self.root = etree.Element("annotation")
+
+        child1 = etree.SubElement(self.root, "folder")
+        child1.text = "Hatcher"
+
+        child2 = etree.SubElement(self.root, "filename")
+        child2.text = filename
+
+        child3 = etree.SubElement(self.root, "source")
+
+        child4 = etree.SubElement(child3, "annotation")
+        child4.text = "Bird_detection"
+        child5 = etree.SubElement(child3, "database")
+        child5.text = "Unknown"
+
+        child6 = etree.SubElement(child3, "image")
+        child6.text = "flickr"
+        child7 = etree.SubElement(child3, "flickrid")
+        child7.text = "35435"
+
+    def set_size(self, witdh, height, channel):
+        size = etree.SubElement(self.root, "size")
+        widthn = etree.SubElement(size, "width")
+        widthn.text = str(witdh)
+        heightn = etree.SubElement(size, "height")
+        heightn.text = str(height)
+        channeln = etree.SubElement(size, "depth")
+        channeln.text = str(channel)
+
+    def savefile(self, filename):
+        tree = etree.ElementTree(self.root)
+        tree.write(filename, pretty_print=True, xml_declaration=False, encoding='utf-8')
+
+    def add_pic_attr(self, label, xmin, ymin, xmax, ymax,score):
+        object = etree.SubElement(self.root, "object")
+        namen = etree.SubElement(object, "name")
+        namen.text = str(label).replace("tensor(","").replace("., device='cuda:0')","")
+        bndbox = etree.SubElement(object, "bndbox")
+        xminn = etree.SubElement(bndbox, "xmin")
+        xminn.text = str(xmin)
+        yminn = etree.SubElement(bndbox, "ymin")
+        yminn.text = str(ymin)
+        xmaxn = etree.SubElement(bndbox, "xmax")
+        xmaxn.text = str(xmax)
+        ymaxn = etree.SubElement(bndbox, "ymax")
+        ymaxn.text = str(ymax)
+        scoren = etree.SubElement(bndbox, "score")
+        scoren.text = str(score)
+    def add_pic_attr_bird(self, label, xmin, ymin, xmax, ymax):
+        object = etree.SubElement(self.root, "object")
+        namen = etree.SubElement(object, "name")
+        namen.text = label
+        bndbox = etree.SubElement(object, "bndbox")
+        xminn = etree.SubElement(bndbox, "xmin")
+        xminn.text = str(xmin)
+        yminn = etree.SubElement(bndbox, "ymin")
+        yminn.text = str(ymin)
+        xmaxn = etree.SubElement(bndbox, "xmax")
+        xmaxn.text = str(xmax)
+        ymaxn = etree.SubElement(bndbox, "ymax")
+        ymaxn.text = str(ymax)
+    def add_pic_density(self, label, density_h):
+        object = etree.SubElement(self.root, "object")
+        count = etree.SubElement(object, "count")
+        for i in density_h:
+            density = etree.SubElement(object, "density_h")
+            num=0
+            for j in i:
+                density.attrib["att{}".format(num)] = j
+                num+=1
+
+
+
+
+
+def make_xml(filename,size_x,size_y,boxes,labels,scores,savename):
+    anno = GEN_Annotations(filename)
+    anno.set_size(size_y, size_x, 3)
+    for i in range(len(boxes)):
+        xmin = int(boxes[i][0])
+        ymin = int(boxes[i][1])
+        xmax = int(boxes[i][2])
+        ymax = int(boxes[i][3])
+        score = int(scores[i]*1000)
+        anno.add_pic_attr(str(labels[i]), xmin, ymin, xmax, ymax,score)
+    anno.savefile(savename)
+def make_xml_iccv(filename,size_x,size_y,boxes,savename):
+    anno = GEN_Annotations(filename)
+    anno.set_size(size_y, size_x, 3)
+    for i in range(len(boxes)):
+        xmin = int(boxes[i][0])
+        ymin = int(boxes[i][1])
+        xmax = int(boxes[i][2])
+        ymax = int(boxes[i][3])
+        anno.add_pic_attr_bird("bird", xmin, ymin, xmax, ymax)
+    anno.savefile(savename)
+
+
+if __name__ == '__main__':
+    filename = "000001.jpg"
+    anno = GEN_Annotations(filename)
+    anno.set_size(1280, 720, 3)
+    density = np.array([4,4],dtype=np.float64)
+    density = [[0.123,93.132213,0.2132131,0.32132131],[0.123,93.132213,0.2132131,0.32132131],[0.123,93.132213,0.2132131,0.32132131],[0.123,93.132213,0.2132131,0.32132131]]
+    anno.add_pic_density("mouse", density)
+    anno.savefile("00001.xml")

crow_data_generate.py

@@ -0,0 +1,36 @@
+
+import numpy as np
+import os
+import json
+import argparse
+dataset_name = "BirdCounting1"
+root_path ="../data_MPKFNet/data/BirdCounting"
+
+
+parser = argparse.ArgumentParser(description="DensityVideodemo.")
+parser.add_argument("dataset_name", default='', type=str,
+                    help='Video name.')
+
+parser.add_argument(
+    "opts",
+    help="Modify config options using the command-line",
+    default=None,
+    nargs=argparse.REMAINDER,
+)
+args = parser.parse_args()
+print(args)
+dataset_name = args.dataset_name
+datasetjsonpath = os.path.join(root_path ,dataset_name)
+train = os.path.join(root_path,dataset_name,'train_data','images')
+test = os.path.join(root_path,dataset_name,'test_data','images')
+trainfile = os.listdir(train)
+train_json = [os.path.join(train,file) for file in trainfile]
+with open(datasetjsonpath + "/train_data.json",'w') as f:
+    json.dump(train_json,f)
+
+testfile = os.listdir(test)
+test_json = [os.path.join(test,file) for file in testfile]
+with open(datasetjsonpath + "/test_data.json",'w') as f:
+    json.dump(test_json,f)
+
+print("data is ok")

crow_train.py

@@ -0,0 +1,166 @@
+'''
+author ：  Hatcher W
+Email： hcwang@semi.ac.cn
+use: 密度估计训练
+'''
+import torch
+import torch.nn as nn
+from torch import optim
+from torch.utils.data import DataLoader
+import numpy as np
+from network import CSRNet,PKSNet,PKSNet_attention
+from dataset import People_Dataset
+from configs.config import *
+from configs.config import model_head
+import json
+import os
+from tqdm import tqdm
+import cv2
+os.environ['CUDA_VISIBLE_DEVICES'] = '3'
+device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+train_weigth_path = "../data_MPKFNet/weigth/BirdCounting2/"
+if not os.path.exists(train_weigth_path):
+    os.makedirs(train_weigth_path)
+'''
+    准备数据
+'''
+batch_size = 1
+with open(trainA_json, 'r') as outfile:
+    train_list = json.load(outfile)
+with open(testA_json, 'r') as outfile:
+    test_list = json.load(outfile)
+train_dataset = People_Dataset(train_list)
+train_loader = DataLoader(train_dataset, batch_size=batch_size)
+test_dataset = People_Dataset(test_list)
+test_loader = DataLoader(test_dataset, batch_size=1)
+
+'''
+    创建网络
+'''
+net = PKSNet().to(device)
+
+
+'''
+    训练设置
+'''
+lr=1e-10
+criterion_desity = nn.MSELoss(reduction='sum')
+criterion_spice = nn.CrossEntropyLoss()
+criterion_posture = nn.CrossEntropyLoss()
+optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.95, weight_decay=5*1e-4)
+epoch_start = 0
+# optimizer = optim.Adam(net.parameters(), lr=1e-7,  weight_decay=1e-4)
+if True:
+   # net.load_state_dict(torch.load('./weigth/bestmae.pth'))
+    checkpoint = torch.load(train_weigth_path+'final.pth')
+    net.load_state_dict(checkpoint['net'])
+    #optimizer.load_state_dict(checkpoint['optimizer'])
+    #epoch_start = checkpoint['epoch']
+    print("load model weight ok")
+
+'''
+    开始训练
+'''
+mae_best = 1e10
+num_lr=0
+print(net)
+for epoch in range(epoch_start,epochs):
+    net.train()
+    epoch_loss = 0
+    with tqdm(total=len(train_dataset), desc=f'Epoch {epoch + 1}/{epochs}', unit='img') as pbar:
+        if model_head:
+            for img, target, count, posture, species in train_loader:
+                img = img.to(device)
+                target = target.to(device)
+                posture = posture.to(device)
+                species = species.to(device)
+                target_pred, posture_pred, spice_pred = net(img)
+                # print(target_pred.shape, posture_pred.shape, spice_pred.shape)
+                # print(target.shape, posture.shape, species.shape)
+                loss = 0.9*criterion_desity(target_pred, target) + 0.05*criterion_posture(posture_pred, posture) \
+                       + 0.05*criterion_spice(spice_pred, species)
+                epoch_loss += loss.item()
+                pbar.set_postfix({'loss (batch)': loss.item()/img.shape[0]})
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+                pbar.update(img.shape[0])
+        else:
+            for img, target, count in train_loader:
+                img = img.to(device)
+                target = target.to(device)
+                target_pred= net(img)
+                # print(target_pred.shape, posture_pred.shape, spice_pred.shape)
+                # print(target.shape, posture.shape, species.shape)
+                loss = criterion_desity(target_pred, target)
+                epoch_loss += loss.item()
+                pbar.set_postfix({'loss (batch)': loss.item() / img.shape[0]})
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+                pbar.update(img.shape[0])
+    print('train_data: {}, iter loss: {} for single image \n'.format(len(train_list), epoch_loss/len(train_list)))
+
+    mae ,mse= 0,0
+    if model_head:
+        for img, traget, count, posture, species  in tqdm(test_loader):
+            img = img.to(device)
+            target = traget.to(device)
+            targets_pred, posture_pred, spice_pred = net(img)
+            #print(targets_pred.data.sum())
+            #print(target.data.sum())
+            mae += abs(targets_pred.data.sum() - target.data.sum())
+            mse += ((target.data.sum() - targets_pred.data.sum()) * (target.data.sum() - targets_pred.data.sum()))
+    else:
+        for img, traget, count in tqdm(test_loader):
+            img = img.to(device)
+            target = traget.to(device)
+            targets_pred = net(img)
+            #print(targets_pred.data.sum())
+            #print(target.data.sum())
+            mae += abs(targets_pred.data.sum() - target.data.sum())
+            mse += ((target.data.sum() - targets_pred.data.sum()) * (target.data.sum() - targets_pred.data.sum()))
+    mae = mae / len(test_list)
+    mse = mse /len(test_list)
+    if mae < mae_best:
+        num_lr=0
+        print(epoch)
+        state = {
+            "epoch" :epoch,
+            "net"   :net.state_dict(),
+            "optimizer":optimizer.state_dict()
+
+        }
+        torch.save(state, train_weigth_path+'bestmaesheep{},mse{}.pth'.format(int(mae),int(mse)))
+        torch.save(state, train_weigth_path+'bestmaesheep.pth'.format(int(mae)))
+        mae_best = mae
+    else:
+        num_lr+=1
+    if num_lr>5:
+        lr*=0.1
+        num_lr = 0
+        if lr < 1e-10:
+            lr = 1e-10
+        print("lr is down to :"+str(lr))
+        optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.95, weight_decay=5 * 1e-4)
+    print('test_data:{},calculate MAE: \n'.format(len(test_list)))
+    print(' * MAE {mae:.3f} \n'.format(mae=mae))
+    print(' * MSE {mse:.3f} \n'.format(mse=mse))
+
+
+'''
+    模型加载与单张图片测试
+'''
+# net = CSRNet()
+# net.load_state_dict(torch.load('CSRnet.pth'))
+# net.eval()
+# test_img = cv2.imread('./data/part_A_final/test_data/images/IMG_1.jpg.jpg')
+# test_img = cv2.resize(test_img, (224, 224))
+# test_img = np.transpose(test_img, [2, 0, 1])
+# test_img = np.expand_dims(test_img,axis=0).astype(np.float)
+# with torch.no_grad():
+#     out = net(torch.from_numpy(test_img).float())
+#     out = out.squeeze(0).numpy()
+# cv2.imwrite('result.png',out)
+
+

dataset.py

@@ -0,0 +1,200 @@
+'''
+author ：  Hatcher W
+Email： hcwang@semi.ac.cn
+use: 密度估计数据集预处理
+'''
+from torch.utils.data import Dataset
+import random
+import numpy as np
+import cv2
+from scipy.io import loadmat
+import json
+from configs.config import *
+from configs.config import model_head
+import matplotlib.pyplot as plt
+import torch
+
+class People_Dataset(Dataset):
+    def __init__(self, img_list, train=True):
+        random.shuffle(img_list)
+        self.img_list = img_list
+        self.train = train
+
+    def get_pixels(self, target_path, target_h, target_w):
+        mat_data = loadmat(target_path)
+        positions = mat_data['image_info'][0, 0][0, 0][0]
+        count = mat_data['image_info'][0][0][0][0][1][0][0]
+        if model_head:
+            posture = mat_data['image_info'][0][0][0][0][2][0][0]
+            species = mat_data['image_info'][0][0][0][0][3][0][0]
+            if (species >= 0) & (species <= 100):
+                species = 0
+            if (species >= 101) & (species <= 200):
+                species = 1
+            if (species >= 201) & (species <= 300):
+                species = 2
+            if (species >= 301) & (species <= 400):
+                species = 3
+            if (species >= 401) & (species <= 500):
+                species = 4
+            if (species >= 501) & (species <= 600):
+                species = 5
+
+            if (species >= 601) & (species <= 700):
+                species = 6
+            if (species >= 701) & (species <= 800):
+                species = 7
+            if (species >= 801) & (species <= 900):
+                species = 8
+            if (species >= 901) & (species <= 1000):
+                species = 9
+            if (species >= 1001) & (species <= 1100):
+                species = 10
+            if (species >= 1101) & (species <= 1200):
+                species = 11
+
+            if (species >= 1201):
+                species = 12
+        pixels = np.zeros((768, 1024))
+        for p in positions:
+            p[0], p[1] = int(p[0] * 1024 / target_w), int(p[1] * 768 / target_h)
+            if p[0] >= 1024 or p[1] >= 768:
+                pass
+            else:
+                pixels[int(p[1]), int(p[0])] += 1
+        density = cv2.GaussianBlur(pixels, (15, 15), 0)
+
+
+        if model_head:
+            temp = (density, count,posture,species)
+        else:
+            temp = (density,count)
+        return temp
+
+    def __len__(self):
+        return len(self.img_list)
+
+    def __getitem__(self, index):
+        assert index <= len(self), 'index range error'
+        img_path = self.img_list[index]
+        img_temppath = img_path.replace('../data_MPKFNet/data/test_data/images/','')
+        target_path = img_path.replace('.jpg', '.mat').replace('images', 'ground_truth').replace('HatIMG_', 'GT_HatIMG_')
+        img = cv2.imread(img_path)
+        #print(img_path)
+        if img is None:
+            print(img_path)
+        target_h, target_w = img.shape[0], img.shape[1]
+
+        img = cv2.resize(img, (1024, 768)) / 255.0
+        # img[:, :, 0] = (img[:,:,0] - 0.485) / 0.229
+        # img[:, :, 1] = (img[:, :, 0] - 0.456) / 0.224
+        # img[:, :, 2] = (img[:, :, 0] - 0.406) / 0.225
+        if model_head:
+            density, count,posture, species= self.get_pixels(target_path, target_h, target_w)
+        else:
+            density, count = self.get_pixels(target_path, target_h, target_w)
+        #density = cv2.resize(density, (1024//16, 768//16))*256
+        density = cv2.resize(density, (1024//8, 768//8))*64
+        #density = cv2.resize(density, (1024 // 4, 768 // 4)) *16
+        #density = cv2.resize(density, (1024 // 2, 768 // 2)) * 4
+        #density = cv2.resize(density, (1024 , 768 ))
+        density = np.expand_dims(density, 0)
+        img = np.transpose(img, [2, 0, 1])
+        if model_head:
+            temp = (torch.from_numpy(img).type(torch.FloatTensor), torch.from_numpy(density).type(torch.FloatTensor), torch.tensor(int(count)),\
+               torch.tensor(np.array(posture)).type(torch.long),torch.tensor(np.array(species)).type(torch.long))
+        else:
+            temp = (torch.from_numpy(img).type(torch.FloatTensor), torch.from_numpy(density).type(torch.FloatTensor), torch.tensor(int(count)))
+        return temp
+
+
+class INf_Dataset(Dataset):
+    def __init__(self, img_list, train=True):
+        random.shuffle(img_list)
+        self.img_list = img_list
+        self.train = train
+
+    def __len__(self):
+        return len(self.img_list)
+
+    def __getitem__(self, index):
+        assert index <= len(self), 'index range error'
+        img_path = self.img_list[index]
+        img_temppath = img_path.replace('../data_MPKFNet/data/test_data/images/','')
+        target_path = img_path.replace('.jpg', '.mat').replace('images', 'ground_truth').replace('HatIMG_', 'GT_HatIMG_')
+        img = cv2.imread(img_path)
+
+        target_h, target_w = img.shape[0], img.shape[1]
+        img = cv2.resize(img, (1024, 768)) / 255.0
+        # img[:, :, 0] = (img[:,:,0] - 0.485) / 0.229
+        # img[:, :, 1] = (img[:, :, 0] - 0.456) / 0.224
+        # img[:, :, 2] = (img[:, :, 0] - 0.406) / 0.225
+        if model_head:
+            density, count, posture, species = self.get_pixels(target_path, target_h, target_w)
+        else:
+            density, count = self.get_pixels(target_path, target_h, target_w)
+        density = cv2.resize(density, (1024//8, 768//8))*64
+        #density = cv2.resize(density, (1024,768))
+        density = np.expand_dims(density, 0)
+        img = np.transpose(img, [2, 0, 1])
+        if model_head:
+            temp = (torch.from_numpy(img).type(torch.FloatTensor), torch.from_numpy(density).type(torch.FloatTensor),
+                    torch.tensor(int(count)), \
+                    torch.tensor(np.array(posture)).type(torch.long), torch.tensor(np.array(species)).type(torch.long))
+        else:
+            temp = (torch.from_numpy(img).type(torch.FloatTensor), torch.from_numpy(density).type(torch.FloatTensor),
+                    torch.tensor(int(count)))
+        return temp
+
+if __name__ == '__main__':
+    with open(trainA_json, 'r') as outfile:
+        train_list = json.load(outfile)
+        count_all =0
+    dataset = People_Dataset(train_list)
+    if model_head:
+        for img, target, count,posture, species in dataset:
+            #print(img.shape, target.shape)
+            img = img.numpy()
+
+            target = target.squeeze(0).numpy()
+
+            count = count.numpy()
+            # posture = posture.numpy()
+            # species = species.numpy()
+            # print(posture)
+            # print(species)
+            # print("posture is {} ,species is {}".format(posture,species))
+            img = np.transpose(img, [1,2,0])
+            #
+            # plt.figure('img')
+            # plt.imshow((img[:, :, ::-1]*255).astype(np.uint8))
+            # plt.show()
+            # plt.figure('target')
+            # plt.imshow(target)
+            # plt.show()
+            print(abs(target.sum() - count))
+            count_all =count_all+ abs(target.sum() - count)
+        print(count_all)
+           # print(target.sum() - count)
+    else:
+        for img, target, count in dataset:
+            print(img.shape, target.shape)
+            img = img.numpy()
+
+            target = target.squeeze(0).numpy()
+
+            count = count.numpy()
+
+            img = np.transpose(img, [1, 2, 0])
+            #
+            plt.figure('img')
+            plt.imshow((img[:, :, ::-1] * 255).astype(np.uint8))
+            plt.show()
+            plt.figure('target')
+            plt.imshow(target)
+            plt.show()
+            print(target.sum() - count)
+
+
+
+

detection.py

@@ -0,0 +1,522 @@
+'''
+author ：  Hatcher W
+Email： hcwang@semi.ac.cn
+use: 目标检测推理
+input:需要推理的文件目录
+output：推理结果
+'''
+import glob
+import os
+import time
+import cv2
+import torch
+from torch.autograd import Variable
+from PIL import Image
+from vizer.draw import draw_boxes
+import random
+from ssd.config import cfg
+from ssd.data.datasets import COCODataset, VOCDataset
+import argparse
+import numpy as np
+
+from ssd.data.transforms import build_transforms
+from ssd.modeling.detector import build_detection_model
+from ssd.utils import mkdir
+from ssd.utils.checkpoint import CheckPointer
+from creat_xml_fordetection import GEN_Annotations,make_xml
+from Dcl.models.LoadModel import MainModel
+from Dcl.DclConfig import LoadConfig, load_data_transformers
+os.environ['CUDA_VISIBLE_DEVICES'] = '7'
+from PIL import ImageFile
+import time
+import cv2
+import torch,os
+import numpy as np
+from numpy import random
+import os,sys
+root_path = os.getcwd()
+sys.path.insert(0,root_path+"/Y5")
+from Y5.models.experimental import attempt_load
+from Y5.utils.datasets import  letterbox
+from Y5.utils.general import   non_max_suppression, \
+    scale_coords
+from Y5.utils.plots import plot_one_box
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+fine_class_names = ("Anser albifrons","Fulica atra","Leucogeranus leucogeranus","Egretta garzetta","Platalea leucorodia","Mergus albellus","Grus monacha","Tringa ochropus","Numenius arquata","Anser indicus","Anas zonorhyncha","Ardea cinerea","Anas strepera","Mareca penelope","Tadorna ferruginea","Ardea alba","Numenius madagascariensis","Botaurus stellaris","Ciconia boyciana","Anser fabalis","Recurvirostra avosetta","Podiceps cristatus","Vanellus vanellus","Aythya fuligula","Tringa erythropus","Calidris alpina","Ciconia nigra","Gallinula chloropus","Porphyrio poliocephalus","Larus ridibundus","Anser cygnoides","Charadrius alexandrinus","Pluvialis squatarola","Grus grus" ,"Anser anser","Pluvialis fulva","Charadrius dubius","Anas crecca","Anas platyrhynchos","Spatula clypeata","Phalacrocorax carbo","Mergus merganser","Tadorna tadorna","Tringa nebulariaw","Bucephala clangula","Gallinago gallinago","Anser erythropus","Cygnus columbianus","Ichthyaetus ichthyaetus","Anas acuta","Ardea intermedia","Lalage melaschistos","Zosterops japonicus","Acridotheres cristatellus","Circus spilonotus","Motacilla alba","Garrulax sannio","Pycnonotus sinensis","Aythya nyroca","Lonchura striata","Turdus eunomus","Limosa lapponica","Lonchura punctulata","Phoenicurus auroreus","Ardea purpurea","Aethopyga christinae","Ardeola bacchus","Prinia inornata","Cuculus canorus","Parus major","Upupa epops","Asio fiammeus","Anthus roseatus","Phylloscopus reguloides","Phylloscopus fuscatus","Himantopus himantopus","Saxicola torquata","Dicrurus macrocercus","Eophona migratoria","Milvus migrans","Oriolus chinensis","Anthus cervinus","Falco tinnunculus","Stachyris ruficeps","Aegithalos concinnus","Turdus naumanni","Urocissa erythrorhyncha","Sibirionetta formosa","Phasianus colchicus","Anthus rubescens","Parus venustulus","Motacilla citreola","Emberiza aureola","Motacilla cinerea","Sturnus cineraceus","Picus canus","Vanellus cinereus","Emberiza spodocephala","Bambusicola thoracicus","Streptopelia tranquebarica","Actitis hypoleucos","Hirundo rustica","Carduelis sinica","Hirundo daurica","Luscinia svecicus","Monticola solitarius","Anthus richardi","Ixobrychus cinnamomeus","Spizixos semitorques","Mareca falcata","Parus monticolus","Passer montanus","Bubulcus coromandus","Turdus merula","Buteo buteo","Alcedo atthis","Horornis fortipes","Accipiter nisus","Copsychus saularis","Circus melanoleucos","Anthus hodgsoni","Anthus spinoletta","Spodiopsar sericeus","Cuculus micropterus","Motacilla flava","Pica pica","Tachybaptus ruficollis","Emberiza pusilla","Alauda gulgula","Lanius sphenocercus","Riparia riparia","Nycticorax nycticorax","Terpsiphone paradisi","Cuculus sparverioides","Falco Peregrinus","Aix galericulata","Charadrius placidus","Gallinago stenura","Mergus squamatus","Spilopelia chinensis","Lanius schach","Pomatorhinus ruficollis","Abroscopus albogularis","Turnix suscitator","Cisticola juncidis","Paradoxornis webbianus")
+class Detector:
+
+    def __init__(self, device):
+        super(Detector, self).__init__()
+        self.img_size = 640
+        self.conf_thres = 0.3  #0.3
+        self.iou_thres = 0.4   #0.4
+        self.stride = 1
+        #self.weights = "/data1/Hatcher/model/yolov5/runs/train/exp15/weights/best.pt"
+        #self.weights = '/data1/Hatcher/model/yolov5/runs/train/exp10/weights/best.pt'
+        self.weights = '/data1/Hatcher/model/data_MPKFNet/weigth/detection/Y/best.pt'
+        self.device = device
+        #self.device = select_device(self.device)
+        model = attempt_load(self.weights, map_location=self.device)
+        model.to(self.device).eval()
+        model.half()
+        self.model = model
+        self.names = model.module.names if hasattr(
+            model, 'module') else model.names
+
+    def preprocess(self, img):
+
+        img0 = img.copy()
+
+        img = letterbox(img, new_shape=self.img_size, stride=self.stride)[0]
+        img = img[:, :, ::-1].transpose(2, 0, 1)
+        img = np.ascontiguousarray(img)
+        img = torch.from_numpy(img).to(self.device)
+        img = img.half()  # °ë¾«¶È
+        img /= 255.0  # ͼÏñ¹éÒ»»¯
+        if img.ndimension() == 3:
+            img = img.unsqueeze(0)
+
+        return img0, img
+    #im 原图  imo 1024 img 640
+    def detect(self, im):
+        img_detect = cv2.resize(im, (1024, 768))
+        im0, img = self.preprocess(img_detect)
+
+        pred = self.model(img, augment=False)[0]
+        pred = pred.float()
+        pred = non_max_suppression(pred, self.conf_thres, self.iou_thres)
+
+        pred_boxes = []
+        scores = []
+        labels = []
+        for det in pred:
+
+            if det is not None and len(det):
+                det[:, :4] = scale_coords(
+                    img.shape[2:], det[:, :4], im0.shape).round()
+
+                pred_boxes.append(det[:,:4])
+                scores.append(det[:,4])
+                labels.append((det[:,-1]))
+                pred_boxes = pred_boxes[0]
+                scores = scores[0]
+                labels = labels[0]
+                # for *x, conf, cls_id in det:
+                #     lbl = self.names[int(cls_id)]
+                #     if not lbl in ['person', 'car', 'truck']:
+                #         continue
+                #     x1, y1 = int(x[0]), int(x[1])
+                #     x2, y2 = int(x[2]), int(x[3])
+                #     pred_boxes.append(
+                #         (x1, y1, x2, y2, lbl, conf))
+
+        return im, pred_boxes,scores,labels
+class args_temp():
+    dataset = 'HCData'
+    backbone = 'resnet50'
+    batch_size = 1
+    num_workers = 8
+    version = 'val'
+    resize_resolution = 512
+    crop_resolution = 448
+    save_suffix = None
+    swap_num = [7,7]
+class Fine_grained:
+    def __init__(self,device):
+        self.device = device
+        self.DclPreWeightPath = '../data_MPKFNet/weigth/FinedGrainedRecognition/Test/weights_19_499_0.9998_1.0000.pth'
+        # parser = argparse.ArgumentParser(description='dcl parameters')
+        # parser.add_argument('--data', dest='dataset', default='HCData', type=str)
+        # parser.add_argument('--backbone', dest='backbone', default='resnet50', type=str)
+        # parser.add_argument('--b', dest='batch_size', default=1, type=int)
+        # parser.add_argument('--nw', dest='num_workers', default=8, type=int)
+        # parser.add_argument('--ver', dest='version', default='val', type=str)
+        # parser.add_argument('--size', dest='resize_resolution', default=512, type=int)
+        # parser.add_argument('--crop', dest='crop_resolution', default=448, type=int)
+        # parser.add_argument('--ss', dest='save_suffix', default=None, type=str)
+        # parser.add_argument('--acc_report', dest='acc_report', action='store_true')
+        # parser.add_argument('--swap_num', default=[7, 7], nargs=2, metavar=('swap1', 'swap2'), type=int,
+        #                     help='specify a range')
+        self.args = args_temp()
+        self.Config = LoadConfig(self.args, 's')
+        self.transformers = load_data_transformers(self.args.resize_resolution, self.args.crop_resolution, self.args.swap_num)
+        self.net = MainModel(self.Config).to(device)
+        model_dict = self.net.state_dict()
+        pretrained_dict = torch.load(self.DclPreWeightPath)
+        model_dict.update(pretrained_dict)
+        self.net.load_state_dict(model_dict)
+
+        print("load model weight is ok")
+        self.net.train(False)
+    def inference(self,img,imagelist,labellist,scorelist):
+        with torch.no_grad():
+            ratio_x = (3840.0 / 1024.0)
+            ratio_y = (2160.0 / 768.0)
+            for i,box in enumerate(imagelist):
+                # box1=[0,0,0,0]
+                # box1[1] = box[1]*ratio_y
+                # box1[3] = box[3]*ratio_y
+                # box1[0] = box[0]*ratio_x
+                # box1[2] = box[2] * ratio_x
+                # Img = img.copy()[int(box1[1]):int(box1[3]), int(box1[0]):int(box1[2]), :]
+                # if (int(box1[3])-int(box1[1])>10) and (int(box1[2])-int(box1[0])>10):
+                Img = img.copy()[int(box[1]):int(box[3]), int(box[0]):int(box[2]), :]
+                #cv2.imwrite("/data/Hatcher/detectout6/{}.jpg".format(str(i)),Img)
+                #height_or, width_or = Img.shape[:2]
+                #Img = cv2.resize(Img,(int(width_or/ratio_x),int(height_or/ratio_y)))
+                if (int(box[3]) - int(box[1]) > 10) and (int(box[2]) - int(box[0]) > 10):
+                    t1 = time.time()
+                    Img = Image.fromarray(cv2.cvtColor(Img, cv2.COLOR_BGR2RGB))
+                    Img = self.transformers['test_totensor'](Img)
+                    Img = Img.unsqueeze(0)
+                    inputs = Variable(Img.cuda())
+                    t2 = time.time()
+
+                    print("fine data deal is ", (t2 - t1) * 1000)
+                    outputs = self.net(inputs)
+                    t3 = time.time()
+
+                    print("fine inference is ", (t3 - t2) * 1000)
+                    top3_val, top3_pos = torch.topk(outputs[0], 3)  # topk返回的是一个元组，第一个元素是每一行最大的三个值（默认是按行取最大值），第二个元素是每一行最大的值的排序
+                    Species = top3_pos.cpu().numpy()[0][0]
+                    ConfidenceLevel = top3_val.cpu().numpy()[0][0]
+                    ConfidenceLevel = ConfidenceLevel/26.0
+                    if ConfidenceLevel >1:
+                        ConfidenceLevel = 1.0* random.uniform(0, 1)
+
+                    print("hhhhhhhhhhhhhhhhhhhhhhh is ",ConfidenceLevel)
+                    labellist[i] = Species
+                    scorelist[i] = ConfidenceLevel
+                    t4 = time.time()
+
+                    print("fine data final is ", (t4 - t3) * 1000)
+
+            return imagelist, labellist, scorelist
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="SSD Demo.")
+    parser.add_argument(
+        "--config-file",
+        default="ssd/config/hat-512.yaml",
+        metavar="FILE",
+        help="path to config file",
+        type=str,
+    )
+    parser.add_argument("--ckpt", type=str, default="../data_MPKFNet/weigth/detection/detection.pth", help="Trained weights.")
+    parser.add_argument("--score_threshold", type=float, default=0.3)
+    parser.add_argument("--images_dir", default='/data1/Hatcher/model/data_MPKFNet/data/VOC/VOCdevkit/VOC16/JPEGImages/', type=str, help='Specify a image dir to do prediction.')
+    parser.add_argument("--output_dir", default='/data/Hatcher/detectout13', type=str, help='Specify a image dir to save predicted images.')
+    parser.add_argument("--dataset_type", default="voc", type=str, help='Specify dataset type. Currently support voc and coco.')
+
+    parser.add_argument(
+        "opts",
+        help="Modify config options using the command-line",
+        default=None,
+        nargs=argparse.REMAINDER,
+    )
+    args = parser.parse_args()
+    return args
+
+@torch.no_grad()
+def run_demo_SSD(cfg, ckpt, score_threshold, images_dir, output_dir, dataset_type):
+    if dataset_type == "voc":
+        class_names = VOCDataset.class_names
+    elif dataset_type == 'coco':
+        class_names = COCODataset.class_names
+    else:
+        raise NotImplementedError('Not implemented now.')
+    device = torch.device(cfg.MODEL.DEVICE)
+
+    model = build_detection_model(cfg)
+    model = model.to(device)
+    checkpointer = CheckPointer(model, save_dir=cfg.OUTPUT_DIR)
+    checkpointer.load(ckpt, use_latest=ckpt is None)
+    weight_file = ckpt if ckpt else checkpointer.get_checkpoint_file()
+    print('Loaded weights from {}'.format(weight_file))
+
+    image_paths = glob.glob(os.path.join(images_dir, '*.jpg'))
+    mkdir(output_dir)
+
+    cpu_device = torch.device("cpu")
+    transforms = build_transforms(cfg, is_train=False)
+    model.eval()
+    for i, image_path in enumerate(image_paths):
+        start = time.time()
+        image_name = os.path.basename(image_path)
+        a = Image.open(image_path)
+        image = np.array(Image.open(image_path).convert("RGB"))
+        height, width = image.shape[:2]
+        images = transforms(image)[0].unsqueeze(0)
+        load_time = time.time() - start
+
+        start = time.time()
+        result = model(images.to(device))[0]
+        inference_time = time.time() - start
+
+        result = result.resize((width, height)).to(cpu_device).numpy()
+        boxes, labels, scores = result['boxes'], result['labels'], result['scores']
+
+        indices = scores > score_threshold
+        boxes = boxes[indices]
+        labels = labels[indices]
+        scores = scores[indices]
+        meters = ' | '.join(
+            [
+                'objects {:02d}'.format(len(boxes)),
+                'load {:03d}ms'.format(round(load_time * 1000)),
+                'inference {:03d}ms'.format(round(inference_time * 1000)),
+                'FPS {}'.format(round(1.0 / inference_time))
+            ]
+        )
+        print('({:04d}/{:04d}) {}: {}'.format(i + 1, len(image_paths), image_name, meters))
+
+        drawn_image = draw_boxes(image, boxes, labels, scores, class_names).astype(np.uint8)
+        xml_outputpath = os.path.join(output_dir, image_name).replace("/result/","/result_xml/").replace(".jpg",".xml")
+        make_xml(image_name,image.shape[0],image.shape[1],boxes,labels,scores,xml_outputpath,class_names)
+        Image.fromarray(drawn_image).save(os.path.join(output_dir, image_name))
+@torch.no_grad()
+def run_demo_fine_SSD(cfg, ckpt, score_threshold, images_dir, output_dir, dataset_type):
+    if dataset_type == "voc":
+        class_names = VOCDataset.class_names
+    elif dataset_type == 'coco':
+        class_names = COCODataset.class_names
+    else:
+        raise NotImplementedError('Not implemented now.')
+    device = torch.device(cfg.MODEL.DEVICE)
+    fine1 = Fine_grained(device)
+    model = build_detection_model(cfg)
+    model = model.to(device)
+    checkpointer = CheckPointer(model, save_dir=cfg.OUTPUT_DIR)
+    checkpointer.load(ckpt, use_latest=ckpt is None)
+    weight_file = ckpt if ckpt else checkpointer.get_checkpoint_file()
+    print('Loaded weights from {}'.format(weight_file))
+
+    image_paths = glob.glob(os.path.join(images_dir, '*.jpg'))
+    mkdir(output_dir)
+
+    cpu_device = torch.device("cpu")
+    transforms = build_transforms(cfg, is_train=False)
+    model.eval()
+    for i_D, image_path in enumerate(image_paths):
+        start = time.time()
+        image_name = os.path.basename(image_path)
+        a = Image.open(image_path)
+        image = np.array(Image.open(image_path).convert("RGB"))
+        #image_or = np.array(Image.open(image_path.replace("/JPEGImages/","/JPEGOriginal/")).convert("RGB"))
+        height, width = image.shape[:2]
+        images = transforms(image)[0].unsqueeze(0)
+        load_time = time.time() - start
+
+        start = time.time()
+        result = model(images.to(device))[0]
+        inference_time = time.time() - start
+
+        result = result.resize((width, height)).to(cpu_device).numpy()
+        boxes, labels, scores = result['boxes'], result['labels'], result['scores']
+
+        indices_score = (scores > score_threshold)
+        indices_label_left = labels < 4
+        indices_label_right = labels > 2
+        indices = []
+        for i in range(len(indices_score)):
+            indices.append(indices_score[i] and indices_label_left[i] and indices_label_right[i])
+        boxes = boxes[indices]
+        labels = labels[indices]
+        scores = scores[indices]
+        for i in range(len(boxes)):
+            if boxes[i][0] < 0:
+                boxes[i][0] = 0
+            if boxes[i][1] < 0:
+                boxes[i][1] = 0
+            if boxes[i][2] < 0:
+                boxes[i][2] = 0
+            if boxes[i][3] < 0:
+                boxes[i][3] = 0
+        for i in range(len(labels)):
+            labels[i] = 99999
+        boxes_fine, labels_fine, scores_fine = fine1.inference(image, boxes, labels, scores)
+        meters = ' | '.join(
+            [
+                'objects {:02d}'.format(len(boxes_fine)),
+                'load {:03d}ms'.format(round(load_time * 1000)),
+                'inference {:03d}ms'.format(round(inference_time * 1000)),
+                'FPS {}'.format(round(1.0 / inference_time))
+            ]
+        )
+        print('({:04d}/{:04d}) {}: {}'.format(i_D + 1, len(image_paths), image_name, meters))
+
+        drawn_image = draw_boxes(image, boxes_fine, labels_fine, scores_fine).astype(np.uint8)
+        xml_outputpath = os.path.join(output_dir, image_name).replace("/result/","/result_xml/").replace(".jpg",".xml")
+        #make_xml(image_name,image.shape[0],image.shape[1],boxes,labels,scores,xml_outputpath,class_names)
+        Image.fromarray(drawn_image).save(os.path.join(output_dir, image_name))
+def run_demo_fine_Y5(cfg, ckpt, score_threshold, images_dir, output_dir, dataset_type):
+    device = torch.device(cfg.MODEL.DEVICE)
+    fine1 = Fine_grained(device)
+    detector = Detector(device)
+    image_paths = glob.glob(os.path.join(images_dir, '*.jpg'))
+    mkdir(output_dir)
+
+    cpu_device = torch.device("cpu")
+
+    for i_D, image_path in enumerate(image_paths):
+
+        image_name = os.path.basename(image_path)
+        #a = Image.open(image_path)
+        img = cv2.imread(image_path)  # BG
+        height_or, width_or = img.shape[:2]
+        t3 = time.time()
+        im0, pred_boxes, scores, labels = detector.detect(img)
+        t4 = time.time()
+        print("推理： ", (t4 - t3) * 1000)
+        print(type(labels))
+        if len(labels) == 0:
+            continue
+        indices_label_left = labels < 1
+        indices_label_right = labels <1
+        indices = []
+        for i in range(len(indices_label_left)):
+            indices.append(indices_label_left[i] and indices_label_right[i])
+
+        pred_boxes = pred_boxes[indices]
+        labels = labels[indices]
+        scores = scores[indices]  # 因为莫名奇妙多了一个list维度，所以加了这个0
+        if len(labels)>0:
+            boxes = pred_boxes[0]
+            labels = labels[0]
+            scores = scores[0]
+            print(len(scores))
+            ratio_x = (width_or / 1024.0)
+            ratio_y = (height_or / 768.0)
+
+            for i in range(len(boxes)):
+                boxes[i][0] = boxes[i][0] * ratio_x
+                boxes[i][2] = boxes[i][2] * ratio_x
+                boxes[i][1] = boxes[i][1] * ratio_y
+                boxes[i][3] = boxes[i][3] * ratio_y
+                if boxes[i][0] < 0:
+                    boxes[i][0] = 0
+                if boxes[i][1] < 0:
+                    boxes[i][1] = 0
+                if boxes[i][2] < 0:
+                    boxes[i][2] = 0
+                if boxes[i][3] < 0:
+                    boxes[i][3] = 0
+            for i in range(len(labels)):
+                labels[i] = 99999
+            boxes_fine, labels_fine, scores_fine = fine1.inference(img, boxes, labels, scores)
+            meters = ' | '.join(
+                [
+                    'objects {:02d}'.format(len(boxes_fine)),
+                ]
+            )
+            print('({:04d}/{:04d}) {}: {}'.format(i_D + 1, len(image_paths), image_name, meters))
+
+            for i in range(len(scores)):
+                xyxy = boxes[i]
+                conf = scores[i]
+                cls = labels[i]
+                print(cls)
+                #label = f'{int(cls)} {conf:.2f}'
+                #if int(cls)==3:
+                if int(cls)==99999:
+                    label = f'{int(cls)} {conf:.2f}'
+                else:
+                    label = f'{fine_class_names[int(cls)]} {conf:.2f}'
+                # if int(cls) == 11:
+                #     label = f'{"苍鹭"} {conf:.2f}'
+                # if label == 3:
+                #     label = "白鹭"
+                # if label ==11:
+                #     label = "苍鹭"
+                palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1)
+                colors = [int((p * (cls ** 2 - cls + 1)) % 255) for p in palette]
+                plot_one_box(xyxy, img, label=label, color=colors, line_thickness=3)
+        else:
+            boxes, labels, scores = [],[],[]
+        #drawn_image = draw_boxes(image, boxes_fine, labels_fine, scores_fine).astype(np.uint8)
+        xml_outputpath = os.path.join(output_dir, image_name).replace("/result/","/result_xml/").replace(".jpg",".xml")
+        make_xml(image_name,im0.shape[0],im0.shape[1],boxes,labels,scores,xml_outputpath)
+        cv2.imwrite(os.path.join(output_dir, image_name),im0)
+        #Image.fromarray(drawn_image).save(os.path.join(output_dir, image_name))
+def run_demo_Y5(cfg, ckpt, score_threshold, images_dir, output_dir, dataset_type):
+    device = torch.device(cfg.MODEL.DEVICE)
+    #fine1 = Fine_grained(device)
+    detector = Detector(device)
+    image_paths = glob.glob(os.path.join(images_dir, '*.jpg'))
+    mkdir(output_dir)
+
+    cpu_device = torch.device("cpu")
+
+    for i_D, image_path in enumerate(image_paths):
+
+        image_name = os.path.basename(image_path)
+        print(image_name)
+        #a = Image.open(image_path)
+        img = cv2.imread(image_path)  # BG
+        height_or, width_or = img.shape[:2]
+        t3 = time.time()
+        im0, pred_boxes, scores, labels = detector.detect(img)
+        t4 = time.time()
+        print("推理： ", (t4 - t3) * 1000)
+        print(type(labels))
+        if len(labels) == 0:
+            continue
+        meters = ' | '.join(
+            [
+                'objects {:02d}'.format(len(pred_boxes)),
+            ]
+        )
+        print('({:04d}/{:04d}) {}: {}'.format(i_D + 1, len(image_paths), image_name, meters))
+        ratio_x = (width_or / 1024.0)
+        ratio_y = (height_or/ 768.0)
+
+
+        for i in range(len(scores)):
+            xyxy = pred_boxes[i]
+            xyxy[0] = xyxy[0] * ratio_x
+            xyxy[2] = xyxy[2] * ratio_x
+            xyxy[1] = xyxy[1] * ratio_y
+            xyxy[3] = xyxy[3] * ratio_y
+            conf = scores[i]
+            cls = labels[i]
+            print(cls)
+            label = f'{int(cls)} {conf:.2f}'
+            palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1)
+            colors = [int((p * ((cls+10) ** 2 - cls + 1)) % 255) for p in palette]
+            plot_one_box(xyxy, im0, label=label, color=colors, line_thickness=6)
+        #drawn_image = draw_boxes(image, boxes_fine, labels_fine, scores_fine).astype(np.uint8)
+        xml_outputpath = os.path.join(output_dir, image_name).replace("/result/","/result_xml/").replace(".jpg",".xml")
+        make_xml(image_name,im0.shape[0],im0.shape[1],pred_boxes,labels,scores,xml_outputpath)
+        cv2.imwrite(os.path.join(output_dir, image_name),im0)
+        #Image.fromarray(drawn_image).save(os.path.join(output_dir, image_name))
+
+
+def main(args):
+
+    cfg.merge_from_file(args.config_file)
+    cfg.merge_from_list(args.opts)
+    cfg.freeze()
+
+    # print("Loaded configuration file {}".format(args.config_file))
+    # with open(args.config_file, "r") as cf:
+    #     config_str = "\n" + cf.read()
+    #     print(config_str)
+    # print("Running with config:\n{}".format(cfg))
+    if not os.path.exists(args.output_dir):
+       os.mkdir(args.output_dir)
+    output_temp = args.output_dir
+    output_temp = output_temp.replace("/result/","/result_xml/")
+    if not os.path.exists(args.output_dir):
+       os.mkdir(args.output_dir)
+    if not os.path.exists(output_temp):
+       os.mkdir(output_temp)
+    run_demo_fine_Y5(cfg=cfg,
+             ckpt=args.ckpt,
+             score_threshold=args.score_threshold,
+             images_dir=args.images_dir,
+             output_dir=args.output_dir,
+             dataset_type=args.dataset_type)
+
+
+if __name__ == '__main__':
+    main(parse_args())

detection_negative_samples.py

@@ -0,0 +1,2 @@
+import os
+import cv2

detection_test.py

@@ -0,0 +1,99 @@
+'''
+author ：  Hatcher W
+Email： hcwang@semi.ac.cn
+use: 目标检测测试结果
+
+'''
+import argparse
+import logging
+import os
+
+import torch
+import torch.utils.data
+
+from ssd.config import cfg
+from ssd.engine.inference import do_evaluation,do_evaluation_byHatcherW
+from ssd.modeling.detector import build_detection_model
+from ssd.utils import dist_util
+from ssd.utils.checkpoint import CheckPointer
+from ssd.utils.dist_util import synchronize
+from ssd.utils.logger import setup_logger
+
+
+def evaluation_byHatcherW(cfg, ckpt, distributed) :
+    logger = logging.getLogger("SSD.inference")
+
+    model = build_detection_model(cfg)
+    checkpointer = CheckPointer(model, save_dir=cfg.OUTPUT_DIR, logger=logger)
+    device = torch.device(cfg.MODEL.DEVICE)
+    model.to(device)
+    checkpointer.load(ckpt, use_latest=ckpt is None)
+    do_evaluation_byHatcherW(cfg, model, distributed)
+def evaluation(cfg, ckpt, distributed):
+    logger = logging.getLogger("SSD.inference")
+
+    model = build_detection_model(cfg)
+    checkpointer = CheckPointer(model, save_dir=cfg.OUTPUT_DIR, logger=logger)
+    device = torch.device(cfg.MODEL.DEVICE)
+    model.to(device)
+    checkpointer.load(ckpt, use_latest=ckpt is None)
+    do_evaluation(cfg, model, distributed)
+
+
+def main():
+    parser = argparse.ArgumentParser(description='SSD Evaluation on VOC and COCO dataset.')
+    parser.add_argument(
+        "--config-file",
+        default="configs/vgg_ssd300_voc0712.yaml",
+        metavar="FILE",
+        help="path to config file",
+        type=str,
+    )
+    parser.add_argument("--local_rank", type=int, default=0)
+    parser.add_argument(
+        "--ckpt",
+        help="The path to the checkpoint for test, default is the latest checkpoint.",
+        default=None,
+        type=str,
+    )
+
+    parser.add_argument("--output_dir", default="eval_results", type=str, help="The directory to store evaluation results.")
+
+    parser.add_argument(
+        "opts",
+        help="Modify config options using the command-line",
+        default=None,
+        nargs=argparse.REMAINDER,
+    )
+    args = parser.parse_args()
+
+    num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
+    distributed = num_gpus > 1
+
+    if torch.cuda.is_available():
+        # This flag allows you to enable the inbuilt cudnn auto-tuner to
+        # find the best algorithm to use for your hardware.
+        torch.backends.cudnn.benchmark = True
+    if distributed:
+        torch.cuda.set_device(args.local_rank)
+        torch.distributed.init_process_group(backend="nccl", init_method="env://")
+        synchronize()
+
+    cfg.merge_from_file(args.config_file)
+    cfg.merge_from_list(args.opts)
+    cfg.freeze()
+
+    logger = setup_logger("SSD", dist_util.get_rank(), cfg.OUTPUT_DIR)
+    logger.info("Using {} GPUs".format(num_gpus))
+    logger.info(args)
+
+    logger.info("Loaded configuration file {}".format(args.config_file))
+    with open(args.config_file, "r") as cf:
+        config_str = "\n" + cf.read()
+        logger.info(config_str)
+    logger.info("Running with config:\n{}".format(cfg))
+    evaluation_byHatcherW(cfg, ckpt=args.ckpt, distributed=distributed)
+
+
+if __name__ == '__main__':
+    main()