commit 54ccfb11feb3ec769f4b16a33940845d100c7222 Author: Semi212 <2729735763@qq.com> Date: Wed Sep 15 11:33:58 2021 +0800 first commit diff --git a/.idea/.gitignore b/.idea/.gitignore new file mode 100755 index 0000000..26d3352 --- /dev/null +++ b/.idea/.gitignore @@ -0,0 +1,3 @@ +# Default ignored files +/shelf/ +/workspace.xml diff --git a/.idea/MPKFNet.iml b/.idea/MPKFNet.iml new file mode 100755 index 0000000..b557b9a --- /dev/null +++ b/.idea/MPKFNet.iml @@ -0,0 +1,12 @@ + + + + + + + + + + \ No newline at end of file diff --git a/.idea/inspectionProfiles/Project_Default.xml b/.idea/inspectionProfiles/Project_Default.xml new file mode 100755 index 0000000..717a144 --- /dev/null +++ b/.idea/inspectionProfiles/Project_Default.xml @@ -0,0 +1,14 @@ + + + + \ No newline at end of file diff --git a/.idea/inspectionProfiles/profiles_settings.xml b/.idea/inspectionProfiles/profiles_settings.xml new file mode 100755 index 0000000..105ce2d --- /dev/null +++ b/.idea/inspectionProfiles/profiles_settings.xml @@ -0,0 +1,6 @@ + + + + \ No newline at end of file diff --git a/.idea/misc.xml b/.idea/misc.xml new file mode 100755 index 0000000..acece97 --- /dev/null +++ b/.idea/misc.xml @@ -0,0 +1,4 @@ + + + + \ No newline at end of file diff --git a/.idea/modules.xml b/.idea/modules.xml new file mode 100755 index 0000000..e0bce84 --- /dev/null +++ b/.idea/modules.xml @@ -0,0 +1,8 @@ + + + + + + + + \ No newline at end of file diff --git a/.idea/vcs.xml b/.idea/vcs.xml new file mode 100644 index 0000000..94a25f7 --- /dev/null +++ b/.idea/vcs.xml @@ -0,0 +1,6 @@ + + + + + + \ No newline at end of file diff --git a/Check_VOC_Hatcher.py b/Check_VOC_Hatcher.py new file mode 100755 index 0000000..1e37273 --- /dev/null +++ b/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)) \ No newline at end of file diff --git a/Dcl/DclConfig.py b/Dcl/DclConfig.py new file mode 100755 index 0000000..9d8ea2c --- /dev/null +++ b/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) + + + + diff --git a/Dcl/config.py b/Dcl/config.py new file mode 100755 index 0000000..129c0e2 --- /dev/null +++ b/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) + + + + diff --git a/Dcl/models/LoadModel.py b/Dcl/models/LoadModel.py new file mode 100755 index 0000000..a17d582 --- /dev/null +++ b/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 diff --git a/Dcl/transforms/__init__.py b/Dcl/transforms/__init__.py new file mode 100755 index 0000000..7986cdd --- /dev/null +++ b/Dcl/transforms/__init__.py @@ -0,0 +1 @@ +from .transforms import * diff --git a/Dcl/transforms/functional.py b/Dcl/transforms/functional.py new file mode 100644 index 0000000..5b16cbf --- /dev/null +++ b/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 diff --git a/Dcl/transforms/transforms.py b/Dcl/transforms/transforms.py new file mode 100755 index 0000000..5fe6d64 --- /dev/null +++ b/Dcl/transforms/transforms.py @@ -0,0 +1,1031 @@ +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 + +from . import functional as F + +__all__ = ["Compose", "ToTensor", "ToPILImage", "Normalize", "Resize", "Scale", "CenterCrop", "Pad", + "Lambda", "RandomApply", "RandomChoice", "RandomOrder", "RandomCrop", "RandomHorizontalFlip", + "RandomVerticalFlip", "RandomResizedCrop", "RandomSizedCrop", "FiveCrop", "TenCrop", "LinearTransformation", + "ColorJitter", "RandomRotation", "RandomAffine", "Grayscale", "RandomGrayscale", "Randomswap"] + +_pil_interpolation_to_str = { + Image.NEAREST: 'PIL.Image.NEAREST', + Image.BILINEAR: 'PIL.Image.BILINEAR', + Image.BICUBIC: 'PIL.Image.BICUBIC', + Image.LANCZOS: 'PIL.Image.LANCZOS', +} + + +class Compose(object): + """Composes several transforms together. + + Args: + transforms (list of ``Transform`` objects): list of transforms to compose. + + Example: + >>> transforms.Compose([ + >>> transforms.CenterCrop(10), + >>> transforms.ToTensor(), + >>> ]) + """ + + def __init__(self, transforms): + self.transforms = transforms + + def __call__(self, img): + for t in self.transforms: + img = t(img) + return img + + def __repr__(self): + format_string = self.__class__.__name__ + '(' + for t in self.transforms: + format_string += '\n' + format_string += ' {0}'.format(t) + format_string += '\n)' + return format_string + + +class ToTensor(object): + """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor. + + Converts a PIL Image or numpy.ndarray (H x W x C) in the range + [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0]. + """ + + def __call__(self, pic): + """ + Args: + pic (PIL Image or numpy.ndarray): Image to be converted to tensor. + + Returns: + Tensor: Converted image. + """ + return F.to_tensor(pic) + + def __repr__(self): + return self.__class__.__name__ + '()' + + +class ToPILImage(object): + """Convert a tensor or an ndarray to PIL Image. + + Converts a torch.*Tensor of shape C x H x W or a numpy ndarray of shape + H x W x C to a PIL Image while preserving the value range. + + Args: + mode (`PIL.Image mode`_): color space and pixel depth of input data (optional). + If ``mode`` is ``None`` (default) there are some assumptions made about the input data: + 1. If the input has 3 channels, the ``mode`` is assumed to be ``RGB``. + 2. If the input has 4 channels, the ``mode`` is assumed to be ``RGBA``. + 3. If the input has 1 channel, the ``mode`` is determined by the data type (i,e, + ``int``, ``float``, ``short``). + + .. _PIL.Image mode: http://pillow.readthedocs.io/en/3.4.x/handbook/concepts.html#modes + """ + def __init__(self, mode=None): + self.mode = mode + + def __call__(self, pic): + """ + Args: + pic (Tensor or numpy.ndarray): Image to be converted to PIL Image. + + Returns: + PIL Image: Image converted to PIL Image. + + """ + return F.to_pil_image(pic, self.mode) + + def __repr__(self): + format_string = self.__class__.__name__ + '(' + if self.mode is not None: + format_string += 'mode={0}'.format(self.mode) + format_string += ')' + return format_string + + +class Normalize(object): + """Normalize a tensor image with mean and standard deviation. + Given mean: ``(M1,...,Mn)`` and std: ``(S1,..,Sn)`` for ``n`` channels, this transform + will normalize each channel of the input ``torch.*Tensor`` i.e. + ``input[channel] = (input[channel] - mean[channel]) / std[channel]`` + + Args: + mean (sequence): Sequence of means for each channel. + std (sequence): Sequence of standard deviations for each channel. + """ + + def __init__(self, mean, std): + self.mean = mean + self.std = std + + def __call__(self, tensor): + """ + Args: + tensor (Tensor): Tensor image of size (C, H, W) to be normalized. + + Returns: + Tensor: Normalized Tensor image. + """ + return F.normalize(tensor, self.mean, self.std) + + def __repr__(self): + return self.__class__.__name__ + '(mean={0}, std={1})'.format(self.mean, self.std) + + +class Randomswap(object): + def __init__(self, size): + self.size = size + if isinstance(size, numbers.Number): + self.size = (int(size), int(size)) + else: + assert len(size) == 2, "Please provide only two dimensions (h, w) for size." + self.size = size + + def __call__(self, img): + return F.swap(img, self.size) + + def __repr__(self): + return self.__class__.__name__ + '(size={0})'.format(self.size) + + +class Resize(object): + """Resize the input PIL Image to the given size. + + Args: + size (sequence or int): Desired output size. If size is a sequence like + (h, w), output size will be matched to this. If size is an int, + smaller edge of the image will be matched to this number. + 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`` + """ + + def __init__(self, size, interpolation=Image.BILINEAR): + assert isinstance(size, int) or (isinstance(size, collections.Iterable) and len(size) == 2) + self.size = size + self.interpolation = interpolation + + def __call__(self, img): + """ + Args: + img (PIL Image): Image to be scaled. + + Returns: + PIL Image: Rescaled image. + """ + return F.resize(img, self.size, self.interpolation) + + def __repr__(self): + interpolate_str = _pil_interpolation_to_str[self.interpolation] + return self.__class__.__name__ + '(size={0}, interpolation={1})'.format(self.size, interpolate_str) + + +class Scale(Resize): + """ + Note: This transform is deprecated in favor of Resize. + """ + def __init__(self, *args, **kwargs): + warnings.warn("The use of the transforms.Scale transform is deprecated, " + + "please use transforms.Resize instead.") + super(Scale, self).__init__(*args, **kwargs) + + +class CenterCrop(object): + """Crops the given PIL Image at the center. + + 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. + """ + + def __init__(self, size): + if isinstance(size, numbers.Number): + self.size = (int(size), int(size)) + else: + self.size = size + + def __call__(self, img): + """ + Args: + img (PIL Image): Image to be cropped. + + Returns: + PIL Image: Cropped image. + """ + return F.center_crop(img, self.size) + + def __repr__(self): + return self.__class__.__name__ + '(size={0})'.format(self.size) + + +class Pad(object): + """Pad the given PIL Image on all sides with the given "pad" value. + + Args: + 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 at 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] + """ + + def __init__(self, padding, fill=0, padding_mode='constant'): + assert isinstance(padding, (numbers.Number, tuple)) + assert isinstance(fill, (numbers.Number, str, tuple)) + assert padding_mode in ['constant', 'edge', 'reflect', 'symmetric'] + 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))) + + self.padding = padding + self.fill = fill + self.padding_mode = padding_mode + + def __call__(self, img): + """ + Args: + img (PIL Image): Image to be padded. + + Returns: + PIL Image: Padded image. + """ + return F.pad(img, self.padding, self.fill, self.padding_mode) + + def __repr__(self): + return self.__class__.__name__ + '(padding={0}, fill={1}, padding_mode={2})'.\ + format(self.padding, self.fill, self.padding_mode) + + +class Lambda(object): + """Apply a user-defined lambda as a transform. + + Args: + lambd (function): Lambda/function to be used for transform. + """ + + def __init__(self, lambd): + assert isinstance(lambd, types.LambdaType) + self.lambd = lambd + + def __call__(self, img): + return self.lambd(img) + + def __repr__(self): + return self.__class__.__name__ + '()' + + +class RandomTransforms(object): + """Base class for a list of transformations with randomness + + Args: + transforms (list or tuple): list of transformations + """ + + def __init__(self, transforms): + assert isinstance(transforms, (list, tuple)) + self.transforms = transforms + + def __call__(self, *args, **kwargs): + raise NotImplementedError() + + def __repr__(self): + format_string = self.__class__.__name__ + '(' + for t in self.transforms: + format_string += '\n' + format_string += ' {0}'.format(t) + format_string += '\n)' + return format_string + + +class RandomApply(RandomTransforms): + """Apply randomly a list of transformations with a given probability + + Args: + transforms (list or tuple): list of transformations + p (float): probability + """ + + def __init__(self, transforms, p=0.5): + super(RandomApply, self).__init__(transforms) + self.p = p + + def __call__(self, img): + if self.p < random.random(): + return img + for t in self.transforms: + img = t(img) + return img + + def __repr__(self): + format_string = self.__class__.__name__ + '(' + format_string += '\n p={}'.format(self.p) + for t in self.transforms: + format_string += '\n' + format_string += ' {0}'.format(t) + format_string += '\n)' + return format_string + + +class RandomOrder(RandomTransforms): + """Apply a list of transformations in a random order + """ + def __call__(self, img): + order = list(range(len(self.transforms))) + random.shuffle(order) + for i in order: + img = self.transforms[i](img) + return img + + +class RandomChoice(RandomTransforms): + """Apply single transformation randomly picked from a list + """ + def __call__(self, img): + t = random.choice(self.transforms) + return t(img) + + +class RandomCrop(object): + """Crop the given PIL Image at a random location. + + 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. + padding (int or sequence, optional): Optional padding on each border + of the image. Default is 0, i.e no padding. If a sequence of length + 4 is provided, it is used to pad left, top, right, bottom borders + respectively. + pad_if_needed (boolean): It will pad the image if smaller than the + desired size to avoid raising an exception. + """ + + def __init__(self, size, padding=0, pad_if_needed=False): + if isinstance(size, numbers.Number): + self.size = (int(size), int(size)) + else: + self.size = size + self.padding = padding + self.pad_if_needed = pad_if_needed + + @staticmethod + def get_params(img, output_size): + """Get parameters for ``crop`` for a random crop. + + Args: + img (PIL Image): Image to be cropped. + output_size (tuple): Expected output size of the crop. + + Returns: + tuple: params (i, j, h, w) to be passed to ``crop`` for random crop. + """ + w, h = img.size + th, tw = output_size + if w == tw and h == th: + return 0, 0, h, w + + i = random.randint(0, h - th) + j = random.randint(0, w - tw) + return i, j, th, tw + + def __call__(self, img): + """ + Args: + img (PIL Image): Image to be cropped. + + Returns: + PIL Image: Cropped image. + """ + if self.padding > 0: + img = F.pad(img, self.padding) + + # pad the width if needed + if self.pad_if_needed and img.size[0] < self.size[1]: + img = F.pad(img, (int((1 + self.size[1] - img.size[0]) / 2), 0)) + # pad the height if needed + if self.pad_if_needed and img.size[1] < self.size[0]: + img = F.pad(img, (0, int((1 + self.size[0] - img.size[1]) / 2))) + + i, j, h, w = self.get_params(img, self.size) + + return F.crop(img, i, j, h, w) + + def __repr__(self): + return self.__class__.__name__ + '(size={0}, padding={1})'.format(self.size, self.padding) + + +class RandomHorizontalFlip(object): + """Horizontally flip the given PIL Image randomly with a given probability. + + Args: + p (float): probability of the image being flipped. Default value is 0.5 + """ + + def __init__(self, p=0.5): + self.p = p + + def __call__(self, img): + """ + Args: + img (PIL Image): Image to be flipped. + + Returns: + PIL Image: Randomly flipped image. + """ + if random.random() < self.p: + return F.hflip(img) + return img + + def __repr__(self): + return self.__class__.__name__ + '(p={})'.format(self.p) + + +class RandomVerticalFlip(object): + """Vertically flip the given PIL Image randomly with a given probability. + + Args: + p (float): probability of the image being flipped. Default value is 0.5 + """ + + def __init__(self, p=0.5): + self.p = p + + def __call__(self, img): + """ + Args: + img (PIL Image): Image to be flipped. + + Returns: + PIL Image: Randomly flipped image. + """ + if random.random() < self.p: + return F.vflip(img) + return img + + def __repr__(self): + return self.__class__.__name__ + '(p={})'.format(self.p) + + +class RandomResizedCrop(object): + """Crop the given PIL Image to random size and aspect ratio. + + A crop of random size (default: of 0.08 to 1.0) of the original size and a random + aspect ratio (default: of 3/4 to 4/3) of the original aspect ratio is made. This crop + is finally resized to given size. + This is popularly used to train the Inception networks. + + Args: + size: expected output size of each edge + scale: range of size of the origin size cropped + ratio: range of aspect ratio of the origin aspect ratio cropped + interpolation: Default: PIL.Image.BILINEAR + """ + + def __init__(self, size, scale=(0.08, 1.0), ratio=(3. / 4., 4. / 3.), interpolation=Image.BILINEAR): + self.size = (size, size) + self.interpolation = interpolation + self.scale = scale + self.ratio = ratio + + @staticmethod + def get_params(img, scale, ratio): + """Get parameters for ``crop`` for a random sized crop. + + Args: + img (PIL Image): Image to be cropped. + scale (tuple): range of size of the origin size cropped + ratio (tuple): range of aspect ratio of the origin aspect ratio cropped + + Returns: + tuple: params (i, j, h, w) to be passed to ``crop`` for a random + sized crop. + """ + for attempt in range(10): + area = img.size[0] * img.size[1] + target_area = random.uniform(*scale) * area + aspect_ratio = random.uniform(*ratio) + + w = int(round(math.sqrt(target_area * aspect_ratio))) + h = int(round(math.sqrt(target_area / aspect_ratio))) + + if random.random() < 0.5: + w, h = h, w + + if w <= img.size[0] and h <= img.size[1]: + i = random.randint(0, img.size[1] - h) + j = random.randint(0, img.size[0] - w) + return i, j, h, w + + # Fallback + w = min(img.size[0], img.size[1]) + i = (img.size[1] - w) // 2 + j = (img.size[0] - w) // 2 + return i, j, w, w + + def __call__(self, img): + """ + Args: + img (PIL Image): Image to be cropped and resized. + + Returns: + PIL Image: Randomly cropped and resized image. + """ + i, j, h, w = self.get_params(img, self.scale, self.ratio) + return F.resized_crop(img, i, j, h, w, self.size, self.interpolation) + + def __repr__(self): + interpolate_str = _pil_interpolation_to_str[self.interpolation] + format_string = self.__class__.__name__ + '(size={0}'.format(self.size) + format_string += ', scale={0}'.format(tuple(round(s, 4) for s in self.scale)) + format_string += ', ratio={0}'.format(tuple(round(r, 4) for r in self.ratio)) + format_string += ', interpolation={0})'.format(interpolate_str) + return format_string + + +class RandomSizedCrop(RandomResizedCrop): + """ + Note: This transform is deprecated in favor of RandomResizedCrop. + """ + def __init__(self, *args, **kwargs): + warnings.warn("The use of the transforms.RandomSizedCrop transform is deprecated, " + + "please use transforms.RandomResizedCrop instead.") + super(RandomSizedCrop, self).__init__(*args, **kwargs) + + +class FiveCrop(object): + """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. See below for an example of how to deal with + this. + + 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 of size (size, size) is made. + + Example: + >>> transform = Compose([ + >>> FiveCrop(size), # this is a list of PIL Images + >>> Lambda(lambda crops: torch.stack([ToTensor()(crop) for crop in crops])) # returns a 4D tensor + >>> ]) + >>> #In your test loop you can do the following: + >>> input, target = batch # input is a 5d tensor, target is 2d + >>> bs, ncrops, c, h, w = input.size() + >>> result = model(input.view(-1, c, h, w)) # fuse batch size and ncrops + >>> result_avg = result.view(bs, ncrops, -1).mean(1) # avg over crops + """ + + def __init__(self, size): + self.size = size + if isinstance(size, numbers.Number): + self.size = (int(size), int(size)) + else: + assert len(size) == 2, "Please provide only two dimensions (h, w) for size." + self.size = size + + def __call__(self, img): + return F.five_crop(img, self.size) + + def __repr__(self): + return self.__class__.__name__ + '(size={0})'.format(self.size) + + +class TenCrop(object): + """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. See below for an example of how to deal with + this. + + 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 + + Example: + >>> transform = Compose([ + >>> TenCrop(size), # this is a list of PIL Images + >>> Lambda(lambda crops: torch.stack([ToTensor()(crop) for crop in crops])) # returns a 4D tensor + >>> ]) + >>> #In your test loop you can do the following: + >>> input, target = batch # input is a 5d tensor, target is 2d + >>> bs, ncrops, c, h, w = input.size() + >>> result = model(input.view(-1, c, h, w)) # fuse batch size and ncrops + >>> result_avg = result.view(bs, ncrops, -1).mean(1) # avg over crops + """ + + def __init__(self, size, vertical_flip=False): + self.size = size + if isinstance(size, numbers.Number): + self.size = (int(size), int(size)) + else: + assert len(size) == 2, "Please provide only two dimensions (h, w) for size." + self.size = size + self.vertical_flip = vertical_flip + + def __call__(self, img): + return F.ten_crop(img, self.size, self.vertical_flip) + + def __repr__(self): + return self.__class__.__name__ + '(size={0}, vertical_flip={1})'.format(self.size, self.vertical_flip) + + +class LinearTransformation(object): + """Transform a tensor image with a square transformation matrix computed + offline. + + Given transformation_matrix, will flatten the torch.*Tensor, compute the dot + product with the transformation matrix and reshape the tensor to its + original shape. + + Applications: + - whitening: zero-center the data, compute the data covariance matrix + [D x D] with np.dot(X.T, X), perform SVD on this matrix and + pass it as transformation_matrix. + + Args: + transformation_matrix (Tensor): tensor [D x D], D = C x H x W + """ + + def __init__(self, transformation_matrix): + if transformation_matrix.size(0) != transformation_matrix.size(1): + raise ValueError("transformation_matrix should be square. Got " + + "[{} x {}] rectangular matrix.".format(*transformation_matrix.size())) + self.transformation_matrix = transformation_matrix + + def __call__(self, tensor): + """ + Args: + tensor (Tensor): Tensor image of size (C, H, W) to be whitened. + + Returns: + Tensor: Transformed image. + """ + if tensor.size(0) * tensor.size(1) * tensor.size(2) != self.transformation_matrix.size(0): + raise ValueError("tensor and transformation matrix have incompatible shape." + + "[{} x {} x {}] != ".format(*tensor.size()) + + "{}".format(self.transformation_matrix.size(0))) + flat_tensor = tensor.view(1, -1) + transformed_tensor = torch.mm(flat_tensor, self.transformation_matrix) + tensor = transformed_tensor.view(tensor.size()) + return tensor + + def __repr__(self): + format_string = self.__class__.__name__ + '(' + format_string += (str(self.transformation_matrix.numpy().tolist()) + ')') + return format_string + + +class ColorJitter(object): + """Randomly change the brightness, contrast and saturation of an image. + + Args: + brightness (float): How much to jitter brightness. brightness_factor + is chosen uniformly from [max(0, 1 - brightness), 1 + brightness]. + contrast (float): How much to jitter contrast. contrast_factor + is chosen uniformly from [max(0, 1 - contrast), 1 + contrast]. + saturation (float): How much to jitter saturation. saturation_factor + is chosen uniformly from [max(0, 1 - saturation), 1 + saturation]. + hue(float): How much to jitter hue. hue_factor is chosen uniformly from + [-hue, hue]. Should be >=0 and <= 0.5. + """ + def __init__(self, brightness=0, contrast=0, saturation=0, hue=0): + self.brightness = brightness + self.contrast = contrast + self.saturation = saturation + self.hue = hue + + @staticmethod + def get_params(brightness, contrast, saturation, hue): + """Get a randomized transform to be applied on image. + + Arguments are same as that of __init__. + + Returns: + Transform which randomly adjusts brightness, contrast and + saturation in a random order. + """ + transforms = [] + if brightness > 0: + brightness_factor = random.uniform(max(0, 1 - brightness), 1 + brightness) + transforms.append(Lambda(lambda img: F.adjust_brightness(img, brightness_factor))) + + if contrast > 0: + contrast_factor = random.uniform(max(0, 1 - contrast), 1 + contrast) + transforms.append(Lambda(lambda img: F.adjust_contrast(img, contrast_factor))) + + if saturation > 0: + saturation_factor = random.uniform(max(0, 1 - saturation), 1 + saturation) + transforms.append(Lambda(lambda img: F.adjust_saturation(img, saturation_factor))) + + if hue > 0: + hue_factor = random.uniform(-hue, hue) + transforms.append(Lambda(lambda img: F.adjust_hue(img, hue_factor))) + + random.shuffle(transforms) + transform = Compose(transforms) + + return transform + + def __call__(self, img): + """ + Args: + img (PIL Image): Input image. + + Returns: + PIL Image: Color jittered image. + """ + transform = self.get_params(self.brightness, self.contrast, + self.saturation, self.hue) + return transform(img) + + def __repr__(self): + format_string = self.__class__.__name__ + '(' + format_string += 'brightness={0}'.format(self.brightness) + format_string += ', contrast={0}'.format(self.contrast) + format_string += ', saturation={0}'.format(self.saturation) + format_string += ', hue={0})'.format(self.hue) + return format_string + + +class RandomRotation(object): + """Rotate the image by angle. + + Args: + degrees (sequence or float or int): Range of degrees to select from. + If degrees is a number instead of sequence like (min, max), the range of degrees + will be (-degrees, +degrees). + 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 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. + """ + + def __init__(self, degrees, resample=False, expand=False, center=None): + if isinstance(degrees, numbers.Number): + if degrees < 0: + raise ValueError("If degrees is a single number, it must be positive.") + self.degrees = (-degrees, degrees) + else: + if len(degrees) != 2: + raise ValueError("If degrees is a sequence, it must be of len 2.") + self.degrees = degrees + + self.resample = resample + self.expand = expand + self.center = center + + @staticmethod + def get_params(degrees): + """Get parameters for ``rotate`` for a random rotation. + + Returns: + sequence: params to be passed to ``rotate`` for random rotation. + """ + angle = random.uniform(degrees[0], degrees[1]) + + return angle + + def __call__(self, img): + """ + img (PIL Image): Image to be rotated. + + Returns: + PIL Image: Rotated image. + """ + + angle = self.get_params(self.degrees) + + return F.rotate(img, angle, self.resample, self.expand, self.center) + + def __repr__(self): + format_string = self.__class__.__name__ + '(degrees={0}'.format(self.degrees) + format_string += ', resample={0}'.format(self.resample) + format_string += ', expand={0}'.format(self.expand) + if self.center is not None: + format_string += ', center={0}'.format(self.center) + format_string += ')' + return format_string + + +class RandomAffine(object): + """Random affine transformation of the image keeping center invariant + + Args: + degrees (sequence or float or int): Range of degrees to select from. + If degrees is a number instead of sequence like (min, max), the range of degrees + will be (-degrees, +degrees). Set to 0 to desactivate rotations. + translate (tuple, optional): tuple of maximum absolute fraction for horizontal + and vertical translations. For example translate=(a, b), then horizontal shift + is randomly sampled in the range -img_width * a < dx < img_width * a and vertical shift is + randomly sampled in the range -img_height * b < dy < img_height * b. Will not translate by default. + scale (tuple, optional): scaling factor interval, e.g (a, b), then scale is + randomly sampled from the range a <= scale <= b. Will keep original scale by default. + shear (sequence or float or int, optional): Range of degrees to select from. + If degrees is a number instead of sequence like (min, max), the range of degrees + will be (-degrees, +degrees). Will not apply shear by default + 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) + """ + + def __init__(self, degrees, translate=None, scale=None, shear=None, resample=False, fillcolor=0): + if isinstance(degrees, numbers.Number): + if degrees < 0: + raise ValueError("If degrees is a single number, it must be positive.") + self.degrees = (-degrees, degrees) + else: + assert isinstance(degrees, (tuple, list)) and len(degrees) == 2, \ + "degrees should be a list or tuple and it must be of length 2." + self.degrees = degrees + + if translate is not None: + assert isinstance(translate, (tuple, list)) and len(translate) == 2, \ + "translate should be a list or tuple and it must be of length 2." + for t in translate: + if not (0.0 <= t <= 1.0): + raise ValueError("translation values should be between 0 and 1") + self.translate = translate + + if scale is not None: + assert isinstance(scale, (tuple, list)) and len(scale) == 2, \ + "scale should be a list or tuple and it must be of length 2." + for s in scale: + if s <= 0: + raise ValueError("scale values should be positive") + self.scale = scale + + if shear is not None: + if isinstance(shear, numbers.Number): + if shear < 0: + raise ValueError("If shear is a single number, it must be positive.") + self.shear = (-shear, shear) + else: + assert isinstance(shear, (tuple, list)) and len(shear) == 2, \ + "shear should be a list or tuple and it must be of length 2." + self.shear = shear + else: + self.shear = shear + + self.resample = resample + self.fillcolor = fillcolor + + @staticmethod + def get_params(degrees, translate, scale_ranges, shears, img_size): + """Get parameters for affine transformation + + Returns: + sequence: params to be passed to the affine transformation + """ + angle = random.uniform(degrees[0], degrees[1]) + if translate is not None: + max_dx = translate[0] * img_size[0] + max_dy = translate[1] * img_size[1] + translations = (np.round(random.uniform(-max_dx, max_dx)), + np.round(random.uniform(-max_dy, max_dy))) + else: + translations = (0, 0) + + if scale_ranges is not None: + scale = random.uniform(scale_ranges[0], scale_ranges[1]) + else: + scale = 1.0 + + if shears is not None: + shear = random.uniform(shears[0], shears[1]) + else: + shear = 0.0 + + return angle, translations, scale, shear + + def __call__(self, img): + """ + img (PIL Image): Image to be transformed. + + Returns: + PIL Image: Affine transformed image. + """ + ret = self.get_params(self.degrees, self.translate, self.scale, self.shear, img.size) + return F.affine(img, *ret, resample=self.resample, fillcolor=self.fillcolor) + + def __repr__(self): + s = '{name}(degrees={degrees}' + if self.translate is not None: + s += ', translate={translate}' + if self.scale is not None: + s += ', scale={scale}' + if self.shear is not None: + s += ', shear={shear}' + if self.resample > 0: + s += ', resample={resample}' + if self.fillcolor != 0: + s += ', fillcolor={fillcolor}' + s += ')' + d = dict(self.__dict__) + d['resample'] = _pil_interpolation_to_str[d['resample']] + return s.format(name=self.__class__.__name__, **d) + + +class Grayscale(object): + """Convert image to grayscale. + + Args: + num_output_channels (int): (1 or 3) number of channels desired for output image + + Returns: + PIL Image: Grayscale version of the input. + - If num_output_channels == 1 : returned image is single channel + - If num_output_channels == 3 : returned image is 3 channel with r == g == b + + """ + + def __init__(self, num_output_channels=1): + self.num_output_channels = num_output_channels + + def __call__(self, img): + """ + Args: + img (PIL Image): Image to be converted to grayscale. + + Returns: + PIL Image: Randomly grayscaled image. + """ + return F.to_grayscale(img, num_output_channels=self.num_output_channels) + + def __repr__(self): + return self.__class__.__name__ + '(num_output_channels={0})'.format(self.num_output_channels) + + +class RandomGrayscale(object): + """Randomly convert image to grayscale with a probability of p (default 0.1). + + Args: + p (float): probability that image should be converted to grayscale. + + Returns: + PIL Image: Grayscale version of the input image with probability p and unchanged + with probability (1-p). + - If input image is 1 channel: grayscale version is 1 channel + - If input image is 3 channel: grayscale version is 3 channel with r == g == b + + """ + + def __init__(self, p=0.1): + self.p = p + + def __call__(self, img): + """ + Args: + img (PIL Image): Image to be converted to grayscale. + + Returns: + PIL Image: Randomly grayscaled image. + """ + num_output_channels = 1 if img.mode == 'L' else 3 + if random.random() < self.p: + return F.to_grayscale(img, num_output_channels=num_output_channels) + return img + + def __repr__(self): + return self.__class__.__name__ + '(p={0})'.format(self.p) diff --git a/Dcl/util/Asoftmax_loss.py b/Dcl/util/Asoftmax_loss.py new file mode 100755 index 0000000..44327eb --- /dev/null +++ b/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 + diff --git a/Dcl/util/__init__.py b/Dcl/util/__init__.py new file mode 100755 index 0000000..e69de29 diff --git a/Dcl/util/autoaugment.py b/Dcl/util/autoaugment.py new file mode 100755 index 0000000..9f39b51 --- /dev/null +++ b/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 diff --git a/Dcl/util/dataset_DCL.py b/Dcl/util/dataset_DCL.py new file mode 100755 index 0000000..2e9c81f --- /dev/null +++ b/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 diff --git a/Dcl/util/eval_model.py b/Dcl/util/eval_model.py new file mode 100755 index 0000000..7a68ce2 --- /dev/null +++ b/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 + diff --git a/Dcl/util/test_tool.py b/Dcl/util/test_tool.py new file mode 100755 index 0000000..3c409e0 --- /dev/null +++ b/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 + + + + + + + diff --git a/Dcl/util/train_model.py b/Dcl/util/train_model.py new file mode 100755 index 0000000..b5a2066 --- /dev/null +++ b/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() + + + diff --git a/Dcl/util/utils.py b/Dcl/util/utils.py new file mode 100755 index 0000000..7552510 --- /dev/null +++ b/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) diff --git a/DclTest.py b/DclTest.py new file mode 100755 index 0000000..e40efe8 --- /dev/null +++ b/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) diff --git a/DensityVideo_Video.py b/DensityVideo_Video.py new file mode 100755 index 0000000..653428f --- /dev/null +++ b/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) diff --git a/DetectionVideo_Video.py b/DetectionVideo_Video.py new file mode 100755 index 0000000..a318f8a --- /dev/null +++ b/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) diff --git a/GYD_RedCamera.py b/GYD_RedCamera.py new file mode 100755 index 0000000..d91a91c --- /dev/null +++ b/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) diff --git a/Hatcher_cam.py b/Hatcher_cam.py new file mode 100644 index 0000000..0c4abc2 --- /dev/null +++ b/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) + # diff --git a/Hatcher_is_cool.py b/Hatcher_is_cool.py new file mode 100755 index 0000000..1ba28c7 --- /dev/null +++ b/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) diff --git a/Hatcher_red.py b/Hatcher_red.py new file mode 100644 index 0000000..9ccc992 --- /dev/null +++ b/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()))) diff --git a/PlayVideo.py b/PlayVideo.py new file mode 100755 index 0000000..ab52178 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/System_I.py b/System_I.py new file mode 100755 index 0000000..14251ca --- /dev/null +++ b/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 + + + + + + + + diff --git a/Y5/data/argoverse_hd.yaml b/Y5/data/argoverse_hd.yaml new file mode 100644 index 0000000..df7a936 --- /dev/null +++ b/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' ] diff --git a/Y5/data/coco.yaml b/Y5/data/coco.yaml new file mode 100644 index 0000000..b9da2bf --- /dev/null +++ b/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) diff --git a/Y5/data/coco128.yaml b/Y5/data/coco128.yaml new file mode 100644 index 0000000..c41bccf --- /dev/null +++ b/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' ] diff --git a/Y5/data/hyp.finetune.yaml b/Y5/data/hyp.finetune.yaml new file mode 100644 index 0000000..1b84cff --- /dev/null +++ b/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 diff --git a/Y5/data/hyp.scratch.yaml b/Y5/data/hyp.scratch.yaml new file mode 100644 index 0000000..44f26b6 --- /dev/null +++ b/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) diff --git a/Y5/data/images/bus.jpg b/Y5/data/images/bus.jpg new file mode 100644 index 0000000..3415f25 Binary files /dev/null and b/Y5/data/images/bus.jpg differ diff --git a/Y5/data/images/zidane.jpg b/Y5/data/images/zidane.jpg new file mode 100644 index 0000000..92d72ea Binary files /dev/null and b/Y5/data/images/zidane.jpg differ diff --git a/Y5/data/scripts/get_argoverse_hd.sh b/Y5/data/scripts/get_argoverse_hd.sh new file mode 100644 index 0000000..9e0db9f --- /dev/null +++ b/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 - "$@" <train.txt +cat 2007_train.txt 2007_val.txt 2007_test.txt 2012_train.txt 2012_val.txt >train.all.txt + +python3 - "$@" <