Update cmc_reid.py

.gitignore

@@ -5,6 +5,4 @@
 *.mp4
 *.t7
 *.pth
-*label*
-result.txt
-rank10_detection*.txt
+label*

cmc_reid.py

@@ -0,0 +1,112 @@
+import shutil
+import matplotlib.pyplot as plt
+import torch
+from torchvision import transforms
+import numpy as np
+import cv2
+import os
+import param
+import param_v1
+from fastreid.my_extractor import MyExtractor
+
+
+def extractidfeature(id_path: str, extractor: MyExtractor):
+    img_id = cv2.imread(id_path)
+    tensor_id = extractor([img_id])
+    return tensor_id.detach().cpu().numpy()
+
+if __name__ == '__main__':
+    my_option = param_v1.Parameters()                  # load model in head
+    extractor = MyExtractor(my_option.weights_reid, use_cuda=True)
+    # gallery_index = 3
+    gallery_path = f"../ReIDMixLearning/mydataset/gallery/"
+    classes = os.listdir(gallery_path)
+    cmc_sum = []
+    cmc_final = []
+    for query_index in range(1, 6+1):                #query是从1开始的
+
+        query_path = f"../ReIDMixLearning/mydataset/query/cam{query_index}/"   # Query path
+                #Gallary path，这里直接对gallary文件夹下所有图片进行分析
+        print("--------------query cam is ", query_index)
+        for cls in classes:
+            if cls != f"{query_index}":
+                print("——————————gallery cam is", cls)
+                gallery_path = f"../ReIDMixLearning/mydataset/gallery/{cls}/"
+                query_feature_list = [] # 计算每个query的特征向量保存在列表中
+                acck_l_list = []        # 二维列表（列表的列表），其元素（也是列表）存储单个query的"acck"("cmc")
+                cmc_k_list = []         # 存储最终的CMC数字（所有query取平均）
+                query_list = os.listdir(query_path)     # 为每张图片生成路径并保存在列表变量中
+                query_list.sort()
+                gallery_list = os.listdir(gallery_path) # 为每张图片生成路径并保存在列表变量中
+                with torch.no_grad():
+                    # 对于每个query，（这句话加在本循环内每个注释）
+                    for path in query_list:
+                        distance_dict = {}
+                        parsed_query_id = path.split("_")[1]
+                        parsed_query_id = parsed_query_id.split(".")[0]
+                        parsed_query_id = int(parsed_query_id)
+                        if parsed_query_id == 4 and cls == "5":
+                            continue
+                        print(parsed_query_id)
+                        id_feature = extractidfeature(query_path+path, extractor)
+                        query_feature_list.append(id_feature)
+                        counter = 0
+                        # 计算每个gallary图片到目标query的特征距离
+                        for p in gallery_list:
+                                if counter%10 == 9:     # 抽样，相当于设置帧间隔
+                                    img_path = gallery_path+p
+                                    gallery_feature = extractidfeature(img_path, extractor)
+                                    distance_dict[p] = np.sum(np.abs(id_feature-gallery_feature))
+                                counter += 1
+                        matched = sorted(distance_dict.items(), key=lambda x : x[1])    # 按特征距离排序
+                        # 储存k从前1到前10名的是否（存在命中的ID）
+                        temp_acck = [int(matched[0][0].split("_")[2].split(".")[0])==parsed_query_id]     # ID相同为True，否则False
+                        # 获取前k个acck情况
+                        for i in range(1,11):
+                            parsed_gallery_id = matched[i][0].split("_")[2].split(".")[0]
+                            parsed_gallery_id = int(parsed_gallery_id)           # 解析出gallary的ID
+                            print(parsed_gallery_id)  # 解析出gallary的ID
+                            temp_acck.append(parsed_gallery_id==parsed_query_id or temp_acck[i-1])    # or用意是表明如果Acc(k-1)已经为True，那么Acck必然为True（递推）
+                        print(temp_acck)
+                        acck_l_list.append(temp_acck)
+                    # 计算不同k下cmc的具体数值
+                    for k in range(1, 11):
+                        acc_counter = 0.0
+                        length = len(acck_l_list)
+                        for results in acck_l_list:
+                            acc_counter += results[k]
+                        cmc_k_list.append(acc_counter/length)
+                    #打印以供检查
+                    print(cmc_k_list)
+                    cmc_sum.append(cmc_k_list)
+                    k_index = [k for k in range(1, 11)]
+                    #plot figure并保存（不会显示）
+                    plt.clf()
+                    plt.plot(k_index, cmc_k_list)
+                    plt.xlabel("k")
+                    plt.ylabel("ACCK")
+                    plt.ylim(0.0,1.05)
+                    plt.title(f"CMC:query_cam{query_index} and gallery {cls}")
+                    for x, y in zip(k_index, cmc_k_list):
+                        plt.text(x, y + 0.02, str(round(y, 3)), ha='center', va='bottom', fontsize=10.5)
+                    plt.draw()
+                    plt.savefig(f"./CMC/query_cam{query_index}_and_gallery_{cls}.jpg")
+    for k in range(10):
+        acc_counter = 0.0
+        length = len(cmc_sum)
+        for results in cmc_sum:
+            acc_counter += results[k]
+        cmc_final.append(acc_counter / length)
+    print(cmc_final)
+    k_index = [k for k in range(1, 11)]
+    # plot figure并保存（不会显示）
+    plt.clf()
+    plt.plot(k_index, cmc_final)
+    plt.xlabel("k")
+    plt.ylabel("ACCK")
+    plt.ylim(0.0, 1.05)
+    plt.title(f"CMC:query and gallery")
+    for x, y in zip(k_index, cmc_final):
+        plt.text(x, y + 0.02, str(round(y,3)), ha='center', va='bottom', fontsize=10.5)
+    plt.draw()
+    plt.savefig(f"./CMC/query_gallery.jpg")

cmc_reid_detector.py

@@ -0,0 +1,120 @@
+import shutil
+import matplotlib.pyplot as plt
+import torch
+from torchvision import transforms
+import numpy as np
+import cv2
+import os
+import param
+import param_v1
+from fastreid.my_extractor import MyExtractor
+
+
+def extractidfeature(id_path: str, extractor: MyExtractor):
+    img_id = cv2.imread(id_path)
+    tensor_id = extractor([img_id])
+    return tensor_id.detach().cpu().numpy()
+
+if __name__ == '__main__':
+    my_option = param_v1.Parameters()                  # load model in head
+    extractor = MyExtractor(my_option.weights_reid,use_cuda=True)    
+    # detector_index = 3
+    detector_path = f"../yolov7-inferonly/detector/"
+    classes = os.listdir(detector_path)
+    cmc_sum=[]
+    cmc_final=[]
+    for query_index in range(1, 6+1):                #query是从1开始的
+        data = []
+        query_path = f"../ReIDMixLearning/mydataset/query/cam{query_index}/"   # Query path
+               #Gallary path，这里直接对gallary文件夹下所有图片进行分析
+        for cls in classes:
+            if cls != f"cam{query_index}":
+                detector_path = f"../yolov7-inferonly/detector/{cls}/"
+                # images = os.listdir(detector_path + cls)
+                # data.append(images)
+
+                query_feature_list = [] # 计算每个query的特征向量保存在列表中
+                acck_l_list = []        # 二维列表（列表的列表），其元素（也是列表）存储单个query的"acck"("cmc")
+                cmc_k_list = []         # 存储最终的CMC数字（所有query取平均）
+                query_list = os.listdir(query_path)     # 为每张图片生成路径并保存在列表变量中
+                query_list.sort()
+                detector_list = os.listdir(detector_path) # 为每张图片生成路径并保存在列表变量中
+                # detector_list = data
+                with torch.no_grad():
+                    # 对于每个query，（这句话加在本循环内每个注释）
+                    for path in query_list:
+                        distance_dict = {}
+                        parsed_query_id = path.split("_")[1]
+                        parsed_query_id = parsed_query_id.split(".")[0]
+                        parsed_query_id = int(parsed_query_id)
+                        if parsed_query_id == 4 and cls == "cam5":
+                            continue
+                        print(parsed_query_id)
+                        id_feature = extractidfeature(query_path+path, extractor)
+                        query_feature_list.append(id_feature)
+                        # counter = 0
+                        # 计算每个gallary图片到目标query的特征距离
+                        for p in detector_list:
+                            # if counter%10 == 9:     # 抽样，相当于设置帧间隔
+                                img_path = detector_path+p
+                                detector_feature = extractidfeature(img_path, extractor)
+                                distance_dict[p] = np.sum(np.abs(id_feature-detector_feature))
+                                # counter += 1
+                        matched = sorted(distance_dict.items(), key=lambda x : x[1])    # 按特征距离排序
+                        # 储存k从前1到前10名的是否（存在命中的ID）
+                        temp_acck = [int(matched[0][0].split("_")[2].split(".")[0])==parsed_query_id]     # ID相同为True，否则False
+                        # 获取前k个acck情况
+                        for i in range(1,11):
+                            parsed_detector_id = matched[i][0].split("_")[2].split(".")[0]
+                            parsed_detector_id=int(parsed_detector_id)
+                            print(parsed_detector_id)# 解析出gallary的ID
+                            temp_acck.append(parsed_detector_id==parsed_query_id or temp_acck[i-1])    # or用意是表明如果Acc(k-1)已经为True，那么Acck必然为True（递推）
+                        print(temp_acck)
+                        acck_l_list.append(temp_acck)
+                    # 计算不同k下cmc的具体数值
+                    for k in range(1, 11):
+                        acc_counter = 0.0
+                        length = len(acck_l_list)
+                        for results in acck_l_list:
+                            acc_counter += results[k]
+                        cmc_k_list.append(acc_counter/length)
+
+                    #打印以供检查
+                    print(cmc_k_list)
+                    cmc_sum.append(cmc_k_list)
+
+                    k_index = [k for k in range(1, 11)]
+                    #plot figure并保存（不会显示）
+                    plt.clf()
+                    plt.plot(k_index, cmc_k_list)
+                    plt.xlabel("k")
+                    plt.ylabel("ACCK")
+                    plt.ylim(0.0,1.05)
+                    plt.title(f"CMC:query_cam{query_index} and detector {cls}")
+                    for x, y in zip(k_index, cmc_k_list):
+                        plt.text(x, y + 0.02, str(round(y,3)), ha='center', va='bottom', fontsize=10.5)
+                    plt.draw()
+                    plt.savefig(f"./CMC_detector/query_cam{query_index}_detector_{cls}.jpg")
+
+    for k in range(10):
+        acc_counter = 0.0
+        length = len(cmc_sum)
+        for results in cmc_sum:
+            acc_counter += results[k]
+        cmc_final.append(acc_counter / length)
+    print(cmc_final)
+    k_index = [k for k in range(1, 11)]
+    # plot figure并保存（不会显示）
+    plt.clf()
+    plt.plot(k_index, cmc_final)
+    plt.xlabel("k")
+    plt.ylabel("ACCK")
+    plt.ylim(0.0, 1.05)
+    plt.title(f"CMC:query and detector")
+    for x, y in zip(k_index, cmc_final):
+        plt.text(x, y + 0.02, str(round(y, 3)), ha='center', va='bottom', fontsize=10.5)
+    plt.draw()
+    plt.savefig(f"./CMC_detector/query_detector.jpg")
+
+
+

deep_sort/deep/feature_extractor.py

@@ -1,11 +1,13 @@
 import torch
-from torch import nn
 import torchvision.transforms as transforms
 import numpy as np
 import cv2
 import logging
 
-from .model import Net,MyNet,RestNet18
+from .model import Net
+from fastreid.config import get_cfg
+from fastreid.engine import DefaultTrainer
+from fastreid.utils.checkpoint import Checkpointer
 
 class Extractor(object):
     def __init__(self, model_path, use_cuda=True):
@@ -47,40 +49,36 @@ class Extractor(object):
             features = self.net(im_batch)
         return features.cpu().numpy()
 
-    
-class MyExtractor(object):
-    def __init__(self, model_path, use_cuda=True):
+class FastReIDExtractor(object):
+    def __init__(self, model_config, model_path, use_cuda=True):
+        cfg = get_cfg()
+        cfg.merge_from_file(model_config)
+        cfg.MODEL.BACKBONE.PRETRAIN = False
+        cfg.MODEL.DEVICE = "cuda" if use_cuda else "cpu"
+        self.net = DefaultTrainer.build_model(cfg)
         self.device = "cuda" if torch.cuda.is_available() and use_cuda else "cpu"
-        self.net = RestNet18()
-        self.net.load_state_dict(torch.load(model_path,map_location=torch.device(self.device)))
-        if self.device=="cuda":
-            self.net = self.net.to(self.device)
+
+        Checkpointer(self.net).load(model_path)
         logger = logging.getLogger("root.tracker")
         logger.info("Loading weights from {}... Done!".format(model_path))
-        self.raw_transformer = transforms.Compose([
-        transforms.ToPILImage(),
-        transforms.Resize((128,64)),# hxw
-        transforms.ToTensor(),
-        ])   
+        self.net.to(self.device)
+        self.net.eval()
+        height, width = cfg.INPUT.SIZE_TEST
+        self.size = (width, height)
         self.norm = transforms.Compose([
             transforms.ToTensor(),
             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
         ])
-        
-
 
+    
     def _preprocess(self, im_crops):
-        """
-        TODO:
-            1. to float with scale from 0 to 1
-            2. resize to (64, 128) as Market1501 dataset did
-            3. concatenate to a numpy array
-            3. to torch Tensor
-            4. normalize
-        """
-        im_batch = torch.cat([self.raw_transformer(im).unsqueeze(0) for im in im_crops], dim=0)
+        def _resize(im, size):
+            return cv2.resize(im.astype(np.float32)/255., size)
+
+        im_batch = torch.cat([self.norm(_resize(im, self.size)).unsqueeze(0) for im in im_crops], dim=0).float()
         return im_batch
 
+
     def __call__(self, im_crops):
         im_batch = self._preprocess(im_crops)
         with torch.no_grad():
@@ -89,3 +87,10 @@ class MyExtractor(object):
         return features.cpu().numpy()
 
 
+
+if __name__ == '__main__':
+    img = cv2.imread("demo.jpg")[:,:,(2,1,0)]
+    extr = Extractor("checkpoint/ckpt.t7")
+    feature = extr(img)
+    print(feature.shape)
+

deep_sort/deep/model.py

@@ -93,136 +93,12 @@ class Net(nn.Module):
         # classifier
         x = self.classifier(x)
         return x
-    
-# Define model
-# input:128x256x3
-# output: 64
-class MyNet(nn.Module):
-    def __init__(self):
-        super(MyNet, self).__init__()
-        self.conv1 = nn.Conv2d(3,8,5,1,2)
-        self.maxpool1 = nn.MaxPool2d(2) # 64x128x8
-        self.bn1   = nn.BatchNorm2d(8)
-        self.conv2 = nn.Conv2d(8,16,5,1,2)
-        self.maxpool2 = nn.MaxPool2d(2) # 32x64x16
-        self.bn2   = nn.BatchNorm2d(16)
-        self.conv3 = nn.Conv2d(16,32,5,1,2)
-        self.conv4 = nn.Conv2d(32,64,5,1,2)
-        self.maxpool3 = nn.MaxPool2d(2) # 16x32x64
-        self.bn3   = nn.BatchNorm2d(64)
-        self.conv5 = nn.Conv2d(64,32,5,1,2) # 16x32x32
-        self.maxpool4 = nn.MaxPool2d(2) # 8x16x32
-        self.bn4   = nn.BatchNorm2d(32)
-        self.conv6 = nn.Conv2d(32,64,8,8) #1x2x64
-        self.flat = nn.Flatten()
-        self.linear = nn.Linear(2*64,64) # 64D-feature ID
-        
-        self.sigmoid = nn.Sigmoid()
-        self.lrelu = nn.LeakyReLU()
 
-    def forward(self, x:torch.Tensor):   
-        x = self.conv1(x)
-        x = self.lrelu(x)
-        x = self.maxpool1(x)
-        x = self.bn1(x)
-        x = self.conv2(x)
-        x = self.lrelu(x)
-        x = self.maxpool2(x)
-        x = self.bn2(x)
-        x = self.conv3(x)
-        x = self.lrelu(x)
-        x = self.conv4(x)
-        x = self.lrelu(x)
-        x = self.maxpool3(x)
-        x = self.bn3(x)
-        x = self.conv5(x)
-        x = self.lrelu(x)
-        x = self.maxpool4(x)
-        x = self.bn4(x)
-        x = self.conv6(x)
-        x = self.sigmoid(x)
-        x = self.flat(x)
-        x = self.linear(x)
-        
-        return x
 
 if __name__ == '__main__':
     net = Net()
     x = torch.randn(4,3,128,64)
     y = net(x)
-    
-#########ResNet18
-class RestNetBasicBlock(nn.Module):
-    def __init__(self, in_channels, out_channels, stride):
-        super(RestNetBasicBlock, self).__init__()
-        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
-        self.bn1 = nn.BatchNorm2d(out_channels)
-        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1)
-        self.bn2 = nn.BatchNorm2d(out_channels)
-
-    def forward(self, x):
-        output = self.conv1(x)
-        output = F.relu(self.bn1(output))
-        output = self.conv2(output)
-        output = self.bn2(output)
-        return F.leaky_relu(x + output)
-
-
-class RestNetDownBlock(nn.Module):
-    def __init__(self, in_channels, out_channels, stride):
-        super(RestNetDownBlock, self).__init__()
-        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride[0], padding=1)
-        self.bn1 = nn.BatchNorm2d(out_channels)
-        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride[1], padding=1)
-        self.bn2 = nn.BatchNorm2d(out_channels)
-        self.extra = nn.Sequential(
-            nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride[0], padding=0),
-            nn.BatchNorm2d(out_channels)
-        )
-
-    def forward(self, x):
-        extra_x = self.extra(x)
-        output = self.conv1(x)
-        out = F.relu(self.bn1(output))
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-        return F.leaky_relu(extra_x + out)
-
-
-class RestNet18(nn.Module):
-    def __init__(self):
-        super(RestNet18, self).__init__()
-        self.conv1 = nn.Conv2d(3, 32, kernel_size=(3, 3), padding=1)
-        self.bn1 = nn.BatchNorm2d(32)
-        self.maxpool = nn.MaxPool2d(2)
-
-        self.layer1 = nn.Sequential(RestNetBasicBlock(32, 32, 1),
-                                    RestNetBasicBlock(32, 32, 1))
-
-        self.layer2 = nn.Sequential(RestNetDownBlock(32, 64, [2, 1]),
-                                    RestNetBasicBlock(64, 64, 1))
-
-        self.layer3 = nn.Sequential(RestNetDownBlock(64, 128, [2, 1]),
-                                    RestNetBasicBlock(128, 128, 1))
-
-        self.layer4 = nn.Sequential(RestNetDownBlock(128, 256, [2, 1]),
-                                    RestNetBasicBlock(256, 256, 1))
-
-        self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1, 1))
-
-        self.fc = nn.Linear(32768, 23)
-
-    def forward(self, x):
-        x = self.conv1(x)
-        x = self.layer1(x)
-        x = self.layer2(x)
-        x = self.layer3(x)
-        x = self.layer4(x)
-        # x = self.avgpool(x)
-        x = x.reshape(x.shape[0], -1)
-        # print(x.size())
-        x = self.fc(x)
-        return x
+    import ipdb; ipdb.set_trace()
 
 

deep_sort/deep_sort.py

@@ -1,7 +1,7 @@
 import numpy as np
 import torch
 
-from .deep.feature_extractor import Extractor,MyExtractor
+from .deep.feature_extractor import Extractor, FastReIDExtractor
 from .sort.nn_matching import NearestNeighborDistanceMetric
 from .sort.preprocessing import non_max_suppression
 from .sort.detection import Detection
@@ -12,14 +12,14 @@ __all__ = ['DeepSort']
 
 
 class DeepSort(object):
-    def __init__(self, model_path, model_config=None, max_dist=1e-2, min_confidence=0.35, nms_max_overlap=0.9, max_iou_distance=0.40, max_age=10, n_init=3, nn_budget=5000, use_cuda=True):
+    def __init__(self, model_path, model_config=None, max_dist=0.2, min_confidence=0.3, nms_max_overlap=1.0, max_iou_distance=0.7, max_age=70, n_init=3, nn_budget=100, use_cuda=True):
         self.min_confidence = min_confidence
         self.nms_max_overlap = nms_max_overlap
 
         if model_config is None:
             self.extractor = Extractor(model_path, use_cuda=use_cuda)
         else:
-            self.extractor = MyExtractor(model_path, use_cuda=use_cuda)
+            self.extractor = FastReIDExtractor(model_config, model_path, use_cuda=use_cuda)
 
         max_cosine_distance = max_dist
         metric = NearestNeighborDistanceMetric("cosine", max_cosine_distance, nn_budget)
@@ -30,7 +30,7 @@ class DeepSort(object):
         # generate detections
         features = self._get_features(bbox_xywh, ori_img)
         bbox_tlwh = self._xywh_to_tlwh(bbox_xywh)
-
+        print(confidences,bbox_tlwh,features)
         detections = [Detection(bbox_tlwh[i], conf, features[i]) for i,conf in enumerate(confidences) if conf>self.min_confidence]
 
         # run on non-maximum supression
@@ -42,21 +42,19 @@ class DeepSort(object):
         # update tracker
         self.tracker.predict()
         self.tracker.update(detections)
-        
+
         # output bbox identities
         outputs = []
-        output_features = []
         for track in self.tracker.tracks:
             if not track.is_confirmed() or track.time_since_update > 1:
                 continue
             box = track.to_tlwh()
             x1,y1,x2,y2 = self._tlwh_to_xyxy(box)
             track_id = track.track_id
-            output_features.append(track.features) 
             outputs.append(np.array([x1,y1,x2,y2,track_id], dtype=np.int))
         if len(outputs) > 0:
             outputs = np.stack(outputs,axis=0)
-        return outputs,output_features
+        return outputs,features
 
 
     """

deep_sort/sort/iou_matching.py

@@ -23,8 +23,6 @@ def iou(bbox, candidates):
         occluded by the candidate.
 
     """
-    length = len(candidates)
-
     bbox_tl, bbox_br = bbox[:2], bbox[:2] + bbox[2:]
     candidates_tl = candidates[:, :2]
     candidates_br = candidates[:, :2] + candidates[:, 2:]
@@ -38,14 +36,7 @@ def iou(bbox, candidates):
     area_intersection = wh.prod(axis=1)
     area_bbox = bbox[2:].prod()
     area_candidates = candidates[:, 2:].prod(axis=1)
-    
-    # should be consious
-    gious = []
-    for i in range(length):
-        gious.append(float((area_bbox + area_candidates[i] - area_intersection[i])/(max(bbox[0]+bbox[2],candidates[i][0]+candidates[i][2])-min(bbox[0],candidates[i][0]))/(max(bbox[1]+bbox[3],candidates[i][1]+candidates[i][3])-min(bbox[1],candidates[i][1]))) )
-        
-    # return area_intersection / (area_bbox + area_candidates - area_intersection)
-    return np.array(gious)
+    return area_intersection / (area_bbox + area_candidates - area_intersection)
 
 
 def iou_cost(tracks, detections, track_indices=None,

deep_sort/sort/tracker.py

@@ -86,8 +86,7 @@ class Tracker:
                 continue
             features += track.features
             targets += [track.track_id for _ in track.features]
-            # track.features = []
-            
+            track.features = []
         self.metric.partial_fit(
             np.asarray(features), np.asarray(targets), active_targets)
 
@@ -100,6 +99,7 @@ class Tracker:
             cost_matrix = linear_assignment.gate_cost_matrix(
                 self.kf, cost_matrix, tracks, dets, track_indices,
                 detection_indices)
+
             return cost_matrix
 
         # Split track set into confirmed and unconfirmed tracks.

detect.py

@@ -1,30 +0,0 @@
-import torch
-
-from models.experimental import attempt_load
-from utils.general import non_max_suppression
-
-import param
-
-
-class YOLOv7(object):
-    def __init__(self,option:param.Parameters) -> None:
-        # load FP32 model
-        self.opt = option
-        self.model = attempt_load(weights=self.opt.weights_yolo, device=self.opt.device)
-
-    def detect(self,imgs:torch.Tensor) -> list:
-        # img.size should be (Batch,3,H,W),every value should be range of [0,1]
-        # Both H and W % 64 == 0
-        # Initialize
-        model = self.model
-        opt   = self.opt
-        # Inference
-        with torch.no_grad():   # Calculating gradients would cause a GPU memory leak
-            pred = model(imgs, augment=opt.augment)[0]
-        # Apply NMS
-        pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
-        return pred
-
-        
-        
-        

example_reid.py

@@ -1,185 +0,0 @@
-import torch
-from torchvision import transforms
-import numpy as np
-import cv2
-import os
-
-import param
-from detect import YOLOv7
-from deep_sort import deep_sort as dsort
-
-def get_img_test(raw_img:np.array):
-    raw_transform = transforms.Compose([transforms.ToPILImage(),
-                                        transforms.Resize((360,640)),
-                                        transforms.Pad((0,(640-360)//2)),])
-    return raw_transform(raw_img)
-
-def reverse_box_get_from_yolo(detection):
-    x1 = int(detection[0]*1920/640)
-    x2 = int(detection[2]*1920/640)
-    y1 = int((detection[1]-140)*1080/360)
-    y2 = int((detection[3]-140)*1080/360)
-    return x1,y1,x2,y2
-
-def opencv_box_plot(img:cv2.Mat,pred_img:np.array):
-    pred_img = pred_img.astype(np.uint)
-    for info in pred_img:
-        x1,y1,x2,y2 = info[0],info[1],info[2],info[3]
-        cv2.rectangle(img,(x1,y1-140),(x2,y2-140),(255,0,0),2)
-
-    cv2.imshow("test",img)
-    cv2.waitKey(0)
-    cv2.imwrite("result.png",img)
-    
-def opencv_match_pointer_plot(img:cv2.Mat,start_base:tuple,text:str,color:tuple):
-    SIZE = 50
-    brush = [0,0]
-    brush[0]=start_base[0]
-    brush[1]=start_base[1]
-    brush_push = [brush[0]+SIZE,brush[1]-SIZE]
-    cv2.line(img,brush,brush_push,color,1)
-    # brush=brush_push
-    # brush_push[0]+=SIZE*2
-    # brush_push[1]+=0
-    # cv2.line(img,brush,brush_push,color,1)    
-    cv2.putText(img,text,brush_push, cv2.FONT_HERSHEY_PLAIN, 1, color, 2)
-    return img
-
-routine = {}
-    
-def opencv_sort_plot(img:cv2.Mat,pred_yolo:np.array,pred_sort:np.array):
-    for info in pred_yolo:
-        x1,y1,x2,y2 = int(info[0]),int(info[1]),int(info[2]),int(info[3])
-        cv2.rectangle(img,(x1,y1),(x2,y2),(0,255,0),1)    
-    for info in pred_sort:
-        x1,y1,x2,y2,id = info[0],info[1],info[2],info[3],info[4]
-        cxy = (int((x1+x2)/2),int((y1+y2)/2))
-        color = (int(id%3*100),int(id%4*75),int(id%5*50))
-        cv2.putText(img,str(id),cxy, cv2.FONT_HERSHEY_PLAIN, 1.0, color, 2)
-        if routine.get(id) is None:
-            routine[id] = [cxy]
-        else:
-            for i in range(1,len(routine[id])):
-                cv2.line(img,routine[id][i-1],routine[id][i],color,1)
-            cv2.line(img,routine[id][-1],cxy,color,1)
-            routine[id].append(cxy)
-    return img
-
-def extractidfeature(id_path:str,extractor):
-    img_id = cv2.imread(id_path)
-    tensor_id = extractor([img_id])
-    return tensor_id
-
-# example        
-if __name__ == '__main__':
-    MAX_BUFFLEN        = 10
-    DISTANCE_THRESHOLD = 0.4
-    SAVE_TXT_FLAG = True
-    
-    my_option = param.Parameters()
-    
-    print(torch.__version__)
-    print(my_option.device)
-    
-    query_path = f"./mydataset/query/cam{my_option.query_index}/"                   # query from .jpg photos
-    gallary_path = f"mydataset/video/cam{my_option.gallary_index}.mp4"              # source from mp4 via yolo
-    output_video_path = f"output/output{my_option.gallary_index}_{my_option.query_index}.mp4"    # output file path
-    output_txt_path = f"output/rank10_detection_{my_option.gallary_index}_{my_option.query_index}.txt"
-    
-    query_features = []
-    paths = os.listdir(query_path)
-    query_match_buff = [[] for p in paths]
-    query_matched = [-1 for p in paths]
-    
-    # load detection model and deepsort model
-    model = YOLOv7(my_option)         
-    deepsort = dsort.DeepSort(model_path=my_option.weights_reid,model_config="self",use_cuda=(torch.device("cuda:0") == my_option.device))
-    
-    # get query feature id
-    for p in paths:
-        query_features.append(extractidfeature(query_path+p,deepsort.extractor))
-        
-    with torch.no_grad():
-        sources = cv2.VideoCapture(gallary_path)
-        target = cv2.VideoWriter(output_video_path,cv2.VideoWriter_fourcc('m', 'p', '4', 'v'),24,(1920,1080))
-        frame_counter = 0
-        while True:
-            ret,frame = sources.read()
-            if ret is False:
-                break
-            
-            img    = get_img_test(frame)                        # get image as torch.Tensor with size of [1,1,640,640]
-            img_tensor = transforms.ToTensor()(img).unsqueeze(dim=0).to(my_option.device)     
-            detections = model.detect(img_tensor)               # get the sequence of result   
-            detections = detections[0].detach().cpu().numpy()   # the single img is index 0
-            
-            for detection in detections:
-                detection[0],detection[1],detection[2],detection[3] = reverse_box_get_from_yolo(detection[0:4])
-            
-            bbox_xywhs = []
-            confs      = []
-            for xyxycc in detections: 
-                xywh   =  deepsort._xyxy_to_xywh(xyxycc[0:4])
-                conf   =  xyxycc[4]
-                clas   =  int(xyxycc[5])
-                bbox_xywhs.append(xywh[:])
-                confs.append(conf)
-            dpsort, features = deepsort.update(bbox_xywh=np.array(bbox_xywhs),confidences=np.array(confs),ori_img=np.array(frame))
-            
-            img_drawing = np.array(frame)
-            img_drawing = opencv_sort_plot(img_drawing,detections,dpsort)
-            
-            # ReID: compute features
-            for i in range(len(dpsort)):
-                track = dpsort[i]
-                MOT_id = track[-1]
-                feature_distances = []
-                for j in range(len(query_features)):
-                    feature_distance = np.mean(np.power(features[i]-query_features[j],2))
-                    feature_distances.append(feature_distance)
-                    
-                # ReID: rank distances for every feature   
-                feature_distances_sorted = sorted(feature_distances)
-                for distance in feature_distances_sorted:
-                    if distance < DISTANCE_THRESHOLD:
-                        index_d = feature_distances.index(distance)
-                        query_match_buff[index_d].append((MOT_id,distance,(frame_counter,track[0],track[1],track[2],track[3]))) 
-                        query_match_buff[index_d].sort(key=lambda x:x[1],reverse=True)
-                        if len(query_match_buff[index_d]) > MAX_BUFFLEN:
-                            query_match_buff[index_d].pop()
-                            count = 0
-                            for term in query_match_buff[index_d]:
-                                if term[0] == MOT_id:
-                                    count+=1
-                            if count > MAX_BUFFLEN//3:
-                                cx = int((track[0]+track[2])/2)
-                                cy = int((track[1]+track[3])/2)
-                                color = (int(index_d%2*100),int(index_d%3*75),int(index_d%4*50))
-                                
-                                # img_drawing=opencv_match_pointer_plot(img_drawing,(cx,track[1]-10),f"{paths[index_d]}",color)
-                                # cv2.rectangle(img_drawing,(track[i][0],track[i][1]),(track[i][2],track[i][3]),color,1)                
-                                cv2.putText(img_drawing,f"q_{paths[index_d]}".split(".")[0],(track[0],track[1]), cv2.FONT_HERSHEY_PLAIN, 1, color, 2)
-                            break
-                    else:
-                        break
-            
-            # Final plot and show video
-            cv2.imshow("test",img_drawing)
-            cv2.waitKey(1)
-            target.write(img_drawing)
-            frame_counter += 1
-        
-        # save the video and end the program
-        target.release()
-        
-        # rank10 output as txt
-        if SAVE_TXT_FLAG is True:
-            with open(output_txt_path,mode="w+") as f:
-                result = []
-                for i in range(len(query_match_buff)):
-                    for raw_info in query_match_buff[i]:
-                        detect_info = raw_info[-1]
-                        result.append((detect_info[0],i,detect_info[1],detect_info[2],detect_info[3],detect_info[4]))
-                        result.sort(key=lambda x:x[0])
-                for info in result:
-                    f.write(str(info[0])+","+str(info[1])+","+str(info[2])+","+str(info[3])+","+str(info[4])+","+str(info[5])+"\n")

fastreid/__init__.py

@@ -0,0 +1,8 @@
+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+
+
+__version__ = "0.2.0"

fastreid/cfgs/Base-AGW.yml

@@ -0,0 +1,19 @@
+_BASE_: "Base-bagtricks.yml"
+
+MODEL:
+  BACKBONE:
+    WITH_NL: True
+
+  HEADS:
+    POOL_LAYER: "gempool"
+
+  LOSSES:
+    NAME: ("CrossEntropyLoss", "TripletLoss")
+    CE:
+      EPSILON: 0.1
+      SCALE: 1.0
+
+    TRI:
+      MARGIN: 0.0
+      HARD_MINING: False
+      SCALE: 1.0

fastreid/cfgs/Base-MGN.yml

@@ -0,0 +1,25 @@
+_BASE_: "Base-Strongerbaseline.yml"
+
+MODEL:
+  META_ARCHITECTURE: 'MGN'
+
+  FREEZE_LAYERS: ["backbone", "b1", "b2", "b3",]
+
+  BACKBONE:
+    WITH_NL: False
+
+  HEADS:
+    EMBEDDING_DIM: 256
+
+  LOSSES:
+    NAME: ("CrossEntropyLoss", "TripletLoss",)
+    CE:
+      EPSILON: 0.1
+      SCALE: 1.0
+
+    TRI:
+      MARGIN: 0.0
+      HARD_MINING: True
+      NORM_FEAT: False
+      SCALE: 1.0
+

fastreid/cfgs/Base-Strongerbaseline.yml

@@ -0,0 +1,59 @@
+_BASE_: "Base-bagtricks.yml"
+
+MODEL:
+  FREEZE_LAYERS: ["backbone"]
+
+  BACKBONE:
+    WITH_NL: True
+
+  HEADS:
+    NECK_FEAT: "after"
+    POOL_LAYER: "gempoolP"
+    CLS_LAYER: "circleSoftmax"
+    SCALE: 64
+    MARGIN: 0.35
+
+  LOSSES:
+    NAME: ("CrossEntropyLoss", "TripletLoss",)
+    CE:
+      EPSILON: 0.1
+      SCALE: 1.0
+    TRI:
+      MARGIN: 0.0
+      HARD_MINING: True
+      NORM_FEAT: False
+      SCALE: 1.0
+
+INPUT:
+  SIZE_TRAIN: [384, 128]
+  SIZE_TEST: [384, 128]
+  DO_AUTOAUG: True
+
+DATALOADER:
+  NUM_INSTANCE: 16
+
+SOLVER:
+  OPT: "Adam"
+  MAX_ITER: 60
+  BASE_LR: 0.00035
+  BIAS_LR_FACTOR: 1.
+  WEIGHT_DECAY: 0.0005
+  WEIGHT_DECAY_BIAS: 0.0005
+  IMS_PER_BATCH: 64
+
+  SCHED: "WarmupCosineAnnealingLR"
+  DELAY_ITERS: 30
+  ETA_MIN_LR: 0.00000077
+
+  WARMUP_FACTOR: 0.01
+  WARMUP_ITERS: 10
+  FREEZE_ITERS: 10
+
+  CHECKPOINT_PERIOD: 30
+
+TEST:
+  EVAL_PERIOD: 30
+  IMS_PER_BATCH: 128
+
+CUDNN_BENCHMARK: True
+

fastreid/cfgs/Base-bagtricks.yml

@@ -0,0 +1,73 @@
+MODEL:
+  META_ARCHITECTURE: "Baseline"
+
+  BACKBONE:
+    NAME: "build_resnet_backbone"
+    NORM: "BN"
+    DEPTH: "50x"
+    LAST_STRIDE: 1
+    FEAT_DIM: 2048
+    WITH_IBN: False
+    PRETRAIN: True
+    PRETRAIN_PATH: "/export/home/lxy/.cache/torch/checkpoints/resnet50-19c8e357.pth"
+
+  HEADS:
+    NAME: "EmbeddingHead"
+    NORM: "BN"
+    WITH_BNNECK: True
+    POOL_LAYER: "avgpool"
+    NECK_FEAT: "before"
+    CLS_LAYER: "linear"
+
+  LOSSES:
+    NAME: ("CrossEntropyLoss", "TripletLoss",)
+
+    CE:
+      EPSILON: 0.1
+      SCALE: 1.
+
+    TRI:
+      MARGIN: 0.3
+      HARD_MINING: True
+      NORM_FEAT: False
+      SCALE: 1.
+
+INPUT:
+  SIZE_TRAIN: [256, 128]
+  SIZE_TEST: [256, 128]
+  REA:
+    ENABLED: True
+    PROB: 0.5
+    MEAN: [123.675, 116.28, 103.53]
+  DO_PAD: True
+
+DATALOADER:
+  PK_SAMPLER: True
+  NAIVE_WAY: True
+  NUM_INSTANCE: 4
+  NUM_WORKERS: 8
+
+SOLVER:
+  OPT: "Adam"
+  MAX_ITER: 120
+  BASE_LR: 0.00035
+  BIAS_LR_FACTOR: 2.
+  WEIGHT_DECAY: 0.0005
+  WEIGHT_DECAY_BIAS: 0.0005
+  IMS_PER_BATCH: 64
+
+  SCHED: "WarmupMultiStepLR"
+  STEPS: [40, 90]
+  GAMMA: 0.1
+
+  WARMUP_FACTOR: 0.01
+  WARMUP_ITERS: 10
+
+  CHECKPOINT_PERIOD: 60
+
+TEST:
+  EVAL_PERIOD: 30
+  IMS_PER_BATCH: 128
+
+CUDNN_BENCHMARK: True
+

fastreid/cfgs/DukeMTMC/AGW_R101-ibn.yml

@@ -0,0 +1,12 @@
+_BASE_: "../Base-AGW.yml"
+
+MODEL:
+  BACKBONE:
+    DEPTH: "101x"
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("DukeMTMC",)
+  TESTS: ("DukeMTMC",)
+
+OUTPUT_DIR: "logs/dukemtmc/agw_R101-ibn"

fastreid/cfgs/DukeMTMC/AGW_R50-ibn.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-AGW.yml"
+
+MODEL:
+  BACKBONE:
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("DukeMTMC",)
+  TESTS: ("DukeMTMC",)
+
+OUTPUT_DIR: "logs/dukemtmc/agw_R50-ibn"

fastreid/cfgs/DukeMTMC/AGW_R50.yml

@@ -0,0 +1,7 @@
+_BASE_: "../Base-AGW.yml"
+
+DATASETS:
+  NAMES: ("DukeMTMC",)
+  TESTS: ("DukeMTMC",)
+
+OUTPUT_DIR: "logs/dukemtmc/agw_R50"

fastreid/cfgs/DukeMTMC/AGW_S50.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-AGW.yml"
+
+MODEL:
+  BACKBONE:
+    NAME: "build_resnest_backbone"
+
+DATASETS:
+  NAMES: ("DukeMTMC",)
+  TESTS: ("DukeMTMC",)
+
+OUTPUT_DIR: "logs/dukemtmc/agw_S50"

fastreid/cfgs/DukeMTMC/bagtricks_R101-ibn.yml

@@ -0,0 +1,12 @@
+_BASE_: "../Base-bagtricks.yml"
+
+MODEL:
+  BACKBONE:
+    DEPTH: "101x"
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("DukeMTMC",)
+  TESTS: ("DukeMTMC",)
+
+OUTPUT_DIR: "logs/dukemtmc/bagtricks_R101-ibn"

fastreid/cfgs/DukeMTMC/bagtricks_R50-ibn.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-bagtricks.yml"
+
+MODEL:
+  BACKBONE:
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("DukeMTMC",)
+  TESTS: ("DukeMTMC",)
+
+OUTPUT_DIR: "logs/dukemtmc/bagtricks_R50-ibn"

fastreid/cfgs/DukeMTMC/bagtricks_R50.yml

@@ -0,0 +1,7 @@
+_BASE_: "../Base-bagtricks.yml"
+
+DATASETS:
+  NAMES: ("DukeMTMC",)
+  TESTS: ("DukeMTMC",)
+
+OUTPUT_DIR: "logs/dukemtmc/bagtricks_R50"

fastreid/cfgs/DukeMTMC/bagtricks_S50.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-bagtricks.yml"
+
+MODEL:
+  BACKBONE:
+    NAME: "build_resnest_backbone"
+
+DATASETS:
+  NAMES: ("DukeMTMC",)
+  TESTS: ("DukeMTMC",)
+
+OUTPUT_DIR: "logs/dukemtmc/bagtricks_S50"

fastreid/cfgs/DukeMTMC/mgn_R50-ibn.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-MGN.yml"
+
+MODEL:
+  BACKBONE:
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("DukeMTMC",)
+  TESTS: ("DukeMTMC",)
+
+OUTPUT_DIR: "logs/dukemtmc/mgn_R50-ibn"

fastreid/cfgs/DukeMTMC/sbs_R101-ibn.yml

@@ -0,0 +1,12 @@
+_BASE_: "../Base-Strongerbaseline.yml"
+
+MODEL:
+  BACKBONE:
+    DEPTH: "101x"
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("DukeMTMC",)
+  TESTS: ("DukeMTMC",)
+
+OUTPUT_DIR: "logs/dukemtmc/sbs_R101-ibn"

fastreid/cfgs/DukeMTMC/sbs_R50-ibn.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-Strongerbaseline.yml"
+
+MODEL:
+  BACKBONE:
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("DukeMTMC",)
+  TESTS: ("DukeMTMC",)
+
+OUTPUT_DIR: "logs/dukemtmc/sbs_R50-ibn"

fastreid/cfgs/DukeMTMC/sbs_R50.yml

@@ -0,0 +1,7 @@
+_BASE_: "../Base-Strongerbaseline.yml"
+
+DATASETS:
+  NAMES: ("DukeMTMC",)
+  TESTS: ("DukeMTMC",)
+
+OUTPUT_DIR: "logs/dukemtmc/sbs_R50"

fastreid/cfgs/DukeMTMC/sbs_S50.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-Strongerbaseline.yml"
+
+MODEL:
+  BACKBONE:
+    NAME: "build_resnest_backbone"
+
+DATASETS:
+  NAMES: ("DukeMTMC",)
+  TESTS: ("DukeMTMC",)
+
+OUTPUT_DIR: "logs/dukemtmc/sbs_S50"

fastreid/cfgs/MSMT17/AGW_R101-ibn.yml

@@ -0,0 +1,12 @@
+_BASE_: "../Base-AGW.yml"
+
+MODEL:
+  BACKBONE:
+    DEPTH: "101x"
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("MSMT17",)
+  TESTS: ("MSMT17",)
+
+OUTPUT_DIR: "logs/msmt17/agw_R101-ibn"

fastreid/cfgs/MSMT17/AGW_R50-ibn.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-AGW.yml"
+
+MODEL:
+  BACKBONE:
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("MSMT17",)
+  TESTS: ("MSMT17",)
+
+OUTPUT_DIR: "logs/msmt17/agw_R50-ibn"

fastreid/cfgs/MSMT17/AGW_R50.yml

@@ -0,0 +1,7 @@
+_BASE_: "../Base-AGW.yml"
+
+DATASETS:
+  NAMES: ("MSMT17",)
+  TESTS: ("MSMT17",)
+
+OUTPUT_DIR: "logs/msmt17/agw_R50"

fastreid/cfgs/MSMT17/AGW_S50.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-AGW.yml"
+
+MODEL:
+  BACKBONE:
+    NAME: "build_resnest_backbone"
+
+DATASETS:
+  NAMES: ("MSMT17",)
+  TESTS: ("MSMT17",)
+
+OUTPUT_DIR: "logs/msmt17/agw_S50"

fastreid/cfgs/MSMT17/bagtricks_R101-ibn.yml

@@ -0,0 +1,13 @@
+_BASE_: "../Base-bagtricks.yml"
+
+MODEL:
+  BACKBONE:
+    DEPTH: "101x"
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("MSMT17",)
+  TESTS: ("MSMT17",)
+
+OUTPUT_DIR: "logs/msmt17/bagtricks_R101-ibn"
+

fastreid/cfgs/MSMT17/bagtricks_R50-ibn.yml

@@ -0,0 +1,12 @@
+_BASE_: "../Base-bagtricks.yml"
+
+MODEL:
+  BACKBONE:
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("MSMT17",)
+  TESTS: ("MSMT17",)
+
+OUTPUT_DIR: "logs/msmt17/bagtricks_R50-ibn"
+

fastreid/cfgs/MSMT17/bagtricks_R50.yml

@@ -0,0 +1,7 @@
+_BASE_: "../Base-bagtricks.yml"
+
+DATASETS:
+  NAMES: ("MSMT17",)
+  TESTS: ("MSMT17",)
+
+OUTPUT_DIR: "logs/msmt17/bagtricks_R50"

fastreid/cfgs/MSMT17/bagtricks_S50.yml

@@ -0,0 +1,12 @@
+_BASE_: "../Base-bagtricks.yml"
+
+MODEL:
+  BACKBONE:
+    NAME: "build_resnest_backbone"
+
+DATASETS:
+  NAMES: ("MSMT17",)
+  TESTS: ("MSMT17",)
+
+OUTPUT_DIR: "logs/msmt17/bagtricks_S50"
+

fastreid/cfgs/MSMT17/mgn_R50-ibn.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-MGN.yml"
+
+MODEL:
+  BACKBONE:
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("MSMT17",)
+  TESTS: ("MSMT17",)
+
+OUTPUT_DIR: "logs/msmt17/mgn_R50-ibn"

fastreid/cfgs/MSMT17/sbs_R101-ibn.yml

@@ -0,0 +1,12 @@
+_BASE_: "../Base-Strongerbaseline.yml"
+
+MODEL:
+  BACKBONE:
+    DEPTH: "101x"
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("MSMT17",)
+  TESTS: ("MSMT17",)
+
+OUTPUT_DIR: "logs/msmt17/sbs_R101-ibn"

fastreid/cfgs/MSMT17/sbs_R50-ibn.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-Strongerbaseline.yml"
+
+MODEL:
+  BACKBONE:
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("MSMT17",)
+  TESTS: ("MSMT17",)
+
+OUTPUT_DIR: "logs/msmt17/sbs_R50-ibn"

fastreid/cfgs/MSMT17/sbs_R50.yml

@@ -0,0 +1,7 @@
+_BASE_: "../Base-Strongerbaseline.yml"
+
+DATASETS:
+  NAMES: ("MSMT17",)
+  TESTS: ("MSMT17",)
+
+OUTPUT_DIR: "logs/msmt17/sbs_R50"

fastreid/cfgs/MSMT17/sbs_S50.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-Strongerbaseline.yml"
+
+MODEL:
+  BACKBONE:
+    NAME: "build_resnest_backbone"
+
+DATASETS:
+  NAMES: ("MSMT17",)
+  TESTS: ("MSMT17",)
+
+OUTPUT_DIR: "logs/msmt17/sbs_S50"

fastreid/cfgs/Market1501/AGW_R101-ibn.yml

@@ -0,0 +1,12 @@
+_BASE_: "../Base-AGW.yml"
+
+MODEL:
+  BACKBONE:
+    DEPTH: "101x"
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("Market1501",)
+  TESTS: ("Market1501",)
+
+OUTPUT_DIR: "logs/market1501/agw_R101-ibn"

fastreid/cfgs/Market1501/AGW_R50-ibn.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-AGW.yml"
+
+MODEL:
+  BACKBONE:
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("Market1501",)
+  TESTS: ("Market1501",)
+
+OUTPUT_DIR: "logs/market1501/agw_R50-ibn"

fastreid/cfgs/Market1501/AGW_R50.yml

@@ -0,0 +1,7 @@
+_BASE_: "../Base-AGW.yml"
+
+DATASETS:
+  NAMES: ("Market1501",)
+  TESTS: ("Market1501",)
+
+OUTPUT_DIR: "logs/market1501/agw_R50"

fastreid/cfgs/Market1501/AGW_S50.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-AGW.yml"
+
+MODEL:
+  BACKBONE:
+    NAME: "build_resnest_backbone"
+
+DATASETS:
+  NAMES: ("Market1501",)
+  TESTS: ("Market1501",)
+
+OUTPUT_DIR: "logs/market1501/agw_S50"

fastreid/cfgs/Market1501/bagtricks_R101-ibn.yml

@@ -0,0 +1,12 @@
+_BASE_: "../Base-bagtricks.yml"
+
+MODEL:
+  BACKBONE:
+    DEPTH: "101x"
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("Market1501",)
+  TESTS: ("Market1501",)
+
+OUTPUT_DIR: "logs/market1501/bagtricks_R101-ibn"

fastreid/cfgs/Market1501/bagtricks_R50-ibn.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-bagtricks.yml"
+
+MODEL:
+  BACKBONE:
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("Market1501",)
+  TESTS: ("Market1501",)
+
+OUTPUT_DIR: "logs/market1501/bagtricks_R50-ibn"

fastreid/cfgs/Market1501/bagtricks_R50.yml

@@ -0,0 +1,7 @@
+_BASE_: "../Base-bagtricks.yml"
+
+DATASETS:
+  NAMES: ("Market1501",)
+  TESTS: ("Market1501",)
+
+OUTPUT_DIR: "logs/market1501/bagtricks_R50"

fastreid/cfgs/Market1501/bagtricks_S50.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-bagtricks.yml"
+
+MODEL:
+  BACKBONE:
+    NAME: "build_resnest_backbone"
+
+DATASETS:
+  NAMES: ("Market1501",)
+  TESTS: ("Market1501",)
+
+OUTPUT_DIR: "logs/market1501/bagtricks_S50"

fastreid/cfgs/Market1501/mgn_R50-ibn.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-MGN.yml"
+
+MODEL:
+  BACKBONE:
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("Market1501",)
+  TESTS: ("Market1501",)
+
+OUTPUT_DIR: "logs/market1501/mgn_R50-ibn"

fastreid/cfgs/Market1501/sbs_R101-ibn.yml

@@ -0,0 +1,12 @@
+_BASE_: "../Base-Strongerbaseline.yml"
+
+MODEL:
+  BACKBONE:
+    DEPTH: "101x"
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("Market1501",)
+  TESTS: ("Market1501",)
+
+OUTPUT_DIR: "logs/market1501/sbs_R101-ibn"

fastreid/cfgs/Market1501/sbs_R50-ibn.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-Strongerbaseline.yml"
+
+MODEL:
+  BACKBONE:
+    WITH_IBN: True
+
+DATASETS:
+  NAMES: ("Market1501",)
+  TESTS: ("Market1501",)
+
+OUTPUT_DIR: "logs/market1501/sbs_R50-ibn"

fastreid/cfgs/Market1501/sbs_R50.yml

@@ -0,0 +1,7 @@
+_BASE_: "../Base-Strongerbaseline.yml"
+
+DATASETS:
+  NAMES: ("Market1501",)
+  TESTS: ("Market1501",)
+
+OUTPUT_DIR: "logs/market1501/sbs_R50"

fastreid/cfgs/Market1501/sbs_S50.yml

@@ -0,0 +1,11 @@
+_BASE_: "../Base-Strongerbaseline.yml"
+
+MODEL:
+  BACKBONE:
+    NAME: "build_resnest_backbone"
+
+DATASETS:
+  NAMES: ("Market1501",)
+  TESTS: ("Market1501",)
+
+OUTPUT_DIR: "logs/market1501/sbs_S50"

fastreid/cfgs/VERIWild/bagtricks_R50-ibn.yml

@@ -0,0 +1,33 @@
+_BASE_: "../Base-bagtricks.yml"
+
+INPUT:
+  SIZE_TRAIN: [256, 256]
+  SIZE_TEST: [256, 256]
+
+MODEL:
+  BACKBONE:
+    WITH_IBN: True
+  HEADS:
+    POOL_LAYER: gempool
+  LOSSES:
+    TRI:
+      HARD_MINING: False
+      MARGIN: 0.0
+
+DATASETS:
+  NAMES: ("VeRiWild",)
+  TESTS: ("SmallVeRiWild", "MediumVeRiWild", "LargeVeRiWild",)
+
+SOLVER:
+  IMS_PER_BATCH: 128
+  MAX_ITER: 60
+  STEPS: [30, 50]
+  WARMUP_ITERS: 10
+
+  CHECKPOINT_PERIOD: 20
+
+TEST:
+  EVAL_PERIOD: 20
+  IMS_PER_BATCH: 128
+
+OUTPUT_DIR: "logs/veriwild/bagtricks_R50-ibn_4gpu"

fastreid/cfgs/VeRi/sbs_R50-ibn.yml

@@ -0,0 +1,32 @@
+_BASE_: "../Base-Strongerbaseline.yml"
+
+INPUT:
+  SIZE_TRAIN: [256, 256]
+  SIZE_TEST: [256, 256]
+
+MODEL:
+  BACKBONE:
+    WITH_IBN: True
+
+SOLVER:
+  OPT: "SGD"
+  BASE_LR: 0.01
+  ETA_MIN_LR: 7.7e-5
+
+  IMS_PER_BATCH: 64
+  MAX_ITER: 60
+  DELAY_ITERS: 30
+  WARMUP_ITERS: 10
+  FREEZE_ITERS: 10
+
+  CHECKPOINT_PERIOD: 20
+
+DATASETS:
+  NAMES: ("VeRi",)
+  TESTS: ("VeRi",)
+
+TEST:
+  EVAL_PERIOD: 20
+  IMS_PER_BATCH: 128
+
+OUTPUT_DIR: "logs/veri/sbs_R50-ibn"

fastreid/cfgs/VehicleID/bagtricks_R50-ibn.yml

@@ -0,0 +1,35 @@
+_BASE_: "../Base-bagtricks.yml"
+
+INPUT:
+  SIZE_TRAIN: [256, 256]
+  SIZE_TEST: [256, 256]
+
+MODEL:
+  BACKBONE:
+    WITH_IBN: True
+  HEADS:
+    POOL_LAYER: gempool
+  LOSSES:
+    TRI:
+      HARD_MINING: False
+      MARGIN: 0.0
+
+DATASETS:
+  NAMES: ("VehicleID",)
+  TESTS: ("SmallVehicleID", "MediumVehicleID", "LargeVehicleID",)
+
+SOLVER:
+  BIAS_LR_FACTOR: 1.
+
+  IMS_PER_BATCH: 512
+  MAX_ITER: 60
+  STEPS: [30, 50]
+  WARMUP_ITERS: 10
+
+  CHECKPOINT_PERIOD: 20
+
+TEST:
+  EVAL_PERIOD: 20
+  IMS_PER_BATCH: 128
+
+OUTPUT_DIR: "logs/vehicleid/bagtricks_R50-ibn_4gpu"

fastreid/config/__init__.py

@@ -0,0 +1,8 @@
+# encoding: utf-8
+"""
+@author:  l1aoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+
+from .config import CfgNode, get_cfg
+from .defaults import _C as cfg

fastreid/config/config.py

@@ -0,0 +1,159 @@
+# encoding: utf-8
+"""
+@author:  l1aoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+
+import logging
+import os
+from typing import Any
+
+import yaml
+from yacs.config import CfgNode as _CfgNode
+
+from ..utils.file_io import PathManager
+
+BASE_KEY = "_BASE_"
+
+
+class CfgNode(_CfgNode):
+    """
+    Our own extended version of :class:`yacs.config.CfgNode`.
+    It contains the following extra features:
+    1. The :meth:`merge_from_file` method supports the "_BASE_" key,
+       which allows the new CfgNode to inherit all the attributes from the
+       base configuration file.
+    2. Keys that start with "COMPUTED_" are treated as insertion-only
+       "computed" attributes. They can be inserted regardless of whether
+       the CfgNode is frozen or not.
+    3. With "allow_unsafe=True", it supports pyyaml tags that evaluate
+       expressions in config. See examples in
+       https://pyyaml.org/wiki/PyYAMLDocumentation#yaml-tags-and-python-types
+       Note that this may lead to arbitrary code execution: you must not
+       load a config file from untrusted sources before manually inspecting
+       the content of the file.
+    """
+
+    @staticmethod
+    def load_yaml_with_base(filename: str, allow_unsafe: bool = False):
+        """
+        Just like `yaml.load(open(filename))`, but inherit attributes from its
+            `_BASE_`.
+        Args:
+            filename (str): the file name of the current config. Will be used to
+                find the base config file.
+            allow_unsafe (bool): whether to allow loading the config file with
+                `yaml.unsafe_load`.
+        Returns:
+            (dict): the loaded yaml
+        """
+        with PathManager.open(filename, "r") as f:
+            try:
+                cfg = yaml.safe_load(f)
+            except yaml.constructor.ConstructorError:
+                if not allow_unsafe:
+                    raise
+                logger = logging.getLogger(__name__)
+                logger.warning(
+                    "Loading config {} with yaml.unsafe_load. Your machine may "
+                    "be at risk if the file contains malicious content.".format(
+                        filename
+                    )
+                )
+                f.close()
+                with open(filename, "r") as f:
+                    cfg = yaml.unsafe_load(f)
+
+        def merge_a_into_b(a, b):
+            # merge dict a into dict b. values in a will overwrite b.
+            for k, v in a.items():
+                if isinstance(v, dict) and k in b:
+                    assert isinstance(
+                        b[k], dict
+                    ), "Cannot inherit key '{}' from base!".format(k)
+                    merge_a_into_b(v, b[k])
+                else:
+                    b[k] = v
+
+        if BASE_KEY in cfg:
+            base_cfg_file = cfg[BASE_KEY]
+            if base_cfg_file.startswith("~"):
+                base_cfg_file = os.path.expanduser(base_cfg_file)
+            if not any(
+                    map(base_cfg_file.startswith, ["/", "https://", "http://"])
+            ):
+                # the path to base cfg is relative to the config file itself.
+                base_cfg_file = os.path.join(
+                    os.path.dirname(filename), base_cfg_file
+                )
+            base_cfg = CfgNode.load_yaml_with_base(
+                base_cfg_file, allow_unsafe=allow_unsafe
+            )
+            del cfg[BASE_KEY]
+
+            merge_a_into_b(cfg, base_cfg)
+            return base_cfg
+        return cfg
+
+    def merge_from_file(self, cfg_filename: str, allow_unsafe: bool = False):
+        """
+        Merge configs from a given yaml file.
+        Args:
+            cfg_filename: the file name of the yaml config.
+            allow_unsafe: whether to allow loading the config file with
+                `yaml.unsafe_load`.
+        """
+        loaded_cfg = CfgNode.load_yaml_with_base(
+            cfg_filename, allow_unsafe=allow_unsafe
+        )
+        loaded_cfg = type(self)(loaded_cfg)
+        self.merge_from_other_cfg(loaded_cfg)
+
+    # Forward the following calls to base, but with a check on the BASE_KEY.
+    def merge_from_other_cfg(self, cfg_other):
+        """
+        Args:
+            cfg_other (CfgNode): configs to merge from.
+        """
+        assert (
+                BASE_KEY not in cfg_other
+        ), "The reserved key '{}' can only be used in files!".format(BASE_KEY)
+        return super().merge_from_other_cfg(cfg_other)
+
+    def merge_from_list(self, cfg_list: list):
+        """
+        Args:
+            cfg_list (list): list of configs to merge from.
+        """
+        keys = set(cfg_list[0::2])
+        assert (
+                BASE_KEY not in keys
+        ), "The reserved key '{}' can only be used in files!".format(BASE_KEY)
+        return super().merge_from_list(cfg_list)
+
+    def __setattr__(self, name: str, val: Any):
+        if name.startswith("COMPUTED_"):
+            if name in self:
+                old_val = self[name]
+                if old_val == val:
+                    return
+                raise KeyError(
+                    "Computed attributed '{}' already exists "
+                    "with a different value! old={}, new={}.".format(
+                        name, old_val, val
+                    )
+                )
+            self[name] = val
+        else:
+            super().__setattr__(name, val)
+
+
+def get_cfg() -> CfgNode:
+    """
+    Get a copy of the default config.
+    Returns:
+        a fastreid CfgNode instance.
+    """
+    from .defaults import _C
+
+    return _C.clone()

fastreid/config/defaults.py

@@ -0,0 +1,273 @@
+from .config import CfgNode as CN
+
+# -----------------------------------------------------------------------------
+# Convention about Training / Test specific parameters
+# -----------------------------------------------------------------------------
+# Whenever an argument can be either used for training or for testing, the
+# corresponding name will be post-fixed by a _TRAIN for a training parameter,
+# or _TEST for a test-specific parameter.
+# For example, the number of images during training will be
+# IMAGES_PER_BATCH_TRAIN, while the number of images for testing will be
+# IMAGES_PER_BATCH_TEST
+
+# -----------------------------------------------------------------------------
+# Config definition
+# -----------------------------------------------------------------------------
+
+_C = CN()
+
+# -----------------------------------------------------------------------------
+# MODEL
+# -----------------------------------------------------------------------------
+_C.MODEL = CN()
+_C.MODEL.DEVICE = "cuda"
+_C.MODEL.META_ARCHITECTURE = 'Baseline'
+_C.MODEL.FREEZE_LAYERS = ['']
+
+# ---------------------------------------------------------------------------- #
+# Backbone options
+# ---------------------------------------------------------------------------- #
+_C.MODEL.BACKBONE = CN()
+
+_C.MODEL.BACKBONE.NAME = "build_resnet_backbone"
+_C.MODEL.BACKBONE.DEPTH = "50x"
+_C.MODEL.BACKBONE.LAST_STRIDE = 1
+# Backbone feature dimension
+_C.MODEL.BACKBONE.FEAT_DIM = 2048
+# Normalization method for the convolution layers.
+_C.MODEL.BACKBONE.NORM = "BN"
+# If use IBN block in backbone
+_C.MODEL.BACKBONE.WITH_IBN = False
+# If use SE block in backbone
+_C.MODEL.BACKBONE.WITH_SE = False
+# If use Non-local block in backbone
+_C.MODEL.BACKBONE.WITH_NL = False
+# If use ImageNet pretrain model
+_C.MODEL.BACKBONE.PRETRAIN = True
+# Pretrain model path
+_C.MODEL.BACKBONE.PRETRAIN_PATH = ''
+
+# ---------------------------------------------------------------------------- #
+# REID HEADS options
+# ---------------------------------------------------------------------------- #
+_C.MODEL.HEADS = CN()
+_C.MODEL.HEADS.NAME = "EmbeddingHead"
+# Normalization method for the convolution layers.
+_C.MODEL.HEADS.NORM = "BN"
+# Number of identity
+_C.MODEL.HEADS.NUM_CLASSES = 0
+# Embedding dimension in head
+_C.MODEL.HEADS.EMBEDDING_DIM = 0
+# If use BNneck in embedding
+_C.MODEL.HEADS.WITH_BNNECK = True
+# Triplet feature using feature before(after) bnneck
+_C.MODEL.HEADS.NECK_FEAT = "before"  # options: before, after
+# Pooling layer type
+_C.MODEL.HEADS.POOL_LAYER = "avgpool"
+
+# Classification layer type
+_C.MODEL.HEADS.CLS_LAYER = "linear"  # "arcSoftmax" or "circleSoftmax"
+
+# Margin and Scale for margin-based classification layer
+_C.MODEL.HEADS.MARGIN = 0.15
+_C.MODEL.HEADS.SCALE = 128
+
+# ---------------------------------------------------------------------------- #
+# REID LOSSES options
+# ---------------------------------------------------------------------------- #
+_C.MODEL.LOSSES = CN()
+_C.MODEL.LOSSES.NAME = ("CrossEntropyLoss",)
+
+# Cross Entropy Loss options
+_C.MODEL.LOSSES.CE = CN()
+# if epsilon == 0, it means no label smooth regularization,
+# if epsilon == -1, it means adaptive label smooth regularization
+_C.MODEL.LOSSES.CE.EPSILON = 0.0
+_C.MODEL.LOSSES.CE.ALPHA = 0.2
+_C.MODEL.LOSSES.CE.SCALE = 1.0
+
+# Triplet Loss options
+_C.MODEL.LOSSES.TRI = CN()
+_C.MODEL.LOSSES.TRI.MARGIN = 0.3
+_C.MODEL.LOSSES.TRI.NORM_FEAT = False
+_C.MODEL.LOSSES.TRI.HARD_MINING = True
+_C.MODEL.LOSSES.TRI.SCALE = 1.0
+
+# Circle Loss options
+_C.MODEL.LOSSES.CIRCLE = CN()
+_C.MODEL.LOSSES.CIRCLE.MARGIN = 0.25
+_C.MODEL.LOSSES.CIRCLE.ALPHA = 128
+_C.MODEL.LOSSES.CIRCLE.SCALE = 1.0
+
+# Focal Loss options
+_C.MODEL.LOSSES.FL = CN()
+_C.MODEL.LOSSES.FL.ALPHA = 0.25
+_C.MODEL.LOSSES.FL.GAMMA = 2
+_C.MODEL.LOSSES.FL.SCALE = 1.0
+
+# Path to a checkpoint file to be loaded to the model. You can find available models in the model zoo.
+_C.MODEL.WEIGHTS = ""
+
+# Values to be used for image normalization
+_C.MODEL.PIXEL_MEAN = [0.485*255, 0.456*255, 0.406*255]
+# Values to be used for image normalization
+_C.MODEL.PIXEL_STD = [0.229*255, 0.224*255, 0.225*255]
+
+
+# -----------------------------------------------------------------------------
+# INPUT
+# -----------------------------------------------------------------------------
+_C.INPUT = CN()
+# Size of the image during training
+_C.INPUT.SIZE_TRAIN = [256, 128]
+# Size of the image during test
+_C.INPUT.SIZE_TEST = [256, 128]
+
+# Random probability for image horizontal flip
+_C.INPUT.DO_FLIP = True
+_C.INPUT.FLIP_PROB = 0.5
+
+# Value of padding size
+_C.INPUT.DO_PAD = True
+_C.INPUT.PADDING_MODE = 'constant'
+_C.INPUT.PADDING = 10
+# Random color jitter
+_C.INPUT.CJ = CN()
+_C.INPUT.CJ.ENABLED = False
+_C.INPUT.CJ.PROB = 0.5
+_C.INPUT.CJ.BRIGHTNESS = 0.15
+_C.INPUT.CJ.CONTRAST = 0.15
+_C.INPUT.CJ.SATURATION = 0.1
+_C.INPUT.CJ.HUE = 0.1
+# Auto augmentation
+_C.INPUT.DO_AUTOAUG = False
+# Augmix augmentation
+_C.INPUT.DO_AUGMIX = False
+# Random Erasing
+_C.INPUT.REA = CN()
+_C.INPUT.REA.ENABLED = False
+_C.INPUT.REA.PROB = 0.5
+_C.INPUT.REA.MEAN = [0.596*255, 0.558*255, 0.497*255]  # [0.485*255, 0.456*255, 0.406*255]
+# Random Patch
+_C.INPUT.RPT = CN()
+_C.INPUT.RPT.ENABLED = False
+_C.INPUT.RPT.PROB = 0.5
+
+# -----------------------------------------------------------------------------
+# Dataset
+# -----------------------------------------------------------------------------
+_C.DATASETS = CN()
+# List of the dataset names for training
+_C.DATASETS.NAMES = ("Market1501",)
+# List of the dataset names for testing
+_C.DATASETS.TESTS = ("Market1501",)
+# Combine trainset and testset joint training
+_C.DATASETS.COMBINEALL = False
+
+# -----------------------------------------------------------------------------
+# DataLoader
+# -----------------------------------------------------------------------------
+_C.DATALOADER = CN()
+# P/K Sampler for data loading
+_C.DATALOADER.PK_SAMPLER = True
+# Naive sampler which don't consider balanced identity sampling
+_C.DATALOADER.NAIVE_WAY = False
+# Number of instance for each person
+_C.DATALOADER.NUM_INSTANCE = 4
+_C.DATALOADER.NUM_WORKERS = 8
+
+# ---------------------------------------------------------------------------- #
+# Solver
+# ---------------------------------------------------------------------------- #
+_C.SOLVER = CN()
+
+# AUTOMATIC MIXED PRECISION
+_C.SOLVER.AMP_ENABLED = False
+
+# Optimizer
+_C.SOLVER.OPT = "Adam"
+
+_C.SOLVER.MAX_ITER = 120
+
+_C.SOLVER.BASE_LR = 3e-4
+_C.SOLVER.BIAS_LR_FACTOR = 1.
+_C.SOLVER.HEADS_LR_FACTOR = 1.
+
+_C.SOLVER.MOMENTUM = 0.9
+
+_C.SOLVER.WEIGHT_DECAY = 0.0005
+_C.SOLVER.WEIGHT_DECAY_BIAS = 0.
+
+# Multi-step learning rate options
+_C.SOLVER.SCHED = "WarmupMultiStepLR"
+_C.SOLVER.GAMMA = 0.1
+_C.SOLVER.STEPS = [30, 55]
+
+# Cosine annealing learning rate options
+_C.SOLVER.DELAY_ITERS = 0
+_C.SOLVER.ETA_MIN_LR = 3e-7
+
+# Warmup options
+_C.SOLVER.WARMUP_FACTOR = 0.1
+_C.SOLVER.WARMUP_ITERS = 10
+_C.SOLVER.WARMUP_METHOD = "linear"
+
+_C.SOLVER.FREEZE_ITERS = 0
+
+# SWA options
+_C.SOLVER.SWA = CN()
+_C.SOLVER.SWA.ENABLED = False
+_C.SOLVER.SWA.ITER = 10
+_C.SOLVER.SWA.PERIOD = 2
+_C.SOLVER.SWA.LR_FACTOR = 10.
+_C.SOLVER.SWA.ETA_MIN_LR = 3.5e-6
+_C.SOLVER.SWA.LR_SCHED = False
+
+_C.SOLVER.CHECKPOINT_PERIOD = 20
+
+# Number of images per batch across all machines.
+# This is global, so if we have 8 GPUs and IMS_PER_BATCH = 16, each GPU will
+# see 2 images per batch
+_C.SOLVER.IMS_PER_BATCH = 64
+
+# This is global, so if we have 8 GPUs and IMS_PER_BATCH = 16, each GPU will
+# see 2 images per batch
+_C.TEST = CN()
+
+_C.TEST.EVAL_PERIOD = 20
+
+# Number of images per batch in one process.
+_C.TEST.IMS_PER_BATCH = 64
+_C.TEST.METRIC = "cosine"
+_C.TEST.ROC_ENABLED = False
+
+# Average query expansion
+_C.TEST.AQE = CN()
+_C.TEST.AQE.ENABLED = False
+_C.TEST.AQE.ALPHA = 3.0
+_C.TEST.AQE.QE_TIME = 1
+_C.TEST.AQE.QE_K = 5
+
+# Re-rank
+_C.TEST.RERANK = CN()
+_C.TEST.RERANK.ENABLED = False
+_C.TEST.RERANK.K1 = 20
+_C.TEST.RERANK.K2 = 6
+_C.TEST.RERANK.LAMBDA = 0.3
+
+# Precise batchnorm
+_C.TEST.PRECISE_BN = CN()
+_C.TEST.PRECISE_BN.ENABLED = False
+_C.TEST.PRECISE_BN.DATASET = 'Market1501'
+_C.TEST.PRECISE_BN.NUM_ITER = 300
+
+# ---------------------------------------------------------------------------- #
+# Misc options
+# ---------------------------------------------------------------------------- #
+_C.OUTPUT_DIR = "logs/"
+
+# Benchmark different cudnn algorithms.
+# If input images have very different sizes, this option will have large overhead
+# for about 10k iterations. It usually hurts total time, but can benefit for certain models.
+# If input images have the same or similar sizes, benchmark is often helpful.
+_C.CUDNN_BENCHMARK = False

fastreid/data/__init__.py

@@ -0,0 +1,7 @@
+# encoding: utf-8
+"""
+@author:  sherlock
+@contact: sherlockliao01@gmail.com
+"""
+
+from .build import build_reid_train_loader, build_reid_test_loader

fastreid/data/build.py

@@ -0,0 +1,115 @@
+# encoding: utf-8
+"""
+@author:  l1aoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+import torch
+from torch._six import string_classes
+int_classes = int
+import collections.abc as container_abcs
+from torch.utils.data import DataLoader
+from fastreid.utils import comm
+
+from . import samplers
+from .common import CommDataset
+from .datasets import DATASET_REGISTRY
+from .transforms import build_transforms
+
+_root = os.getenv("FASTREID_DATASETS", "datasets")
+
+
+def build_reid_train_loader(cfg):
+    cfg = cfg.clone()
+    cfg.defrost()
+
+    train_items = list()
+    for d in cfg.DATASETS.NAMES:
+        dataset = DATASET_REGISTRY.get(d)(root=_root, combineall=cfg.DATASETS.COMBINEALL)
+        if comm.is_main_process():
+            dataset.show_train()
+        train_items.extend(dataset.train)
+
+    iters_per_epoch = len(train_items) // cfg.SOLVER.IMS_PER_BATCH
+    cfg.SOLVER.MAX_ITER *= iters_per_epoch
+    train_transforms = build_transforms(cfg, is_train=True)
+    train_set = CommDataset(train_items, train_transforms, relabel=True)
+
+    num_workers = cfg.DATALOADER.NUM_WORKERS
+    num_instance = cfg.DATALOADER.NUM_INSTANCE
+    mini_batch_size = cfg.SOLVER.IMS_PER_BATCH // comm.get_world_size()
+
+    if cfg.DATALOADER.PK_SAMPLER:
+        if cfg.DATALOADER.NAIVE_WAY:
+            data_sampler = samplers.NaiveIdentitySampler(train_set.img_items,
+                                                         cfg.SOLVER.IMS_PER_BATCH, num_instance)
+        else:
+            data_sampler = samplers.BalancedIdentitySampler(train_set.img_items,
+                                                            cfg.SOLVER.IMS_PER_BATCH, num_instance)
+    else:
+        data_sampler = samplers.TrainingSampler(len(train_set))
+    batch_sampler = torch.utils.data.sampler.BatchSampler(data_sampler, mini_batch_size, True)
+
+    train_loader = torch.utils.data.DataLoader(
+        train_set,
+        num_workers=num_workers,
+        batch_sampler=batch_sampler,
+        collate_fn=fast_batch_collator,
+        pin_memory=True,
+    )
+    return train_loader
+
+
+def build_reid_test_loader(cfg, dataset_name):
+    cfg = cfg.clone()
+    cfg.defrost()
+
+    dataset = DATASET_REGISTRY.get(dataset_name)(root=_root)
+    if comm.is_main_process():
+        dataset.show_test()
+    test_items = dataset.query + dataset.gallery
+
+    test_transforms = build_transforms(cfg, is_train=False)
+    test_set = CommDataset(test_items, test_transforms, relabel=False)
+
+    mini_batch_size = cfg.TEST.IMS_PER_BATCH // comm.get_world_size()
+    data_sampler = samplers.InferenceSampler(len(test_set))
+    batch_sampler = torch.utils.data.BatchSampler(data_sampler, mini_batch_size, False)
+    test_loader = DataLoader(
+        test_set,
+        batch_sampler=batch_sampler,
+        num_workers=2,  # save some memory
+        collate_fn=fast_batch_collator,
+        pin_memory=True,
+    )
+    return test_loader, len(dataset.query)
+
+
+def trivial_batch_collator(batch):
+    """
+    A batch collator that does nothing.
+    """
+    return batch
+
+
+def fast_batch_collator(batched_inputs):
+    """
+    A simple batch collator for most common reid tasks
+    """
+    elem = batched_inputs[0]
+    if isinstance(elem, torch.Tensor):
+        out = torch.zeros((len(batched_inputs), *elem.size()), dtype=elem.dtype)
+        for i, tensor in enumerate(batched_inputs):
+            out[i] += tensor
+        return out
+
+    elif isinstance(elem, container_abcs.Mapping):
+        return {key: fast_batch_collator([d[key] for d in batched_inputs]) for key in elem}
+
+    elif isinstance(elem, float):
+        return torch.tensor(batched_inputs, dtype=torch.float64)
+    elif isinstance(elem, int_classes):
+        return torch.tensor(batched_inputs)
+    elif isinstance(elem, string_classes):
+        return batched_inputs

fastreid/data/common.py

@@ -0,0 +1,55 @@
+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+
+from torch.utils.data import Dataset
+
+from .data_utils import read_image
+
+
+class CommDataset(Dataset):
+    """Image Person ReID Dataset"""
+
+    def __init__(self, img_items, transform=None, relabel=True):
+        self.img_items = img_items
+        self.transform = transform
+        self.relabel = relabel
+
+        pid_set = set()
+        cam_set = set()
+        for i in img_items:
+            pid_set.add(i[1])
+            cam_set.add(i[2])
+
+        self.pids = sorted(list(pid_set))
+        self.cams = sorted(list(cam_set))
+        if relabel:
+            self.pid_dict = dict([(p, i) for i, p in enumerate(self.pids)])
+            self.cam_dict = dict([(p, i) for i, p in enumerate(self.cams)])
+
+    def __len__(self):
+        return len(self.img_items)
+
+    def __getitem__(self, index):
+        img_path, pid, camid = self.img_items[index]
+        img = read_image(img_path)
+        if self.transform is not None: img = self.transform(img)
+        if self.relabel:
+            pid = self.pid_dict[pid]
+            camid = self.cam_dict[camid]
+        return {
+            "images": img,
+            "targets": pid,
+            "camids": camid,
+            "img_paths": img_path,
+        }
+
+    @property
+    def num_classes(self):
+        return len(self.pids)
+
+    @property
+    def num_cameras(self):
+        return len(self.cams)

fastreid/data/data_utils.py

@@ -0,0 +1,45 @@
+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+import numpy as np
+from PIL import Image, ImageOps
+
+from fastreid.utils.file_io import PathManager
+
+
+def read_image(file_name, format=None):
+    """
+    Read an image into the given format.
+    Will apply rotation and flipping if the image has such exif information.
+    Args:
+        file_name (str): image file path
+        format (str): one of the supported image modes in PIL, or "BGR"
+    Returns:
+        image (np.ndarray): an HWC image
+    """
+    with PathManager.open(file_name, "rb") as f:
+        image = Image.open(f)
+
+        # capture and ignore this bug: https://github.com/python-pillow/Pillow/issues/3973
+        try:
+            image = ImageOps.exif_transpose(image)
+        except Exception:
+            pass
+
+        if format is not None:
+            # PIL only supports RGB, so convert to RGB and flip channels over below
+            conversion_format = format
+            if format == "BGR":
+                conversion_format = "RGB"
+            image = image.convert(conversion_format)
+        image = np.asarray(image)
+        if format == "BGR":
+            # flip channels if needed
+            image = image[:, :, ::-1]
+        # PIL squeezes out the channel dimension for "L", so make it HWC
+        if format == "L":
+            image = np.expand_dims(image, -1)
+        image = Image.fromarray(image)
+        return image

fastreid/data/datasets/AirportALERT.py

@@ -0,0 +1,46 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['AirportALERT', ]
+
+
+@DATASET_REGISTRY.register()
+class AirportALERT(ImageDataset):
+    dataset_dir = "AirportALERT"
+    dataset_name = "airport"
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+        self.train_file = os.path.join(self.root, self.dataset_dir, 'filepath.txt')
+
+        required_files = [self.train_file, self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path, self.train_file)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, dir_path, train_file):
+        data = []
+        with open(train_file, "r") as f:
+            img_paths = [line.strip('\n') for line in f.readlines()]
+
+        for path in img_paths:
+            split_path = path.split('\\')
+            img_path = '/'.join(split_path)
+            camid = self.dataset_name + "_" + split_path[0]
+            pid = self.dataset_name + "_" + split_path[1]
+            img_path = os.path.join(dir_path, img_path)
+            if 11001 <= int(split_path[1]) <= 401999:
+                data.append([img_path, pid, camid])
+
+        return data

fastreid/data/datasets/__init__.py

@@ -0,0 +1,41 @@
+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+
+from ...utils.registry import Registry
+
+DATASET_REGISTRY = Registry("DATASET")
+DATASET_REGISTRY.__doc__ = """
+Registry for datasets
+It must returns an instance of :class:`Backbone`.
+"""
+
+# Person re-id datasets
+from .cuhk03 import CUHK03
+from .dukemtmcreid import DukeMTMC
+from .market1501 import Market1501
+from .msmt17 import MSMT17
+from .AirportALERT import AirportALERT
+from .iLIDS import iLIDS
+from .pku import PKU
+from .prai import PRAI
+from .sensereid import SenseReID
+from .sysu_mm import SYSU_mm
+from .thermalworld import Thermalworld
+from .pes3d import PeS3D
+from .caviara import CAVIARa
+from .viper import VIPeR
+from .lpw import LPW
+from .shinpuhkan import Shinpuhkan
+from .wildtracker import WildTrackCrop
+from .cuhk_sysu import cuhkSYSU
+
+# Vehicle re-id datasets
+from .veri import VeRi
+from .vehicleid import VehicleID, SmallVehicleID, MediumVehicleID, LargeVehicleID
+from .veriwild import VeRiWild, SmallVeRiWild, MediumVeRiWild, LargeVeRiWild
+
+
+__all__ = [k for k in globals().keys() if "builtin" not in k and not k.startswith("_")]

fastreid/data/datasets/bases.py

@@ -0,0 +1,172 @@
+# encoding: utf-8
+"""
+@author:  sherlock
+@contact: sherlockliao01@gmail.com
+"""
+
+import copy
+import logging
+import os
+# from tabulate import tabulate
+# from termcolor import colored
+
+logger = logging.getLogger(__name__)
+
+
+class Dataset(object):
+    """An abstract class representing a Dataset.
+    This is the base class for ``ImageDataset`` and ``VideoDataset``.
+    Args:
+        train (list): contains tuples of (img_path(s), pid, camid).
+        query (list): contains tuples of (img_path(s), pid, camid).
+        gallery (list): contains tuples of (img_path(s), pid, camid).
+        transform: transform function.
+        mode (str): 'train', 'query' or 'gallery'.
+        combineall (bool): combines train, query and gallery in a
+            dataset for training.
+        verbose (bool): show information.
+    """
+    _junk_pids = []  # contains useless person IDs, e.g. background, false detections
+
+    def __init__(self, train, query, gallery, transform=None, mode='train',
+                 combineall=False, verbose=True, **kwargs):
+        self.train = train
+        self.query = query
+        self.gallery = gallery
+        self.transform = transform
+        self.mode = mode
+        self.combineall = combineall
+        self.verbose = verbose
+
+        self.num_train_pids = self.get_num_pids(self.train)
+        self.num_train_cams = self.get_num_cams(self.train)
+
+        if self.combineall:
+            self.combine_all()
+
+        if self.mode == 'train':
+            self.data = self.train
+        elif self.mode == 'query':
+            self.data = self.query
+        elif self.mode == 'gallery':
+            self.data = self.gallery
+        else:
+            raise ValueError('Invalid mode. Got {}, but expected to be '
+                             'one of [train | query | gallery]'.format(self.mode))
+
+    def __getitem__(self, index):
+        raise NotImplementedError
+
+    def __len__(self):
+        return len(self.data)
+
+    def __radd__(self, other):
+        """Supports sum([dataset1, dataset2, dataset3])."""
+        if other == 0:
+            return self
+        else:
+            return self.__add__(other)
+
+    def parse_data(self, data):
+        """Parses data list and returns the number of person IDs
+        and the number of camera views.
+        Args:
+            data (list): contains tuples of (img_path(s), pid, camid)
+        """
+        pids = set()
+        cams = set()
+        for _, pid, camid in data:
+            pids.add(pid)
+            cams.add(camid)
+        return len(pids), len(cams)
+
+    def get_num_pids(self, data):
+        """Returns the number of training person identities."""
+        return self.parse_data(data)[0]
+
+    def get_num_cams(self, data):
+        """Returns the number of training cameras."""
+        return self.parse_data(data)[1]
+
+    def show_summary(self):
+        """Shows dataset statistics."""
+        pass
+
+    def combine_all(self):
+        """Combines train, query and gallery in a dataset for training."""
+        combined = copy.deepcopy(self.train)
+
+        def _combine_data(data):
+            for img_path, pid, camid in data:
+                if pid in self._junk_pids:
+                    continue
+                pid = self.dataset_name + "_" + str(pid)
+                camid = self.dataset_name + "_" + str(camid)
+                combined.append((img_path, pid, camid))
+
+        _combine_data(self.query)
+        _combine_data(self.gallery)
+
+        self.train = combined
+        self.num_train_pids = self.get_num_pids(self.train)
+
+    def check_before_run(self, required_files):
+        """Checks if required files exist before going deeper.
+        Args:
+            required_files (str or list): string file name(s).
+        """
+        if isinstance(required_files, str):
+            required_files = [required_files]
+
+        for fpath in required_files:
+            if not os.path.exists(fpath):
+                raise RuntimeError('"{}" is not found'.format(fpath))
+
+
+class ImageDataset(Dataset):
+    """A base class representing ImageDataset.
+    All other image datasets should subclass it.
+    ``__getitem__`` returns an image given index.
+    It will return ``img``, ``pid``, ``camid`` and ``img_path``
+    where ``img`` has shape (channel, height, width). As a result,
+    data in each batch has shape (batch_size, channel, height, width).
+    """
+
+    def __init__(self, train, query, gallery, **kwargs):
+        super(ImageDataset, self).__init__(train, query, gallery, **kwargs)
+
+    def show_train(self):
+        num_train_pids, num_train_cams = self.parse_data(self.train)
+
+        headers = ['subset', '# ids', '# images', '# cameras']
+        csv_results = [['train', num_train_pids, len(self.train), num_train_cams]]
+
+        # tabulate it
+        raise Exception("package tabulate is discard")
+        table = tabulate(
+            csv_results,
+            tablefmt="pipe",
+            headers=headers,
+            numalign="left",
+        )
+        logger.info(f"=> Loaded {self.__class__.__name__} in csv format: \n" + colored(table, "cyan"))
+
+    def show_test(self):
+        num_query_pids, num_query_cams = self.parse_data(self.query)
+        num_gallery_pids, num_gallery_cams = self.parse_data(self.gallery)
+
+        headers = ['subset', '# ids', '# images', '# cameras']
+        csv_results = [
+            ['query', num_query_pids, len(self.query), num_query_cams],
+            ['gallery', num_gallery_pids, len(self.gallery), num_gallery_cams],
+        ]
+
+        # tabulate it
+        raise Exception("package tabulate is discard")
+        table = tabulate(
+            csv_results,
+            tablefmt="pipe",
+            headers=headers,
+            numalign="left",
+        )
+        logger.info(f"=> Loaded {self.__class__.__name__} in csv format: \n" + colored(table, "cyan"))

fastreid/data/datasets/caviara.py

@@ -0,0 +1,46 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from scipy.io import loadmat
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+import pdb
+import random
+import numpy as np
+
+__all__ = ['CAVIARa',]
+
+
+@DATASET_REGISTRY.register()
+class CAVIARa(ImageDataset):
+    dataset_dir = "CAVIARa"
+    dataset_name = "caviara"
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+
+        img_list = glob(os.path.join(train_path, "*.jpg"))
+        for img_path in img_list:
+            img_name = img_path.split('/')[-1]
+            pid = self.dataset_name + "_" + img_name[:4]
+            camid = self.dataset_name + "_cam0"
+            data.append([img_path, pid, camid])
+
+        return data

fastreid/data/datasets/cuhk03.py

@@ -0,0 +1,274 @@
+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: liaoxingyu2@jd.com
+"""
+
+import json
+import os.path as osp
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.utils.file_io import PathManager
+from .bases import ImageDataset
+
+
+@DATASET_REGISTRY.register()
+class CUHK03(ImageDataset):
+    """CUHK03.
+
+    Reference:
+        Li et al. DeepReID: Deep Filter Pairing Neural Network for Person Re-identification. CVPR 2014.
+
+    URL: `<http://www.ee.cuhk.edu.hk/~xgwang/CUHK_identification.html#!>`_
+
+    Dataset statistics:
+        - identities: 1360.
+        - images: 13164.
+        - cameras: 6.
+        - splits: 20 (classic).
+    """
+    dataset_dir = 'cuhk03'
+    dataset_url = None
+    dataset_name = "cuhk03"
+
+    def __init__(self, root='datasets', split_id=0, cuhk03_labeled=True, cuhk03_classic_split=False, **kwargs):
+        self.root = root
+        self.dataset_dir = osp.join(self.root, self.dataset_dir)
+
+        self.data_dir = osp.join(self.dataset_dir, 'cuhk03_release')
+        self.raw_mat_path = osp.join(self.data_dir, 'cuhk-03.mat')
+
+        self.imgs_detected_dir = osp.join(self.dataset_dir, 'images_detected')
+        self.imgs_labeled_dir = osp.join(self.dataset_dir, 'images_labeled')
+
+        self.split_classic_det_json_path = osp.join(self.dataset_dir, 'splits_classic_detected.json')
+        self.split_classic_lab_json_path = osp.join(self.dataset_dir, 'splits_classic_labeled.json')
+
+        self.split_new_det_json_path = osp.join(self.dataset_dir, 'splits_new_detected.json')
+        self.split_new_lab_json_path = osp.join(self.dataset_dir, 'splits_new_labeled.json')
+
+        self.split_new_det_mat_path = osp.join(self.dataset_dir, 'cuhk03_new_protocol_config_detected.mat')
+        self.split_new_lab_mat_path = osp.join(self.dataset_dir, 'cuhk03_new_protocol_config_labeled.mat')
+
+        required_files = [
+            self.dataset_dir,
+            self.data_dir,
+            self.raw_mat_path,
+            self.split_new_det_mat_path,
+            self.split_new_lab_mat_path
+        ]
+        self.check_before_run(required_files)
+
+        self.preprocess_split()
+
+        if cuhk03_labeled:
+            split_path = self.split_classic_lab_json_path if cuhk03_classic_split else self.split_new_lab_json_path
+        else:
+            split_path = self.split_classic_det_json_path if cuhk03_classic_split else self.split_new_det_json_path
+
+        with PathManager.open(split_path) as f:
+            splits = json.load(f)
+        assert split_id < len(splits), 'Condition split_id ({}) < len(splits) ({}) is false'.format(split_id,
+                                                                                                    len(splits))
+        split = splits[split_id]
+
+        train = split['train']
+        tmp_train = []
+        for img_path, pid, camid in train:
+            new_pid = self.dataset_name + "_" + str(pid)
+            new_camid = self.dataset_name + "_" + str(camid)
+            tmp_train.append((img_path, new_pid, new_camid))
+        train = tmp_train
+        del tmp_train
+        query = split['query']
+        gallery = split['gallery']
+
+        super(CUHK03, self).__init__(train, query, gallery, **kwargs)
+
+    def preprocess_split(self):
+        # This function is a bit complex and ugly, what it does is
+        # 1. extract data from cuhk-03.mat and save as png images
+        # 2. create 20 classic splits (Li et al. CVPR'14)
+        # 3. create new split (Zhong et al. CVPR'17)
+        if osp.exists(self.imgs_labeled_dir) \
+                and osp.exists(self.imgs_detected_dir) \
+                and osp.exists(self.split_classic_det_json_path) \
+                and osp.exists(self.split_classic_lab_json_path) \
+                and osp.exists(self.split_new_det_json_path) \
+                and osp.exists(self.split_new_lab_json_path):
+            return
+
+        import h5py
+        from imageio import imwrite
+        from scipy.io import loadmat
+
+        PathManager.mkdirs(self.imgs_detected_dir)
+        PathManager.mkdirs(self.imgs_labeled_dir)
+
+        print('Extract image data from "{}" and save as png'.format(self.raw_mat_path))
+        mat = h5py.File(self.raw_mat_path, 'r')
+
+        def _deref(ref):
+            return mat[ref][:].T
+
+        def _process_images(img_refs, campid, pid, save_dir):
+            img_paths = []  # Note: some persons only have images for one view
+            for imgid, img_ref in enumerate(img_refs):
+                img = _deref(img_ref)
+                if img.size == 0 or img.ndim < 3:
+                    continue  # skip empty cell
+                # images are saved with the following format, index-1 (ensure uniqueness)
+                # campid: index of camera pair (1-5)
+                # pid: index of person in 'campid'-th camera pair
+                # viewid: index of view, {1, 2}
+                # imgid: index of image, (1-10)
+                viewid = 1 if imgid < 5 else 2
+                img_name = '{:01d}_{:03d}_{:01d}_{:02d}.png'.format(campid + 1, pid + 1, viewid, imgid + 1)
+                img_path = osp.join(save_dir, img_name)
+                if not osp.isfile(img_path):
+                    imwrite(img_path, img)
+                img_paths.append(img_path)
+            return img_paths
+
+        def _extract_img(image_type):
+            print('Processing {} images ...'.format(image_type))
+            meta_data = []
+            imgs_dir = self.imgs_detected_dir if image_type == 'detected' else self.imgs_labeled_dir
+            for campid, camp_ref in enumerate(mat[image_type][0]):
+                camp = _deref(camp_ref)
+                num_pids = camp.shape[0]
+                for pid in range(num_pids):
+                    img_paths = _process_images(camp[pid, :], campid, pid, imgs_dir)
+                    assert len(img_paths) > 0, 'campid{}-pid{} has no images'.format(campid, pid)
+                    meta_data.append((campid + 1, pid + 1, img_paths))
+                print('- done camera pair {} with {} identities'.format(campid + 1, num_pids))
+            return meta_data
+
+        meta_detected = _extract_img('detected')
+        meta_labeled = _extract_img('labeled')
+
+        def _extract_classic_split(meta_data, test_split):
+            train, test = [], []
+            num_train_pids, num_test_pids = 0, 0
+            num_train_imgs, num_test_imgs = 0, 0
+            for i, (campid, pid, img_paths) in enumerate(meta_data):
+
+                if [campid, pid] in test_split:
+                    for img_path in img_paths:
+                        camid = int(osp.basename(img_path).split('_')[2]) - 1  # make it 0-based
+                        test.append((img_path, num_test_pids, camid))
+                    num_test_pids += 1
+                    num_test_imgs += len(img_paths)
+                else:
+                    for img_path in img_paths:
+                        camid = int(osp.basename(img_path).split('_')[2]) - 1  # make it 0-based
+                        train.append((img_path, num_train_pids, camid))
+                    num_train_pids += 1
+                    num_train_imgs += len(img_paths)
+            return train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs
+
+        print('Creating classic splits (# = 20) ...')
+        splits_classic_det, splits_classic_lab = [], []
+        for split_ref in mat['testsets'][0]:
+            test_split = _deref(split_ref).tolist()
+
+            # create split for detected images
+            train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs = \
+                _extract_classic_split(meta_detected, test_split)
+            splits_classic_det.append({
+                'train': train,
+                'query': test,
+                'gallery': test,
+                'num_train_pids': num_train_pids,
+                'num_train_imgs': num_train_imgs,
+                'num_query_pids': num_test_pids,
+                'num_query_imgs': num_test_imgs,
+                'num_gallery_pids': num_test_pids,
+                'num_gallery_imgs': num_test_imgs
+            })
+
+            # create split for labeled images
+            train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs = \
+                _extract_classic_split(meta_labeled, test_split)
+            splits_classic_lab.append({
+                'train': train,
+                'query': test,
+                'gallery': test,
+                'num_train_pids': num_train_pids,
+                'num_train_imgs': num_train_imgs,
+                'num_query_pids': num_test_pids,
+                'num_query_imgs': num_test_imgs,
+                'num_gallery_pids': num_test_pids,
+                'num_gallery_imgs': num_test_imgs
+            })
+
+        with PathManager.open(self.split_classic_det_json_path, 'w') as f:
+            json.dump(splits_classic_det, f, indent=4, separators=(',', ': '))
+        with PathManager.open(self.split_classic_lab_json_path, 'w') as f:
+            json.dump(splits_classic_lab, f, indent=4, separators=(',', ': '))
+
+        def _extract_set(filelist, pids, pid2label, idxs, img_dir, relabel):
+            tmp_set = []
+            unique_pids = set()
+            for idx in idxs:
+                img_name = filelist[idx][0]
+                camid = int(img_name.split('_')[2]) - 1  # make it 0-based
+                pid = pids[idx]
+                if relabel:
+                    pid = pid2label[pid]
+                img_path = osp.join(img_dir, img_name)
+                tmp_set.append((img_path, int(pid), camid))
+                unique_pids.add(pid)
+            return tmp_set, len(unique_pids), len(idxs)
+
+        def _extract_new_split(split_dict, img_dir):
+            train_idxs = split_dict['train_idx'].flatten() - 1  # index-0
+            pids = split_dict['labels'].flatten()
+            train_pids = set(pids[train_idxs])
+            pid2label = {pid: label for label, pid in enumerate(train_pids)}
+            query_idxs = split_dict['query_idx'].flatten() - 1
+            gallery_idxs = split_dict['gallery_idx'].flatten() - 1
+            filelist = split_dict['filelist'].flatten()
+            train_info = _extract_set(filelist, pids, pid2label, train_idxs, img_dir, relabel=True)
+            query_info = _extract_set(filelist, pids, pid2label, query_idxs, img_dir, relabel=False)
+            gallery_info = _extract_set(filelist, pids, pid2label, gallery_idxs, img_dir, relabel=False)
+            return train_info, query_info, gallery_info
+
+        print('Creating new split for detected images (767/700) ...')
+        train_info, query_info, gallery_info = _extract_new_split(
+            loadmat(self.split_new_det_mat_path),
+            self.imgs_detected_dir
+        )
+        split = [{
+            'train': train_info[0],
+            'query': query_info[0],
+            'gallery': gallery_info[0],
+            'num_train_pids': train_info[1],
+            'num_train_imgs': train_info[2],
+            'num_query_pids': query_info[1],
+            'num_query_imgs': query_info[2],
+            'num_gallery_pids': gallery_info[1],
+            'num_gallery_imgs': gallery_info[2]
+        }]
+
+        with PathManager.open(self.split_new_det_json_path, 'w') as f:
+            json.dump(split, f, indent=4, separators=(',', ': '))
+
+        print('Creating new split for labeled images (767/700) ...')
+        train_info, query_info, gallery_info = _extract_new_split(
+            loadmat(self.split_new_lab_mat_path),
+            self.imgs_labeled_dir
+        )
+        split = [{
+            'train': train_info[0],
+            'query': query_info[0],
+            'gallery': gallery_info[0],
+            'num_train_pids': train_info[1],
+            'num_train_imgs': train_info[2],
+            'num_query_pids': query_info[1],
+            'num_query_imgs': query_info[2],
+            'num_gallery_pids': gallery_info[1],
+            'num_gallery_imgs': gallery_info[2]
+        }]
+        with PathManager.open(self.split_new_lab_json_path, 'w') as f:
+            json.dump(split, f, indent=4, separators=(',', ': '))

fastreid/data/datasets/cuhk_sysu.py

@@ -0,0 +1,58 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import glob
+import os.path as osp
+import re
+import warnings
+
+from .bases import ImageDataset
+from ..datasets import DATASET_REGISTRY
+
+
+@DATASET_REGISTRY.register()
+class cuhkSYSU(ImageDataset):
+    r"""CUHK SYSU datasets.
+
+    The dataset is collected from two sources: street snap and movie.
+    In street snap, 12,490 images and 6,057 query persons were collected
+    with movable cameras across hundreds of scenes while 5,694 images and
+    2,375 query persons were selected from movies and TV dramas.
+
+    Dataset statistics:
+        - identities: xxx.
+        - images: 12936 (train).
+    """
+    dataset_dir = 'cuhk_sysu'
+    dataset_name = "cuhksysu"
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.dataset_dir = osp.join(self.root, self.dataset_dir)
+
+        self.data_dir = osp.join(self.dataset_dir, "cropped_images")
+
+        required_files = [self.data_dir]
+        self.check_before_run(required_files)
+
+        train = self.process_dir(self.data_dir)
+        query = []
+        gallery = []
+
+        super(cuhkSYSU, self).__init__(train, query, gallery, **kwargs)
+
+    def process_dir(self, dir_path):
+        img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
+        pattern = re.compile(r'p([-\d]+)_s(\d)')
+
+        data = []
+        for img_path in img_paths:
+            pid, _ = map(int, pattern.search(img_path).groups())
+            pid = self.dataset_name + "_" + str(pid)
+            camid = self.dataset_name + "_0"
+            data.append((img_path, pid, camid))
+
+        return data

fastreid/data/datasets/dukemtmcreid.py

@@ -0,0 +1,70 @@
+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: liaoxingyu2@jd.com
+"""
+
+import glob
+import os.path as osp
+import re
+
+from .bases import ImageDataset
+from ..datasets import DATASET_REGISTRY
+
+
+@DATASET_REGISTRY.register()
+class DukeMTMC(ImageDataset):
+    """DukeMTMC-reID.
+
+    Reference:
+        - Ristani et al. Performance Measures and a Data Set for Multi-Target, Multi-Camera Tracking. ECCVW 2016.
+        - Zheng et al. Unlabeled Samples Generated by GAN Improve the Person Re-identification Baseline in vitro. ICCV 2017.
+
+    URL: `<https://github.com/layumi/DukeMTMC-reID_evaluation>`_
+
+    Dataset statistics:
+        - identities: 1404 (train + query).
+        - images:16522 (train) + 2228 (query) + 17661 (gallery).
+        - cameras: 8.
+    """
+    dataset_dir = 'DukeMTMC-reID'
+    dataset_url = 'http://vision.cs.duke.edu/DukeMTMC/data/misc/DukeMTMC-reID.zip'
+    dataset_name = "dukemtmc"
+
+    def __init__(self, root='datasets', **kwargs):
+        # self.root = osp.abspath(osp.expanduser(root))
+        self.root = root
+        self.dataset_dir = osp.join(self.root, self.dataset_dir)
+        self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train')
+        self.query_dir = osp.join(self.dataset_dir, 'query')
+        self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test')
+
+        required_files = [
+            self.dataset_dir,
+            self.train_dir,
+            self.query_dir,
+            self.gallery_dir,
+        ]
+        self.check_before_run(required_files)
+
+        train = self.process_dir(self.train_dir)
+        query = self.process_dir(self.query_dir, is_train=False)
+        gallery = self.process_dir(self.gallery_dir, is_train=False)
+
+        super(DukeMTMC, self).__init__(train, query, gallery, **kwargs)
+
+    def process_dir(self, dir_path, is_train=True):
+        img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
+        pattern = re.compile(r'([-\d]+)_c(\d)')
+
+        data = []
+        for img_path in img_paths:
+            pid, camid = map(int, pattern.search(img_path).groups())
+            assert 1 <= camid <= 8
+            camid -= 1  # index starts from 0
+            if is_train:
+                pid = self.dataset_name + "_" + str(pid)
+                camid = self.dataset_name + "_" + str(camid)
+            data.append((img_path, pid, camid))
+
+        return data

fastreid/data/datasets/iLIDS.py

@@ -0,0 +1,43 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['iLIDS', ]
+
+
+@DATASET_REGISTRY.register()
+class iLIDS(ImageDataset):
+    dataset_dir = "iLIDS"
+    dataset_name = "ilids"
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+        file_path = os.listdir(train_path)
+        for pid_dir in file_path:
+            img_file = os.path.join(train_path, pid_dir)
+            img_paths = glob(os.path.join(img_file, "*.png"))
+            for img_path in img_paths:
+                split_path = img_path.split('/')
+                pid = self.dataset_name + "_" + split_path[-2]
+                camid = self.dataset_name + "_" + split_path[-1].split('_')[0]
+                data.append([img_path, pid, camid])
+        return data

fastreid/data/datasets/lpw.py

@@ -0,0 +1,47 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['LPW', ]
+
+
+@DATASET_REGISTRY.register()
+class LPW(ImageDataset):
+    dataset_dir = "pep_256x128"
+    dataset_name = "lpw"
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+
+        file_path_list = ['scen1', 'scen2', 'scen3']
+
+        for scene in file_path_list:
+            cam_list = os.listdir(os.path.join(train_path, scene))
+            for cam in cam_list:
+                camid = self.dataset_name + "_" + cam
+                pid_list = os.listdir(os.path.join(train_path, scene, cam))
+                for pid_dir in pid_list:
+                    img_paths = glob(os.path.join(train_path, scene, cam, pid_dir, "*.jpg"))
+                    for img_path in img_paths:
+                        pid = self.dataset_name + "_" + scene + "-" + pid_dir
+                        data.append([img_path, pid, camid])
+        return data

fastreid/data/datasets/market1501.py

@@ -0,0 +1,90 @@
+# encoding: utf-8
+"""
+@author:  sherlock
+@contact: sherlockliao01@gmail.com
+"""
+
+import glob
+import os.path as osp
+import re
+import warnings
+
+from .bases import ImageDataset
+from ..datasets import DATASET_REGISTRY
+
+
+@DATASET_REGISTRY.register()
+class Market1501(ImageDataset):
+    """Market1501.
+
+    Reference:
+        Zheng et al. Scalable Person Re-identification: A Benchmark. ICCV 2015.
+
+    URL: `<http://www.liangzheng.org/Project/project_reid.html>`_
+
+    Dataset statistics:
+        - identities: 1501 (+1 for background).
+        - images: 12936 (train) + 3368 (query) + 15913 (gallery).
+    """
+    _junk_pids = [0, -1]
+    dataset_dir = ''
+    dataset_url = 'http://188.138.127.15:81/Datasets/Market-1501-v15.09.15.zip'
+    dataset_name = "market1501"
+
+    def __init__(self, root='datasets', market1501_500k=False, **kwargs):
+        # self.root = osp.abspath(osp.expanduser(root))
+        self.root = root
+        self.dataset_dir = osp.join(self.root, self.dataset_dir)
+
+        # allow alternative directory structure
+        self.data_dir = self.dataset_dir
+        data_dir = osp.join(self.data_dir, 'Market-1501-v15.09.15')
+        if osp.isdir(data_dir):
+            self.data_dir = data_dir
+        else:
+            warnings.warn('The current data structure is deprecated. Please '
+                          'put data folders such as "bounding_box_train" under '
+                          '"Market-1501-v15.09.15".')
+
+        self.train_dir = osp.join(self.data_dir, 'bounding_box_train')
+        self.query_dir = osp.join(self.data_dir, 'query')
+        self.gallery_dir = osp.join(self.data_dir, 'bounding_box_test')
+        self.extra_gallery_dir = osp.join(self.data_dir, 'images')
+        self.market1501_500k = market1501_500k
+
+        required_files = [
+            self.data_dir,
+            self.train_dir,
+            self.query_dir,
+            self.gallery_dir,
+        ]
+        if self.market1501_500k:
+            required_files.append(self.extra_gallery_dir)
+        self.check_before_run(required_files)
+
+        train = self.process_dir(self.train_dir)
+        query = self.process_dir(self.query_dir, is_train=False)
+        gallery = self.process_dir(self.gallery_dir, is_train=False)
+        if self.market1501_500k:
+            gallery += self.process_dir(self.extra_gallery_dir, is_train=False)
+
+        super(Market1501, self).__init__(train, query, gallery, **kwargs)
+
+    def process_dir(self, dir_path, is_train=True):
+        img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
+        pattern = re.compile(r'([-\d]+)_c(\d)')
+
+        data = []
+        for img_path in img_paths:
+            pid, camid = map(int, pattern.search(img_path).groups())
+            if pid == -1:
+                continue  # junk images are just ignored
+            assert 0 <= pid <= 1501  # pid == 0 means background
+            assert 1 <= camid <= 6
+            camid -= 1  # index starts from 0
+            if is_train:
+                pid = self.dataset_name + "_" + str(pid)
+                camid = self.dataset_name + "_" + str(camid)
+            data.append((img_path, pid, camid))
+
+        return data

fastreid/data/datasets/msmt17.py

@@ -0,0 +1,114 @@
+# encoding: utf-8
+"""
+@author:  l1aoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+
+import sys
+import os
+import os.path as osp
+
+from .bases import ImageDataset
+from ..datasets import DATASET_REGISTRY
+##### Log #####
+# 22.01.2019
+# - add v2
+# - v1 and v2 differ in dir names
+# - note that faces in v2 are blurred
+TRAIN_DIR_KEY = 'train_dir'
+TEST_DIR_KEY = 'test_dir'
+VERSION_DICT = {
+    'MSMT17_V1': {
+        TRAIN_DIR_KEY: 'train',
+        TEST_DIR_KEY: 'test',
+    },
+    'MSMT17_V2': {
+        TRAIN_DIR_KEY: 'mask_train_v2',
+        TEST_DIR_KEY: 'mask_test_v2',
+    }
+}
+
+
+@DATASET_REGISTRY.register()
+class MSMT17(ImageDataset):
+    """MSMT17.
+    Reference:
+        Wei et al. Person Transfer GAN to Bridge Domain Gap for Person Re-Identification. CVPR 2018.
+    URL: `<http://www.pkuvmc.com/publications/msmt17.html>`_
+
+    Dataset statistics:
+        - identities: 4101.
+        - images: 32621 (train) + 11659 (query) + 82161 (gallery).
+        - cameras: 15.
+    """
+    # dataset_dir = 'MSMT17_V2'
+    dataset_url = None
+    dataset_name = 'msmt17'
+
+    def __init__(self, root='datasets', **kwargs):
+        self.dataset_dir = root
+
+        has_main_dir = False
+        for main_dir in VERSION_DICT:
+            if osp.exists(osp.join(self.dataset_dir, main_dir)):
+                train_dir = VERSION_DICT[main_dir][TRAIN_DIR_KEY]
+                test_dir = VERSION_DICT[main_dir][TEST_DIR_KEY]
+                has_main_dir = True
+                break
+        assert has_main_dir, 'Dataset folder not found'
+
+        self.train_dir = osp.join(self.dataset_dir, main_dir, train_dir)
+        self.test_dir = osp.join(self.dataset_dir, main_dir, test_dir)
+        self.list_train_path = osp.join(self.dataset_dir, main_dir, 'list_train.txt')
+        self.list_val_path = osp.join(self.dataset_dir, main_dir, 'list_val.txt')
+        self.list_query_path = osp.join(self.dataset_dir, main_dir, 'list_query.txt')
+        self.list_gallery_path = osp.join(self.dataset_dir, main_dir, 'list_gallery.txt')
+
+        required_files = [
+            self.dataset_dir,
+            self.train_dir,
+            self.test_dir
+        ]
+        self.check_before_run(required_files)
+
+        train = self.process_dir(self.train_dir, self.list_train_path)
+        val = self.process_dir(self.train_dir, self.list_val_path)
+        query = self.process_dir(self.test_dir, self.list_query_path, is_train=False)
+        gallery = self.process_dir(self.test_dir, self.list_gallery_path, is_train=False)
+
+        num_train_pids = self.get_num_pids(train)
+        query_tmp = []
+        for img_path, pid, camid in query:
+            query_tmp.append((img_path, pid+num_train_pids, camid))
+        del query
+        query = query_tmp
+
+        gallery_temp = []
+        for img_path, pid, camid in gallery:
+            gallery_temp.append((img_path, pid+num_train_pids, camid))
+        del gallery
+        gallery = gallery_temp
+
+        # Note: to fairly compare with published methods on the conventional ReID setting,
+        #       do not add val images to the training set.
+        if 'combineall' in kwargs and kwargs['combineall']:
+            train += val
+        super(MSMT17, self).__init__(train, query, gallery, **kwargs)
+
+    def process_dir(self, dir_path, list_path, is_train=True):
+        with open(list_path, 'r') as txt:
+            lines = txt.readlines()
+
+        data = []
+
+        for img_idx, img_info in enumerate(lines):
+            img_path, pid = img_info.split(' ')
+            pid = int(pid)  # no need to relabel
+            camid = int(img_path.split('_')[2]) - 1  # index starts from 0
+            img_path = osp.join(dir_path, img_path)
+            if is_train:
+                pid = self.dataset_name + "_" + str(pid)
+                camid = self.dataset_name + "_" + str(camid)
+            data.append((img_path, pid, camid))
+
+        return data

fastreid/data/datasets/pes3d.py

@@ -0,0 +1,46 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from scipy.io import loadmat
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+import pdb
+import random
+import numpy as np
+
+__all__ = ['PeS3D',]
+
+
+@DATASET_REGISTRY.register()
+class PeS3D(ImageDataset):
+    dataset_dir = "3DPeS"
+    dataset_name = "pes3d"
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+
+        pid_list = os.listdir(train_path)
+        for pid_dir in pid_list:
+            pid = self.dataset_name + "_" + pid_dir
+            img_list = glob(os.path.join(train_path, pid_dir,  "*.bmp"))
+            for img_path in img_list:
+                camid = self.dataset_name + "_cam0"
+                data.append([img_path, pid, camid])
+        return data

fastreid/data/datasets/pku.py

@@ -0,0 +1,42 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['PKU', ]
+
+
+@DATASET_REGISTRY.register()
+class PKU(ImageDataset):
+    dataset_dir = "PKUv1a_128x48"
+    dataset_name = 'pku'
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+        img_paths = glob(os.path.join(train_path, "*.png"))
+
+        for img_path in img_paths:
+            split_path = img_path.split('/')
+            img_info = split_path[-1].split('_')
+            pid = self.dataset_name + "_" + img_info[0]
+            camid = self.dataset_name + "_" + img_info[1]
+            data.append([img_path, pid, camid])
+        return data

fastreid/data/datasets/prai.py

@@ -0,0 +1,44 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from scipy.io import loadmat
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+import pdb
+
+__all__ = ['PRAI',]
+
+
+@DATASET_REGISTRY.register()
+class PRAI(ImageDataset):
+    dataset_dir = "PRAI-1581"
+    dataset_name = 'prai'
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir, 'images')
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+        img_paths = glob(os.path.join(train_path, "*.jpg"))
+        for img_path in img_paths:
+            split_path = img_path.split('/')
+            img_info = split_path[-1].split('_')
+            pid = self.dataset_name + "_" + img_info[0]
+            camid = self.dataset_name + "_" + img_info[1]
+            data.append([img_path, pid, camid])
+        return data
+

fastreid/data/datasets/sensereid.py

@@ -0,0 +1,45 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['SenseReID', ]
+
+
+@DATASET_REGISTRY.register()
+class SenseReID(ImageDataset):
+    dataset_dir = "SenseReID"
+    dataset_name = "senseid"
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+        file_path_list = ['test_gallery', 'test_prob']
+
+        for file_path in file_path_list:
+            sub_file = os.path.join(train_path, file_path)
+            img_name = glob(os.path.join(sub_file, "*.jpg"))
+            for img_path in img_name:
+                img_name = img_path.split('/')[-1]
+                img_info = img_name.split('_')
+                pid = self.dataset_name + "_" + img_info[0]
+                camid = self.dataset_name + "_" + img_info[1].split('.')[0]
+                data.append([img_path, pid, camid])
+        return data

fastreid/data/datasets/shinpuhkan.py

@@ -0,0 +1,46 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['Shinpuhkan', ]
+
+
+@DATASET_REGISTRY.register()
+class Shinpuhkan(ImageDataset):
+    dataset_dir = "shinpuhkan"
+    dataset_name = 'shinpuhkan'
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+
+        for root, dirs, files in os.walk(train_path):
+            img_names = list(filter(lambda x: x.endswith(".jpg"), files))
+            # fmt: off
+            if len(img_names) == 0: continue
+            # fmt: on
+            for img_name in img_names:
+                img_path = os.path.join(root, img_name)
+                split_path = img_name.split('_')
+                pid = self.dataset_name + "_" + split_path[0]
+                camid = self.dataset_name + "_" + split_path[2]
+                data.append((img_path, pid, camid))
+
+        return data

fastreid/data/datasets/sysu_mm.py

@@ -0,0 +1,48 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from scipy.io import loadmat
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+import pdb
+
+__all__ = ['SYSU_mm', ]
+
+
+@DATASET_REGISTRY.register()
+class SYSU_mm(ImageDataset):
+    dataset_dir = "SYSU-MM01"
+    dataset_name = "sysumm01"
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+
+        file_path_list = ['cam1', 'cam2', 'cam4', 'cam5']
+        
+        for file_path in file_path_list:
+            camid = self.dataset_name + "_" + file_path
+            pid_list = os.listdir(os.path.join(train_path, file_path))
+            for pid_dir in pid_list:
+                pid = self.dataset_name + "_" + pid_dir
+                img_list = glob(os.path.join(train_path, file_path, pid_dir, "*.jpg"))
+                for img_path in img_list:
+                    data.append([img_path, pid, camid])
+        return data
+

fastreid/data/datasets/thermalworld.py

@@ -0,0 +1,45 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from scipy.io import loadmat
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+import pdb
+import random
+import numpy as np
+
+__all__ = ['Thermalworld',]
+
+
+@DATASET_REGISTRY.register()
+class Thermalworld(ImageDataset):
+    dataset_dir = "thermalworld_rgb"
+    dataset_name = "thermalworld"
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+        pid_list = os.listdir(train_path)
+        for pid_dir in pid_list:
+            pid = self.dataset_name + "_" + pid_dir
+            img_list = glob(os.path.join(train_path, pid_dir,  "*.jpg"))
+            for img_path in img_list:
+                camid = self.dataset_name + "_cam0"
+                data.append([img_path, pid, camid])
+        return data

fastreid/data/datasets/vehicleid.py

@@ -0,0 +1,124 @@
+# encoding: utf-8
+"""
+@author:  Jinkai Zheng
+@contact: 1315673509@qq.com
+"""
+
+import os.path as osp
+import random
+
+from .bases import ImageDataset
+from ..datasets import DATASET_REGISTRY
+
+
+@DATASET_REGISTRY.register()
+class VehicleID(ImageDataset):
+    """VehicleID.
+
+    Reference:
+        Liu et al. Deep relative distance learning: Tell the difference between similar vehicles. CVPR 2016.
+
+    URL: `<https://pkuml.org/resources/pku-vehicleid.html>`_
+
+    Train dataset statistics:
+        - identities: 13164.
+        - images: 113346.
+    """
+    dataset_dir = "vehicleid"
+    dataset_name = "vehicleid"
+
+    def __init__(self, root='datasets', test_list='', **kwargs):
+        self.dataset_dir = osp.join(root, self.dataset_dir)
+
+        self.image_dir = osp.join(self.dataset_dir, 'image')
+        self.train_list = osp.join(self.dataset_dir, 'train_test_split/train_list.txt')
+        if test_list:
+            self.test_list = test_list
+        else:
+            self.test_list = osp.join(self.dataset_dir, 'train_test_split/test_list_13164.txt')
+
+        required_files = [
+            self.dataset_dir,
+            self.image_dir,
+            self.train_list,
+            self.test_list,
+        ]
+        self.check_before_run(required_files)
+
+        train = self.process_dir(self.train_list, is_train=True)
+        query, gallery = self.process_dir(self.test_list, is_train=False)
+
+        super(VehicleID, self).__init__(train, query, gallery, **kwargs)
+
+    def process_dir(self, list_file, is_train=True):
+        img_list_lines = open(list_file, 'r').readlines()
+
+        dataset = []
+        for idx, line in enumerate(img_list_lines):
+            line = line.strip()
+            vid = int(line.split(' ')[1])
+            imgid = line.split(' ')[0]
+            img_path = osp.join(self.image_dir, imgid + '.jpg')
+            if is_train:
+                vid = self.dataset_name + "_" + str(vid)
+            dataset.append((img_path, vid, int(imgid)))
+
+        if is_train: return dataset
+        else:
+            random.shuffle(dataset)
+            vid_container = set()
+            query = []
+            gallery = []
+            for sample in dataset:
+                if sample[1] not in vid_container:
+                    vid_container.add(sample[1])
+                    gallery.append(sample)
+                else:
+                    query.append(sample)
+
+            return query, gallery
+
+
+@DATASET_REGISTRY.register()
+class SmallVehicleID(VehicleID):
+    """VehicleID.
+    Small test dataset statistics:
+        - identities: 800.
+        - images: 6493.
+    """
+
+    def __init__(self, root='datasets', **kwargs):
+        dataset_dir = osp.join(root, self.dataset_dir)
+        self.test_list = osp.join(dataset_dir, 'train_test_split/test_list_800.txt')
+
+        super(SmallVehicleID, self).__init__(root, self.test_list, **kwargs)
+
+
+@DATASET_REGISTRY.register()
+class MediumVehicleID(VehicleID):
+    """VehicleID.
+    Medium test dataset statistics:
+        - identities: 1600.
+        - images: 13377.
+    """
+
+    def __init__(self, root='datasets', **kwargs):
+        dataset_dir = osp.join(root, self.dataset_dir)
+        self.test_list = osp.join(dataset_dir, 'train_test_split/test_list_1600.txt')
+
+        super(MediumVehicleID, self).__init__(root, self.test_list, **kwargs)
+
+
+@DATASET_REGISTRY.register()
+class LargeVehicleID(VehicleID):
+    """VehicleID.
+    Large test dataset statistics:
+        - identities: 2400.
+        - images: 19777.
+    """
+
+    def __init__(self, root='datasets', **kwargs):
+        dataset_dir = osp.join(root, self.dataset_dir)
+        self.test_list = osp.join(dataset_dir, 'train_test_split/test_list_2400.txt')
+
+        super(LargeVehicleID, self).__init__(root, self.test_list, **kwargs)

fastreid/data/datasets/veri.py

@@ -0,0 +1,68 @@
+# encoding: utf-8
+"""
+@author:  Jinkai Zheng
+@contact: 1315673509@qq.com
+"""
+
+import glob
+import os.path as osp
+import re
+
+from .bases import ImageDataset
+from ..datasets import DATASET_REGISTRY
+
+
+@DATASET_REGISTRY.register()
+class VeRi(ImageDataset):
+    """VeRi.
+
+    Reference:
+        Liu et al. A Deep Learning based Approach for Progressive Vehicle Re-Identification. ECCV 2016.
+
+    URL: `<https://vehiclereid.github.io/VeRi/>`_
+
+    Dataset statistics:
+        - identities: 775.
+        - images: 37778 (train) + 1678 (query) + 11579 (gallery).
+    """
+    dataset_dir = "veri"
+    dataset_name = "veri"
+
+    def __init__(self, root='datasets', **kwargs):
+        self.dataset_dir = osp.join(root, self.dataset_dir)
+
+        self.train_dir = osp.join(self.dataset_dir, 'image_train')
+        self.query_dir = osp.join(self.dataset_dir, 'image_query')
+        self.gallery_dir = osp.join(self.dataset_dir, 'image_test')
+
+        required_files = [
+            self.dataset_dir,
+            self.train_dir,
+            self.query_dir,
+            self.gallery_dir,
+        ]
+        self.check_before_run(required_files)
+
+        train = self.process_dir(self.train_dir)
+        query = self.process_dir(self.query_dir, is_train=False)
+        gallery = self.process_dir(self.gallery_dir, is_train=False)
+
+        super(VeRi, self).__init__(train, query, gallery, **kwargs)
+
+    def process_dir(self, dir_path, is_train=True):
+        img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
+        pattern = re.compile(r'([\d]+)_c(\d\d\d)')
+
+        data = []
+        for img_path in img_paths:
+            pid, camid = map(int, pattern.search(img_path).groups())
+            if pid == -1: continue  # junk images are just ignored
+            assert 1 <= pid <= 776
+            assert 1 <= camid <= 20
+            camid -= 1  # index starts from 0
+            if is_train:
+                pid = self.dataset_name + "_" + str(pid)
+                camid = self.dataset_name + "_" + str(camid)
+            data.append((img_path, pid, camid))
+
+        return data

fastreid/data/datasets/veriwild.py

@@ -0,0 +1,138 @@
+# encoding: utf-8
+"""
+@author:  Jinkai Zheng
+@contact: 1315673509@qq.com
+"""
+
+import os.path as osp
+
+from .bases import ImageDataset
+from ..datasets import DATASET_REGISTRY
+
+
+@DATASET_REGISTRY.register()
+class VeRiWild(ImageDataset):
+    """VeRi-Wild.
+
+    Reference:
+        Lou et al. A Large-Scale Dataset for Vehicle Re-Identification in the Wild. CVPR 2019.
+
+    URL: `<https://github.com/PKU-IMRE/VERI-Wild>`_
+
+    Train dataset statistics:
+        - identities: 30671.
+        - images: 277797.
+    """
+    dataset_dir = "VERI-Wild"
+    dataset_name = "veriwild"
+
+    def __init__(self, root='datasets', query_list='', gallery_list='', **kwargs):
+        self.dataset_dir = osp.join(root, self.dataset_dir)
+
+        self.image_dir = osp.join(self.dataset_dir, 'images')
+        self.train_list = osp.join(self.dataset_dir, 'train_test_split/train_list.txt')
+        self.vehicle_info = osp.join(self.dataset_dir, 'train_test_split/vehicle_info.txt')
+        if query_list and gallery_list:
+            self.query_list = query_list
+            self.gallery_list = gallery_list
+        else:
+            self.query_list = osp.join(self.dataset_dir, 'train_test_split/test_10000_query.txt')
+            self.gallery_list = osp.join(self.dataset_dir, 'train_test_split/test_10000.txt')
+
+        required_files = [
+            self.image_dir,
+            self.train_list,
+            self.query_list,
+            self.gallery_list,
+            self.vehicle_info,
+        ]
+        self.check_before_run(required_files)
+
+        self.imgid2vid, self.imgid2camid, self.imgid2imgpath = self.process_vehicle(self.vehicle_info)
+
+        train = self.process_dir(self.train_list)
+        query = self.process_dir(self.query_list, is_train=False)
+        gallery = self.process_dir(self.gallery_list, is_train=False)
+
+        super(VeRiWild, self).__init__(train, query, gallery, **kwargs)
+
+    def process_dir(self, img_list, is_train=True):
+        img_list_lines = open(img_list, 'r').readlines()
+
+        dataset = []
+        for idx, line in enumerate(img_list_lines):
+            line = line.strip()
+            vid = int(line.split('/')[0])
+            imgid = line.split('/')[1]
+            if is_train:
+                vid = self.dataset_name + "_" + str(vid)
+            dataset.append((self.imgid2imgpath[imgid], vid, int(self.imgid2camid[imgid])))
+
+        assert len(dataset) == len(img_list_lines)
+        return dataset
+
+    def process_vehicle(self, vehicle_info):
+        imgid2vid = {}
+        imgid2camid = {}
+        imgid2imgpath = {}
+        vehicle_info_lines = open(vehicle_info, 'r').readlines()
+
+        for idx, line in enumerate(vehicle_info_lines[1:]):
+            vid = line.strip().split('/')[0]
+            imgid = line.strip().split(';')[0].split('/')[1]
+            camid = line.strip().split(';')[1]
+            img_path = osp.join(self.image_dir, vid, imgid + '.jpg')
+            imgid2vid[imgid] = vid
+            imgid2camid[imgid] = camid
+            imgid2imgpath[imgid] = img_path
+
+        assert len(imgid2vid) == len(vehicle_info_lines) - 1
+        return imgid2vid, imgid2camid, imgid2imgpath
+
+
+@DATASET_REGISTRY.register()
+class SmallVeRiWild(VeRiWild):
+    """VeRi-Wild.
+    Small test dataset statistics:
+        - identities: 3000.
+        - images: 41861.
+    """
+
+    def __init__(self, root='datasets', **kwargs):
+        dataset_dir = osp.join(root, self.dataset_dir)
+        self.query_list = osp.join(dataset_dir, 'train_test_split/test_3000_query.txt')
+        self.gallery_list = osp.join(dataset_dir, 'train_test_split/test_3000.txt')
+
+        super(SmallVeRiWild, self).__init__(root, self.query_list, self.gallery_list, **kwargs)
+
+
+@DATASET_REGISTRY.register()
+class MediumVeRiWild(VeRiWild):
+    """VeRi-Wild.
+    Medium test dataset statistics:
+        - identities: 5000.
+        - images: 69389.
+    """
+
+    def __init__(self, root='datasets', **kwargs):
+        dataset_dir = osp.join(root, self.dataset_dir)
+        self.query_list = osp.join(dataset_dir, 'train_test_split/test_5000_query.txt')
+        self.gallery_list = osp.join(dataset_dir, 'train_test_split/test_5000.txt')
+
+        super(MediumVeRiWild, self).__init__(root, self.query_list, self.gallery_list, **kwargs)
+
+
+@DATASET_REGISTRY.register()
+class LargeVeRiWild(VeRiWild):
+    """VeRi-Wild.
+    Large test dataset statistics:
+        - identities: 10000.
+        - images: 138517.
+    """
+
+    def __init__(self, root='datasets', **kwargs):
+        dataset_dir = osp.join(root, self.dataset_dir)
+        self.query_list = osp.join(dataset_dir, 'train_test_split/test_10000_query.txt')
+        self.gallery_list = osp.join(dataset_dir, 'train_test_split/test_10000.txt')
+
+        super(LargeVeRiWild, self).__init__(root, self.query_list, self.gallery_list, **kwargs)

fastreid/data/datasets/viper.py

@@ -0,0 +1,45 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['VIPeR', ]
+
+
+@DATASET_REGISTRY.register()
+class VIPeR(ImageDataset):
+    dataset_dir = "VIPeR"
+    dataset_name = "viper"
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+
+        file_path_list = ['cam_a', 'cam_b']
+
+        for file_path in file_path_list:
+            camid = self.dataset_name + "_" + file_path
+            img_list = glob(os.path.join(train_path, file_path, "*.bmp"))
+            for img_path in img_list:
+                img_name = img_path.split('/')[-1]
+                pid = self.dataset_name + "_" + img_name.split('_')[0]
+                data.append([img_path, pid, camid])
+
+        return data

fastreid/data/datasets/wildtracker.py

@@ -0,0 +1,59 @@
+# encoding: utf-8
+"""
+@author:  wangguanan
+@contact: guan.wang0706@gmail.com
+"""
+
+import glob
+import os
+
+from .bases import ImageDataset
+from ..datasets import DATASET_REGISTRY
+
+
+@DATASET_REGISTRY.register()
+class WildTrackCrop(ImageDataset):
+    """WildTrack.
+    Reference:
+        WILDTRACK: A Multi-camera HD Dataset for Dense Unscripted Pedestrian Detection
+            T. Chavdarova; P. Baqué; A. Maksai; S. Bouquet; C. Jose et al.
+    URL: `<https://www.epfl.ch/labs/cvlab/data/data-wildtrack/>`_
+    Dataset statistics:
+        - identities: 313
+        - images: 33979 (train only)
+        - cameras: 7
+    Args:
+        data_path(str): path to WildTrackCrop dataset
+        combineall(bool): combine train and test sets as train set if True
+    """
+    dataset_url = None
+    dataset_dir = 'Wildtrack_crop_dataset'
+    dataset_name = 'wildtrack'
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.dataset_dir = os.path.join(self.root, self.dataset_dir)
+
+        self.train_dir = os.path.join(self.dataset_dir, "crop")
+
+        train = self.process_dir(self.train_dir)
+        query = []
+        gallery = []
+
+        super(WildTrackCrop, self).__init__(train, query, gallery, **kwargs)
+
+    def process_dir(self, dir_path):
+        r"""
+        :param dir_path: directory path saving images
+        Returns
+            data(list) = [img_path, pid, camid]
+        """
+        data = []
+        for dir_name in os.listdir(dir_path):
+            img_lists = glob.glob(os.path.join(dir_path, dir_name, "*.png"))
+            for img_path in img_lists:
+                pid = self.dataset_name + "_" + dir_name
+                camid = img_path.split('/')[-1].split('_')[0]
+                camid = self.dataset_name + "_" + camid
+                data.append([img_path, pid, camid])
+        return data

fastreid/data/samplers/__init__.py

@@ -0,0 +1,8 @@
+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+
+from .triplet_sampler import BalancedIdentitySampler, NaiveIdentitySampler
+from .data_sampler import TrainingSampler, InferenceSampler

fastreid/data/samplers/data_sampler.py

@@ -0,0 +1,85 @@
+# encoding: utf-8
+"""
+@author:  l1aoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+import itertools
+from typing import Optional
+
+import numpy as np
+from torch.utils.data import Sampler
+
+from fastreid.utils import comm
+
+
+class TrainingSampler(Sampler):
+    """
+    In training, we only care about the "infinite stream" of training data.
+    So this sampler produces an infinite stream of indices and
+    all workers cooperate to correctly shuffle the indices and sample different indices.
+    The samplers in each worker effectively produces `indices[worker_id::num_workers]`
+    where `indices` is an infinite stream of indices consisting of
+    `shuffle(range(size)) + shuffle(range(size)) + ...` (if shuffle is True)
+    or `range(size) + range(size) + ...` (if shuffle is False)
+    """
+
+    def __init__(self, size: int, shuffle: bool = True, seed: Optional[int] = None):
+        """
+        Args:
+            size (int): the total number of data of the underlying dataset to sample from
+            shuffle (bool): whether to shuffle the indices or not
+            seed (int): the initial seed of the shuffle. Must be the same
+                across all workers. If None, will use a random seed shared
+                among workers (require synchronization among all workers).
+        """
+        self._size = size
+        assert size > 0
+        self._shuffle = shuffle
+        if seed is None:
+            seed = comm.shared_random_seed()
+        self._seed = int(seed)
+
+        self._rank = comm.get_rank()
+        self._world_size = comm.get_world_size()
+
+    def __iter__(self):
+        start = self._rank
+        yield from itertools.islice(self._infinite_indices(), start, None, self._world_size)
+
+    def _infinite_indices(self):
+        np.random.seed(self._seed)
+        while True:
+            if self._shuffle:
+                yield from np.random.permutation(self._size)
+            else:
+                yield from np.arange(self._size)
+
+
+class InferenceSampler(Sampler):
+    """
+    Produce indices for inference.
+    Inference needs to run on the __exact__ set of samples,
+    therefore when the total number of samples is not divisible by the number of workers,
+    this sampler produces different number of samples on different workers.
+    """
+
+    def __init__(self, size: int):
+        """
+        Args:
+            size (int): the total number of data of the underlying dataset to sample from
+        """
+        self._size = size
+        assert size > 0
+        self._rank = comm.get_rank()
+        self._world_size = comm.get_world_size()
+
+        shard_size = (self._size - 1) // self._world_size + 1
+        begin = shard_size * self._rank
+        end = min(shard_size * (self._rank + 1), self._size)
+        self._local_indices = range(begin, end)
+
+    def __iter__(self):
+        yield from self._local_indices
+
+    def __len__(self):
+        return len(self._local_indices)

fastreid/data/samplers/triplet_sampler.py

@@ -0,0 +1,169 @@
+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: liaoxingyu2@jd.com
+"""
+
+import copy
+import itertools
+from collections import defaultdict
+from typing import Optional
+
+import numpy as np
+from torch.utils.data.sampler import Sampler
+
+from fastreid.utils import comm
+
+
+def no_index(a, b):
+    assert isinstance(a, list)
+    return [i for i, j in enumerate(a) if j != b]
+
+
+class BalancedIdentitySampler(Sampler):
+    def __init__(self, data_source: str, batch_size: int, num_instances: int, seed: Optional[int] = None):
+        self.data_source = data_source
+        self.batch_size = batch_size
+        self.num_instances = num_instances
+        self.num_pids_per_batch = batch_size // self.num_instances
+
+        self.index_pid = defaultdict(list)
+        self.pid_cam = defaultdict(list)
+        self.pid_index = defaultdict(list)
+
+        for index, info in enumerate(data_source):
+            pid = info[1]
+            camid = info[2]
+            self.index_pid[index] = pid
+            self.pid_cam[pid].append(camid)
+            self.pid_index[pid].append(index)
+
+        self.pids = sorted(list(self.pid_index.keys()))
+        self.num_identities = len(self.pids)
+
+        if seed is None:
+            seed = comm.shared_random_seed()
+        self._seed = int(seed)
+
+        self._rank = comm.get_rank()
+        self._world_size = comm.get_world_size()
+
+    def __iter__(self):
+        start = self._rank
+        yield from itertools.islice(self._infinite_indices(), start, None, self._world_size)
+
+    def _infinite_indices(self):
+        np.random.seed(self._seed)
+        while True:
+            # Shuffle identity list
+            identities = np.random.permutation(self.num_identities)
+
+            # If remaining identities cannot be enough for a batch,
+            # just drop the remaining parts
+            drop_indices = self.num_identities % self.num_pids_per_batch
+            if drop_indices: identities = identities[:-drop_indices]
+
+            ret = []
+            for kid in identities:
+                i = np.random.choice(self.pid_index[self.pids[kid]])
+                _, i_pid, i_cam = self.data_source[i]
+                ret.append(i)
+                pid_i = self.index_pid[i]
+                cams = self.pid_cam[pid_i]
+                index = self.pid_index[pid_i]
+                select_cams = no_index(cams, i_cam)
+
+                if select_cams:
+                    if len(select_cams) >= self.num_instances:
+                        cam_indexes = np.random.choice(select_cams, size=self.num_instances - 1, replace=False)
+                    else:
+                        cam_indexes = np.random.choice(select_cams, size=self.num_instances - 1, replace=True)
+                    for kk in cam_indexes:
+                        ret.append(index[kk])
+                else:
+                    select_indexes = no_index(index, i)
+                    if not select_indexes:
+                        # Only one image for this identity
+                        ind_indexes = [0] * (self.num_instances - 1)
+                    elif len(select_indexes) >= self.num_instances:
+                        ind_indexes = np.random.choice(select_indexes, size=self.num_instances - 1, replace=False)
+                    else:
+                        ind_indexes = np.random.choice(select_indexes, size=self.num_instances - 1, replace=True)
+
+                    for kk in ind_indexes:
+                        ret.append(index[kk])
+
+                if len(ret) == self.batch_size:
+                    yield from ret
+                    ret = []
+
+
+class NaiveIdentitySampler(Sampler):
+    """
+    Randomly sample N identities, then for each identity,
+    randomly sample K instances, therefore batch size is N*K.
+    Args:
+    - data_source (list): list of (img_path, pid, camid).
+    - num_instances (int): number of instances per identity in a batch.
+    - batch_size (int): number of examples in a batch.
+    """
+
+    def __init__(self, data_source: str, batch_size: int, num_instances: int, seed: Optional[int] = None):
+        self.data_source = data_source
+        self.batch_size = batch_size
+        self.num_instances = num_instances
+        self.num_pids_per_batch = batch_size // self.num_instances
+
+        self.index_pid = defaultdict(list)
+        self.pid_cam = defaultdict(list)
+        self.pid_index = defaultdict(list)
+
+        for index, info in enumerate(data_source):
+            pid = info[1]
+            camid = info[2]
+            self.index_pid[index] = pid
+            self.pid_cam[pid].append(camid)
+            self.pid_index[pid].append(index)
+
+        self.pids = sorted(list(self.pid_index.keys()))
+        self.num_identities = len(self.pids)
+
+        if seed is None:
+            seed = comm.shared_random_seed()
+        self._seed = int(seed)
+
+        self._rank = comm.get_rank()
+        self._world_size = comm.get_world_size()
+
+    def __iter__(self):
+        start = self._rank
+        yield from itertools.islice(self._infinite_indices(), start, None, self._world_size)
+
+    def _infinite_indices(self):
+        np.random.seed(self._seed)
+        while True:
+            avai_pids = copy.deepcopy(self.pids)
+            batch_idxs_dict = {}
+
+            batch_indices = []
+            while len(avai_pids) >= self.num_pids_per_batch:
+                selected_pids = np.random.choice(avai_pids, self.num_pids_per_batch, replace=False).tolist()
+                for pid in selected_pids:
+                    # Register pid in batch_idxs_dict if not
+                    if pid not in batch_idxs_dict:
+                        idxs = copy.deepcopy(self.pid_index[pid])
+                        if len(idxs) < self.num_instances:
+                            idxs = np.random.choice(idxs, size=self.num_instances, replace=True).tolist()
+                        np.random.shuffle(idxs)
+                        batch_idxs_dict[pid] = idxs
+
+                    avai_idxs = batch_idxs_dict[pid]
+                    for _ in range(self.num_instances):
+                        batch_indices.append(avai_idxs.pop(0))
+
+                    if len(avai_idxs) < self.num_instances: avai_pids.remove(pid)
+
+                assert len(batch_indices) == self.batch_size, f"batch indices have wrong " \
+                                                              f"length with {len(batch_indices)}!"
+                yield from batch_indices
+                batch_indices = []

fastreid/data/transforms/__init__.py

@@ -0,0 +1,10 @@
+# encoding: utf-8
+"""
+@author:  sherlock
+@contact: sherlockliao01@gmail.com
+"""
+
+
+from .build import build_transforms
+from .transforms import *
+from .autoaugment import *

fastreid/data/transforms/autoaugment.py

@@ -0,0 +1,812 @@
+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+
+""" AutoAugment, RandAugment, and AugMix for PyTorch
+This code implements the searched ImageNet policies with various tweaks and improvements and
+does not include any of the search code.
+AA and RA Implementation adapted from:
+    https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/autoaugment.py
+AugMix adapted from:
+    https://github.com/google-research/augmix
+Papers:
+    AutoAugment: Learning Augmentation Policies from Data - https://arxiv.org/abs/1805.09501
+    Learning Data Augmentation Strategies for Object Detection - https://arxiv.org/abs/1906.11172
+    RandAugment: Practical automated data augmentation... - https://arxiv.org/abs/1909.13719
+    AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty - https://arxiv.org/abs/1912.02781
+Hacked together by Ross Wightman
+"""
+import math
+import random
+import re
+
+import PIL
+import numpy as np
+from PIL import Image, ImageOps, ImageEnhance
+
+_PIL_VER = tuple([int(x) for x in PIL.__version__.split('.')[:2]])
+
+_FILL = (128, 128, 128)
+
+# This signifies the max integer that the controller RNN could predict for the
+# augmentation scheme.
+_MAX_LEVEL = 10.
+
+_HPARAMS_DEFAULT = dict(
+    translate_const=57,
+    img_mean=_FILL,
+)
+
+_RANDOM_INTERPOLATION = (Image.BILINEAR, Image.BICUBIC)
+
+
+def _interpolation(kwargs):
+    interpolation = kwargs.pop('resample', Image.BILINEAR)
+    if isinstance(interpolation, (list, tuple)):
+        return random.choice(interpolation)
+    else:
+        return interpolation
+
+
+def _check_args_tf(kwargs):
+    if 'fillcolor' in kwargs and _PIL_VER < (5, 0):
+        kwargs.pop('fillcolor')
+    kwargs['resample'] = _interpolation(kwargs)
+
+
+def shear_x(img, factor, **kwargs):
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, factor, 0, 0, 1, 0), **kwargs)
+
+
+def shear_y(img, factor, **kwargs):
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, 0, factor, 1, 0), **kwargs)
+
+
+def translate_x_rel(img, pct, **kwargs):
+    pixels = pct * img.size[0]
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs)
+
+
+def translate_y_rel(img, pct, **kwargs):
+    pixels = pct * img.size[1]
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs)
+
+
+def translate_x_abs(img, pixels, **kwargs):
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs)
+
+
+def translate_y_abs(img, pixels, **kwargs):
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs)
+
+
+def rotate(img, degrees, **kwargs):
+    _check_args_tf(kwargs)
+    if _PIL_VER >= (5, 2):
+        return img.rotate(degrees, **kwargs)
+    elif _PIL_VER >= (5, 0):
+        w, h = img.size
+        post_trans = (0, 0)
+        rotn_center = (w / 2.0, h / 2.0)
+        angle = -math.radians(degrees)
+        matrix = [
+            round(math.cos(angle), 15),
+            round(math.sin(angle), 15),
+            0.0,
+            round(-math.sin(angle), 15),
+            round(math.cos(angle), 15),
+            0.0,
+        ]
+
+        def transform(x, y, matrix):
+            (a, b, c, d, e, f) = matrix
+            return a * x + b * y + c, d * x + e * y + f
+
+        matrix[2], matrix[5] = transform(
+            -rotn_center[0] - post_trans[0], -rotn_center[1] - post_trans[1], matrix
+        )
+        matrix[2] += rotn_center[0]
+        matrix[5] += rotn_center[1]
+        return img.transform(img.size, Image.AFFINE, matrix, **kwargs)
+    else:
+        return img.rotate(degrees, resample=kwargs['resample'])
+
+
+def auto_contrast(img, **__):
+    return ImageOps.autocontrast(img)
+
+
+def invert(img, **__):
+    return ImageOps.invert(img)
+
+
+def equalize(img, **__):
+    return ImageOps.equalize(img)
+
+
+def solarize(img, thresh, **__):
+    return ImageOps.solarize(img, thresh)
+
+
+def solarize_add(img, add, thresh=128, **__):
+    lut = []
+    for i in range(256):
+        if i < thresh:
+            lut.append(min(255, i + add))
+        else:
+            lut.append(i)
+    if img.mode in ("L", "RGB"):
+        if img.mode == "RGB" and len(lut) == 256:
+            lut = lut + lut + lut
+        return img.point(lut)
+    else:
+        return img
+
+
+def posterize(img, bits_to_keep, **__):
+    if bits_to_keep >= 8:
+        return img
+    return ImageOps.posterize(img, bits_to_keep)
+
+
+def contrast(img, factor, **__):
+    return ImageEnhance.Contrast(img).enhance(factor)
+
+
+def color(img, factor, **__):
+    return ImageEnhance.Color(img).enhance(factor)
+
+
+def brightness(img, factor, **__):
+    return ImageEnhance.Brightness(img).enhance(factor)
+
+
+def sharpness(img, factor, **__):
+    return ImageEnhance.Sharpness(img).enhance(factor)
+
+
+def _randomly_negate(v):
+    """With 50% prob, negate the value"""
+    return -v if random.random() > 0.5 else v
+
+
+def _rotate_level_to_arg(level, _hparams):
+    # range [-30, 30]
+    level = (level / _MAX_LEVEL) * 30.
+    level = _randomly_negate(level)
+    return level,
+
+
+def _enhance_level_to_arg(level, _hparams):
+    # range [0.1, 1.9]
+    return (level / _MAX_LEVEL) * 1.8 + 0.1,
+
+
+def _enhance_increasing_level_to_arg(level, _hparams):
+    # the 'no change' level is 1.0, moving away from that towards 0. or 2.0 increases the enhancement blend
+    # range [0.1, 1.9]
+    level = (level / _MAX_LEVEL) * .9
+    level = 1.0 + _randomly_negate(level)
+    return level,
+
+
+def _shear_level_to_arg(level, _hparams):
+    # range [-0.3, 0.3]
+    level = (level / _MAX_LEVEL) * 0.3
+    level = _randomly_negate(level)
+    return level,
+
+
+def _translate_abs_level_to_arg(level, hparams):
+    translate_const = hparams['translate_const']
+    level = (level / _MAX_LEVEL) * float(translate_const)
+    level = _randomly_negate(level)
+    return level,
+
+
+def _translate_rel_level_to_arg(level, hparams):
+    # default range [-0.45, 0.45]
+    translate_pct = hparams.get('translate_pct', 0.45)
+    level = (level / _MAX_LEVEL) * translate_pct
+    level = _randomly_negate(level)
+    return level,
+
+
+def _posterize_level_to_arg(level, _hparams):
+    # As per Tensorflow TPU EfficientNet impl
+    # range [0, 4], 'keep 0 up to 4 MSB of original image'
+    # intensity/severity of augmentation decreases with level
+    return int((level / _MAX_LEVEL) * 4),
+
+
+def _posterize_increasing_level_to_arg(level, hparams):
+    # As per Tensorflow models research and UDA impl
+    # range [4, 0], 'keep 4 down to 0 MSB of original image',
+    # intensity/severity of augmentation increases with level
+    return 4 - _posterize_level_to_arg(level, hparams)[0],
+
+
+def _posterize_original_level_to_arg(level, _hparams):
+    # As per original AutoAugment paper description
+    # range [4, 8], 'keep 4 up to 8 MSB of image'
+    # intensity/severity of augmentation decreases with level
+    return int((level / _MAX_LEVEL) * 4) + 4,
+
+
+def _solarize_level_to_arg(level, _hparams):
+    # range [0, 256]
+    # intensity/severity of augmentation decreases with level
+    return int((level / _MAX_LEVEL) * 256),
+
+
+def _solarize_increasing_level_to_arg(level, _hparams):
+    # range [0, 256]
+    # intensity/severity of augmentation increases with level
+    return 256 - _solarize_level_to_arg(level, _hparams)[0],
+
+
+def _solarize_add_level_to_arg(level, _hparams):
+    # range [0, 110]
+    return int((level / _MAX_LEVEL) * 110),
+
+
+LEVEL_TO_ARG = {
+    'AutoContrast': None,
+    'Equalize': None,
+    'Invert': None,
+    'Rotate': _rotate_level_to_arg,
+    # There are several variations of the posterize level scaling in various Tensorflow/Google repositories/papers
+    'Posterize': _posterize_level_to_arg,
+    'PosterizeIncreasing': _posterize_increasing_level_to_arg,
+    'PosterizeOriginal': _posterize_original_level_to_arg,
+    'Solarize': _solarize_level_to_arg,
+    'SolarizeIncreasing': _solarize_increasing_level_to_arg,
+    'SolarizeAdd': _solarize_add_level_to_arg,
+    'Color': _enhance_level_to_arg,
+    'ColorIncreasing': _enhance_increasing_level_to_arg,
+    'Contrast': _enhance_level_to_arg,
+    'ContrastIncreasing': _enhance_increasing_level_to_arg,
+    'Brightness': _enhance_level_to_arg,
+    'BrightnessIncreasing': _enhance_increasing_level_to_arg,
+    'Sharpness': _enhance_level_to_arg,
+    'SharpnessIncreasing': _enhance_increasing_level_to_arg,
+    'ShearX': _shear_level_to_arg,
+    'ShearY': _shear_level_to_arg,
+    'TranslateX': _translate_abs_level_to_arg,
+    'TranslateY': _translate_abs_level_to_arg,
+    'TranslateXRel': _translate_rel_level_to_arg,
+    'TranslateYRel': _translate_rel_level_to_arg,
+}
+
+NAME_TO_OP = {
+    'AutoContrast': auto_contrast,
+    'Equalize': equalize,
+    'Invert': invert,
+    'Rotate': rotate,
+    'Posterize': posterize,
+    'PosterizeIncreasing': posterize,
+    'PosterizeOriginal': posterize,
+    'Solarize': solarize,
+    'SolarizeIncreasing': solarize,
+    'SolarizeAdd': solarize_add,
+    'Color': color,
+    'ColorIncreasing': color,
+    'Contrast': contrast,
+    'ContrastIncreasing': contrast,
+    'Brightness': brightness,
+    'BrightnessIncreasing': brightness,
+    'Sharpness': sharpness,
+    'SharpnessIncreasing': sharpness,
+    'ShearX': shear_x,
+    'ShearY': shear_y,
+    'TranslateX': translate_x_abs,
+    'TranslateY': translate_y_abs,
+    'TranslateXRel': translate_x_rel,
+    'TranslateYRel': translate_y_rel,
+}
+
+
+class AugmentOp:
+
+    def __init__(self, name, prob=0.5, magnitude=10, hparams=None):
+        hparams = hparams or _HPARAMS_DEFAULT
+        self.aug_fn = NAME_TO_OP[name]
+        self.level_fn = LEVEL_TO_ARG[name]
+        self.prob = prob
+        self.magnitude = magnitude
+        self.hparams = hparams.copy()
+        self.kwargs = dict(
+            fillcolor=hparams['img_mean'] if 'img_mean' in hparams else _FILL,
+            resample=hparams['interpolation'] if 'interpolation' in hparams else _RANDOM_INTERPOLATION,
+        )
+
+        # If magnitude_std is > 0, we introduce some randomness
+        # in the usually fixed policy and sample magnitude from a normal distribution
+        # with mean `magnitude` and std-dev of `magnitude_std`.
+        # NOTE This is my own hack, being tested, not in papers or reference impls.
+        self.magnitude_std = self.hparams.get('magnitude_std', 0)
+
+    def __call__(self, img):
+        if self.prob < 1.0 and random.random() > self.prob:
+            return img
+        magnitude = self.magnitude
+        if self.magnitude_std and self.magnitude_std > 0:
+            magnitude = random.gauss(magnitude, self.magnitude_std)
+        magnitude = min(_MAX_LEVEL, max(0, magnitude))  # clip to valid range
+        level_args = self.level_fn(magnitude, self.hparams) if self.level_fn is not None else tuple()
+        return self.aug_fn(img, *level_args, **self.kwargs)
+
+
+def auto_augment_policy_v0(hparams):
+    # ImageNet v0 policy from TPU EfficientNet impl, cannot find a paper reference.
+    policy = [
+        [('Equalize', 0.8, 1), ('ShearY', 0.8, 4)],
+        [('Color', 0.4, 9), ('Equalize', 0.6, 3)],
+        [('Color', 0.4, 1), ('Rotate', 0.6, 8)],
+        [('Solarize', 0.8, 3), ('Equalize', 0.4, 7)],
+        [('Solarize', 0.4, 2), ('Solarize', 0.6, 2)],
+        [('Color', 0.2, 0), ('Equalize', 0.8, 8)],
+        [('Equalize', 0.4, 8), ('SolarizeAdd', 0.8, 3)],
+        [('ShearX', 0.2, 9), ('Rotate', 0.6, 8)],
+        [('Color', 0.6, 1), ('Equalize', 1.0, 2)],
+        [('Invert', 0.4, 9), ('Rotate', 0.6, 0)],
+        [('Equalize', 1.0, 9), ('ShearY', 0.6, 3)],
+        [('Color', 0.4, 7), ('Equalize', 0.6, 0)],
+        [('Posterize', 0.4, 6), ('AutoContrast', 0.4, 7)],
+        [('Solarize', 0.6, 8), ('Color', 0.6, 9)],
+        [('Solarize', 0.2, 4), ('Rotate', 0.8, 9)],
+        [('Rotate', 1.0, 7), ('TranslateYRel', 0.8, 9)],
+        [('ShearX', 0.0, 0), ('Solarize', 0.8, 4)],
+        [('ShearY', 0.8, 0), ('Color', 0.6, 4)],
+        [('Color', 1.0, 0), ('Rotate', 0.6, 2)],
+        [('Equalize', 0.8, 4), ('Equalize', 0.0, 8)],
+        [('Equalize', 1.0, 4), ('AutoContrast', 0.6, 2)],
+        [('ShearY', 0.4, 7), ('SolarizeAdd', 0.6, 7)],
+        [('Posterize', 0.8, 2), ('Solarize', 0.6, 10)],  # This results in black image with Tpu posterize
+        [('Solarize', 0.6, 8), ('Equalize', 0.6, 1)],
+        [('Color', 0.8, 6), ('Rotate', 0.4, 5)],
+    ]
+    pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
+    return pc
+
+
+def auto_augment_policy_v0r(hparams):
+    # ImageNet v0 policy from TPU EfficientNet impl, with variation of Posterize used
+    # in Google research implementation (number of bits discarded increases with magnitude)
+    policy = [
+        [('Equalize', 0.8, 1), ('ShearY', 0.8, 4)],
+        [('Color', 0.4, 9), ('Equalize', 0.6, 3)],
+        [('Color', 0.4, 1), ('Rotate', 0.6, 8)],
+        [('Solarize', 0.8, 3), ('Equalize', 0.4, 7)],
+        [('Solarize', 0.4, 2), ('Solarize', 0.6, 2)],
+        [('Color', 0.2, 0), ('Equalize', 0.8, 8)],
+        [('Equalize', 0.4, 8), ('SolarizeAdd', 0.8, 3)],
+        [('ShearX', 0.2, 9), ('Rotate', 0.6, 8)],
+        [('Color', 0.6, 1), ('Equalize', 1.0, 2)],
+        [('Invert', 0.4, 9), ('Rotate', 0.6, 0)],
+        [('Equalize', 1.0, 9), ('ShearY', 0.6, 3)],
+        [('Color', 0.4, 7), ('Equalize', 0.6, 0)],
+        [('PosterizeIncreasing', 0.4, 6), ('AutoContrast', 0.4, 7)],
+        [('Solarize', 0.6, 8), ('Color', 0.6, 9)],
+        [('Solarize', 0.2, 4), ('Rotate', 0.8, 9)],
+        [('Rotate', 1.0, 7), ('TranslateYRel', 0.8, 9)],
+        [('ShearX', 0.0, 0), ('Solarize', 0.8, 4)],
+        [('ShearY', 0.8, 0), ('Color', 0.6, 4)],
+        [('Color', 1.0, 0), ('Rotate', 0.6, 2)],
+        [('Equalize', 0.8, 4), ('Equalize', 0.0, 8)],
+        [('Equalize', 1.0, 4), ('AutoContrast', 0.6, 2)],
+        [('ShearY', 0.4, 7), ('SolarizeAdd', 0.6, 7)],
+        [('PosterizeIncreasing', 0.8, 2), ('Solarize', 0.6, 10)],
+        [('Solarize', 0.6, 8), ('Equalize', 0.6, 1)],
+        [('Color', 0.8, 6), ('Rotate', 0.4, 5)],
+    ]
+    pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
+    return pc
+
+
+def auto_augment_policy_original(hparams):
+    # ImageNet policy from https://arxiv.org/abs/1805.09501
+    policy = [
+        [('PosterizeOriginal', 0.4, 8), ('Rotate', 0.6, 9)],
+        [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
+        [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
+        [('PosterizeOriginal', 0.6, 7), ('PosterizeOriginal', 0.6, 6)],
+        [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
+        [('Equalize', 0.4, 4), ('Rotate', 0.8, 8)],
+        [('Solarize', 0.6, 3), ('Equalize', 0.6, 7)],
+        [('PosterizeOriginal', 0.8, 5), ('Equalize', 1.0, 2)],
+        [('Rotate', 0.2, 3), ('Solarize', 0.6, 8)],
+        [('Equalize', 0.6, 8), ('PosterizeOriginal', 0.4, 6)],
+        [('Rotate', 0.8, 8), ('Color', 0.4, 0)],
+        [('Rotate', 0.4, 9), ('Equalize', 0.6, 2)],
+        [('Equalize', 0.0, 7), ('Equalize', 0.8, 8)],
+        [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
+        [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
+        [('Rotate', 0.8, 8), ('Color', 1.0, 2)],
+        [('Color', 0.8, 8), ('Solarize', 0.8, 7)],
+        [('Sharpness', 0.4, 7), ('Invert', 0.6, 8)],
+        [('ShearX', 0.6, 5), ('Equalize', 1.0, 9)],
+        [('Color', 0.4, 0), ('Equalize', 0.6, 3)],
+        [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
+        [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
+        [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
+        [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
+        [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
+    ]
+    pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
+    return pc
+
+
+def auto_augment_policy_originalr(hparams):
+    # ImageNet policy from https://arxiv.org/abs/1805.09501 with research posterize variation
+    policy = [
+        [('PosterizeIncreasing', 0.4, 8), ('Rotate', 0.6, 9)],
+        [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
+        [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
+        [('PosterizeIncreasing', 0.6, 7), ('PosterizeIncreasing', 0.6, 6)],
+        [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
+        [('Equalize', 0.4, 4), ('Rotate', 0.8, 8)],
+        [('Solarize', 0.6, 3), ('Equalize', 0.6, 7)],
+        [('PosterizeIncreasing', 0.8, 5), ('Equalize', 1.0, 2)],
+        [('Rotate', 0.2, 3), ('Solarize', 0.6, 8)],
+        [('Equalize', 0.6, 8), ('PosterizeIncreasing', 0.4, 6)],
+        [('Rotate', 0.8, 8), ('Color', 0.4, 0)],
+        [('Rotate', 0.4, 9), ('Equalize', 0.6, 2)],
+        [('Equalize', 0.0, 7), ('Equalize', 0.8, 8)],
+        [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
+        [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
+        [('Rotate', 0.8, 8), ('Color', 1.0, 2)],
+        [('Color', 0.8, 8), ('Solarize', 0.8, 7)],
+        [('Sharpness', 0.4, 7), ('Invert', 0.6, 8)],
+        [('ShearX', 0.6, 5), ('Equalize', 1.0, 9)],
+        [('Color', 0.4, 0), ('Equalize', 0.6, 3)],
+        [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
+        [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
+        [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
+        [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
+        [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
+    ]
+    pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
+    return pc
+
+
+def auto_augment_policy(name="original"):
+    hparams = _HPARAMS_DEFAULT
+    if name == 'original':
+        return auto_augment_policy_original(hparams)
+    elif name == 'originalr':
+        return auto_augment_policy_originalr(hparams)
+    elif name == 'v0':
+        return auto_augment_policy_v0(hparams)
+    elif name == 'v0r':
+        return auto_augment_policy_v0r(hparams)
+    else:
+        assert False, 'Unknown AA policy (%s)' % name
+
+
+class AutoAugment:
+
+    def __init__(self, total_iter):
+        self.total_iter = total_iter
+        self.gamma = 0
+        self.policy = auto_augment_policy()
+
+    def __call__(self, img):
+        if random.uniform(0, 1) > self.gamma:
+            sub_policy = random.choice(self.policy)
+            self.gamma = min(1.0, self.gamma + 1.0 / self.total_iter)
+            for op in sub_policy:
+                img = op(img)
+            return img
+        else:
+            return img
+
+
+def auto_augment_transform(config_str, hparams):
+    """
+    Create a AutoAugment transform
+    :param config_str: String defining configuration of auto augmentation. Consists of multiple sections separated by
+    dashes ('-'). The first section defines the AutoAugment policy (one of 'v0', 'v0r', 'original', 'originalr').
+    The remaining sections, not order sepecific determine
+        'mstd' -  float std deviation of magnitude noise applied
+    Ex 'original-mstd0.5' results in AutoAugment with original policy, magnitude_std 0.5
+    :param hparams: Other hparams (kwargs) for the AutoAugmentation scheme
+    :return: A PyTorch compatible Transform
+    """
+    config = config_str.split('-')
+    policy_name = config[0]
+    config = config[1:]
+    for c in config:
+        cs = re.split(r'(\d.*)', c)
+        if len(cs) < 2:
+            continue
+        key, val = cs[:2]
+        if key == 'mstd':
+            # noise param injected via hparams for now
+            hparams.setdefault('magnitude_std', float(val))
+        else:
+            assert False, 'Unknown AutoAugment config section'
+    aa_policy = auto_augment_policy(policy_name)
+    return AutoAugment(aa_policy)
+
+
+_RAND_TRANSFORMS = [
+    'AutoContrast',
+    'Equalize',
+    'Invert',
+    'Rotate',
+    'Posterize',
+    'Solarize',
+    'SolarizeAdd',
+    'Color',
+    'Contrast',
+    'Brightness',
+    'Sharpness',
+    'ShearX',
+    'ShearY',
+    'TranslateXRel',
+    'TranslateYRel',
+    # 'Cutout'  # NOTE I've implement this as random erasing separately
+]
+
+_RAND_INCREASING_TRANSFORMS = [
+    'AutoContrast',
+    'Equalize',
+    'Invert',
+    'Rotate',
+    'PosterizeIncreasing',
+    'SolarizeIncreasing',
+    'SolarizeAdd',
+    'ColorIncreasing',
+    'ContrastIncreasing',
+    'BrightnessIncreasing',
+    'SharpnessIncreasing',
+    'ShearX',
+    'ShearY',
+    'TranslateXRel',
+    'TranslateYRel',
+    # 'Cutout'  # NOTE I've implement this as random erasing separately
+]
+
+# These experimental weights are based loosely on the relative improvements mentioned in paper.
+# They may not result in increased performance, but could likely be tuned to so.
+_RAND_CHOICE_WEIGHTS_0 = {
+    'Rotate': 0.3,
+    'ShearX': 0.2,
+    'ShearY': 0.2,
+    'TranslateXRel': 0.1,
+    'TranslateYRel': 0.1,
+    'Color': .025,
+    'Sharpness': 0.025,
+    'AutoContrast': 0.025,
+    'Solarize': .005,
+    'SolarizeAdd': .005,
+    'Contrast': .005,
+    'Brightness': .005,
+    'Equalize': .005,
+    'Posterize': 0,
+    'Invert': 0,
+}
+
+
+def _select_rand_weights(weight_idx=0, transforms=None):
+    transforms = transforms or _RAND_TRANSFORMS
+    assert weight_idx == 0  # only one set of weights currently
+    rand_weights = _RAND_CHOICE_WEIGHTS_0
+    probs = [rand_weights[k] for k in transforms]
+    probs /= np.sum(probs)
+    return probs
+
+
+def rand_augment_ops(magnitude=10, hparams=None, transforms=None):
+    hparams = hparams or _HPARAMS_DEFAULT
+    transforms = transforms or _RAND_TRANSFORMS
+    return [AugmentOp(
+        name, prob=0.5, magnitude=magnitude, hparams=hparams) for name in transforms]
+
+
+class RandAugment:
+    def __init__(self, ops, num_layers=2, choice_weights=None):
+        self.ops = ops
+        self.num_layers = num_layers
+        self.choice_weights = choice_weights
+
+    def __call__(self, img):
+        # no replacement when using weighted choice
+        ops = np.random.choice(
+            self.ops, self.num_layers, replace=self.choice_weights is None, p=self.choice_weights)
+        for op in ops:
+            img = op(img)
+        return img
+
+
+def rand_augment_transform(config_str, hparams):
+    """
+    Create a RandAugment transform
+    :param config_str: String defining configuration of random augmentation. Consists of multiple sections separated by
+    dashes ('-'). The first section defines the specific variant of rand augment (currently only 'rand'). The remaining
+    sections, not order sepecific determine
+        'm' - integer magnitude of rand augment
+        'n' - integer num layers (number of transform ops selected per image)
+        'w' - integer probabiliy weight index (index of a set of weights to influence choice of op)
+        'mstd' -  float std deviation of magnitude noise applied
+        'inc' - integer (bool), use augmentations that increase in severity with magnitude (default: 0)
+    Ex 'rand-m9-n3-mstd0.5' results in RandAugment with magnitude 9, num_layers 3, magnitude_std 0.5
+    'rand-mstd1-w0' results in magnitude_std 1.0, weights 0, default magnitude of 10 and num_layers 2
+    :param hparams: Other hparams (kwargs) for the RandAugmentation scheme
+    :return: A PyTorch compatible Transform
+    """
+    magnitude = _MAX_LEVEL  # default to _MAX_LEVEL for magnitude (currently 10)
+    num_layers = 2  # default to 2 ops per image
+    weight_idx = None  # default to no probability weights for op choice
+    transforms = _RAND_TRANSFORMS
+    config = config_str.split('-')
+    assert config[0] == 'rand'
+    config = config[1:]
+    for c in config:
+        cs = re.split(r'(\d.*)', c)
+        if len(cs) < 2:
+            continue
+        key, val = cs[:2]
+        if key == 'mstd':
+            # noise param injected via hparams for now
+            hparams.setdefault('magnitude_std', float(val))
+        elif key == 'inc':
+            if bool(val):
+                transforms = _RAND_INCREASING_TRANSFORMS
+        elif key == 'm':
+            magnitude = int(val)
+        elif key == 'n':
+            num_layers = int(val)
+        elif key == 'w':
+            weight_idx = int(val)
+        else:
+            assert False, 'Unknown RandAugment config section'
+    ra_ops = rand_augment_ops(magnitude=magnitude, hparams=hparams, transforms=transforms)
+    choice_weights = None if weight_idx is None else _select_rand_weights(weight_idx)
+    return RandAugment(ra_ops, num_layers, choice_weights=choice_weights)
+
+
+_AUGMIX_TRANSFORMS = [
+    'AutoContrast',
+    'ColorIncreasing',  # not in paper
+    'ContrastIncreasing',  # not in paper
+    'BrightnessIncreasing',  # not in paper
+    'SharpnessIncreasing',  # not in paper
+    'Equalize',
+    'Rotate',
+    'PosterizeIncreasing',
+    'SolarizeIncreasing',
+    'ShearX',
+    'ShearY',
+    'TranslateXRel',
+    'TranslateYRel',
+]
+
+
+def augmix_ops(magnitude=10, hparams=None, transforms=None):
+    hparams = hparams or _HPARAMS_DEFAULT
+    transforms = transforms or _AUGMIX_TRANSFORMS
+    return [AugmentOp(
+        name, prob=1.0, magnitude=magnitude, hparams=hparams) for name in transforms]
+
+
+class AugMixAugment:
+    """ AugMix Transform
+    Adapted and improved from impl here: https://github.com/google-research/augmix/blob/master/imagenet.py
+    From paper: 'AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty -
+    https://arxiv.org/abs/1912.02781
+    """
+
+    def __init__(self, ops, alpha=1., width=3, depth=-1, blended=False):
+        self.ops = ops
+        self.alpha = alpha
+        self.width = width
+        self.depth = depth
+        self.blended = blended  # blended mode is faster but not well tested
+
+    def _calc_blended_weights(self, ws, m):
+        ws = ws * m
+        cump = 1.
+        rws = []
+        for w in ws[::-1]:
+            alpha = w / cump
+            cump *= (1 - alpha)
+            rws.append(alpha)
+        return np.array(rws[::-1], dtype=np.float32)
+
+    def _apply_blended(self, img, mixing_weights, m):
+        # This is my first crack and implementing a slightly faster mixed augmentation. Instead
+        # of accumulating the mix for each chain in a Numpy array and then blending with original,
+        # it recomputes the blending coefficients and applies one PIL image blend per chain.
+        # TODO the results appear in the right ballpark but they differ by more than rounding.
+        img_orig = img.copy()
+        ws = self._calc_blended_weights(mixing_weights, m)
+        for w in ws:
+            depth = self.depth if self.depth > 0 else np.random.randint(1, 4)
+            ops = np.random.choice(self.ops, depth, replace=True)
+            img_aug = img_orig  # no ops are in-place, deep copy not necessary
+            for op in ops:
+                img_aug = op(img_aug)
+            img = Image.blend(img, img_aug, w)
+        return img
+
+    def _apply_basic(self, img, mixing_weights, m):
+        # This is a literal adaptation of the paper/official implementation without normalizations and
+        # PIL <-> Numpy conversions between every op. It is still quite CPU compute heavy compared to the
+        # typical augmentation transforms, could use a GPU / Kornia implementation.
+        img_shape = img.size[0], img.size[1], len(img.getbands())
+        mixed = np.zeros(img_shape, dtype=np.float32)
+        for mw in mixing_weights:
+            depth = self.depth if self.depth > 0 else np.random.randint(1, 4)
+            ops = np.random.choice(self.ops, depth, replace=True)
+            img_aug = img  # no ops are in-place, deep copy not necessary
+            for op in ops:
+                img_aug = op(img_aug)
+            mixed += mw * np.asarray(img_aug, dtype=np.float32)
+        np.clip(mixed, 0, 255., out=mixed)
+        mixed = Image.fromarray(mixed.astype(np.uint8))
+        return Image.blend(img, mixed, m)
+
+    def __call__(self, img):
+        mixing_weights = np.float32(np.random.dirichlet([self.alpha] * self.width))
+        m = np.float32(np.random.beta(self.alpha, self.alpha))
+        if self.blended:
+            mixed = self._apply_blended(img, mixing_weights, m)
+        else:
+            mixed = self._apply_basic(img, mixing_weights, m)
+        return mixed
+
+
+def augment_and_mix_transform(config_str, hparams):
+    """ Create AugMix PyTorch transform
+    :param config_str: String defining configuration of random augmentation. Consists of multiple sections separated by
+    dashes ('-'). The first section defines the specific variant of rand augment (currently only 'rand'). The remaining
+    sections, not order sepecific determine
+        'm' - integer magnitude (severity) of augmentation mix (default: 3)
+        'w' - integer width of augmentation chain (default: 3)
+        'd' - integer depth of augmentation chain (-1 is random [1, 3], default: -1)
+        'b' - integer (bool), blend each branch of chain into end result without a final blend, less CPU (default: 0)
+        'mstd' -  float std deviation of magnitude noise applied (default: 0)
+    Ex 'augmix-m5-w4-d2' results in AugMix with severity 5, chain width 4, chain depth 2
+    :param hparams: Other hparams (kwargs) for the Augmentation transforms
+    :return: A PyTorch compatible Transform
+    """
+    magnitude = 3
+    width = 3
+    depth = -1
+    alpha = 1.
+    blended = False
+    config = config_str.split('-')
+    assert config[0] == 'augmix'
+    config = config[1:]
+    for c in config:
+        cs = re.split(r'(\d.*)', c)
+        if len(cs) < 2:
+            continue
+        key, val = cs[:2]
+        if key == 'mstd':
+            # noise param injected via hparams for now
+            hparams.setdefault('magnitude_std', float(val))
+        elif key == 'm':
+            magnitude = int(val)
+        elif key == 'w':
+            width = int(val)
+        elif key == 'd':
+            depth = int(val)
+        elif key == 'a':
+            alpha = float(val)
+        elif key == 'b':
+            blended = bool(val)
+        else:
+            assert False, 'Unknown AugMix config section'
+    ops = augmix_ops(magnitude=magnitude, hparams=hparams)
+    return AugMixAugment(ops, alpha=alpha, width=width, depth=depth, blended=blended)

fastreid/data/transforms/build.py

@@ -0,0 +1,71 @@
+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+
+import torchvision.transforms as T
+
+from .transforms import *
+from .autoaugment import AutoAugment
+
+
+def build_transforms(cfg, is_train=True):
+    res = []
+
+    if is_train:
+        size_train = cfg.INPUT.SIZE_TRAIN
+
+        # augmix augmentation
+        do_augmix = cfg.INPUT.DO_AUGMIX
+
+        # auto augmentation
+        do_autoaug = cfg.INPUT.DO_AUTOAUG
+        total_iter = cfg.SOLVER.MAX_ITER
+
+        # horizontal filp
+        do_flip = cfg.INPUT.DO_FLIP
+        flip_prob = cfg.INPUT.FLIP_PROB
+
+        # padding
+        do_pad = cfg.INPUT.DO_PAD
+        padding = cfg.INPUT.PADDING
+        padding_mode = cfg.INPUT.PADDING_MODE
+
+        # color jitter
+        do_cj = cfg.INPUT.CJ.ENABLED
+        cj_prob = cfg.INPUT.CJ.PROB
+        cj_brightness = cfg.INPUT.CJ.BRIGHTNESS
+        cj_contrast = cfg.INPUT.CJ.CONTRAST
+        cj_saturation = cfg.INPUT.CJ.SATURATION
+        cj_hue = cfg.INPUT.CJ.HUE
+
+        # random erasing
+        do_rea = cfg.INPUT.REA.ENABLED
+        rea_prob = cfg.INPUT.REA.PROB
+        rea_mean = cfg.INPUT.REA.MEAN
+        # random patch
+        do_rpt = cfg.INPUT.RPT.ENABLED
+        rpt_prob = cfg.INPUT.RPT.PROB
+
+        if do_autoaug:
+            res.append(AutoAugment(total_iter))
+        res.append(T.Resize(size_train, interpolation=3))
+        if do_flip:
+            res.append(T.RandomHorizontalFlip(p=flip_prob))
+        if do_pad:
+            res.extend([T.Pad(padding, padding_mode=padding_mode),
+                        T.RandomCrop(size_train)])
+        if do_cj:
+            res.append(T.RandomApply([T.ColorJitter(cj_brightness, cj_contrast, cj_saturation, cj_hue)], p=cj_prob))
+        if do_augmix:
+            res.append(AugMix())
+        if do_rea:
+            res.append(RandomErasing(probability=rea_prob, mean=rea_mean))
+        if do_rpt:
+            res.append(RandomPatch(prob_happen=rpt_prob))
+    else:
+        size_test = cfg.INPUT.SIZE_TEST
+        res.append(T.Resize(size_test, interpolation=3))
+    res.append(ToTensor())
+    return T.Compose(res)

fastreid/data/transforms/functional.py

@@ -0,0 +1,190 @@
+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+
+import numpy as np
+import torch
+from PIL import Image, ImageOps, ImageEnhance
+
+
+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 isinstance(pic, np.ndarray):
+        assert len(pic.shape) in (2, 3)
+        # handle numpy array
+        if pic.ndim == 2:
+            pic = pic[:, :, None]
+
+        img = torch.from_numpy(pic.transpose((2, 0, 1)))
+        # backward compatibility
+        if isinstance(img, torch.ByteTensor):
+            return img.float()
+        else:
+            return img
+
+    # 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, LA, 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()
+    else:
+        return img
+
+
+def int_parameter(level, maxval):
+    """Helper function to scale `val` between 0 and maxval .
+    Args:
+      level: Level of the operation that will be between [0, `PARAMETER_MAX`].
+      maxval: Maximum value that the operation can have. This will be scaled to
+        level/PARAMETER_MAX.
+    Returns:
+      An int that results from scaling `maxval` according to `level`.
+    """
+    return int(level * maxval / 10)
+
+
+def float_parameter(level, maxval):
+    """Helper function to scale `val` between 0 and maxval.
+    Args:
+      level: Level of the operation that will be between [0, `PARAMETER_MAX`].
+      maxval: Maximum value that the operation can have. This will be scaled to
+        level/PARAMETER_MAX.
+    Returns:
+      A float that results from scaling `maxval` according to `level`.
+    """
+    return float(level) * maxval / 10.
+
+
+def sample_level(n):
+    return np.random.uniform(low=0.1, high=n)
+
+
+def autocontrast(pil_img, *args):
+    return ImageOps.autocontrast(pil_img)
+
+
+def equalize(pil_img, *args):
+    return ImageOps.equalize(pil_img)
+
+
+def posterize(pil_img, level, *args):
+    level = int_parameter(sample_level(level), 4)
+    return ImageOps.posterize(pil_img, 4 - level)
+
+
+def rotate(pil_img, level, *args):
+    degrees = int_parameter(sample_level(level), 30)
+    if np.random.uniform() > 0.5:
+        degrees = -degrees
+    return pil_img.rotate(degrees, resample=Image.BILINEAR)
+
+
+def solarize(pil_img, level, *args):
+    level = int_parameter(sample_level(level), 256)
+    return ImageOps.solarize(pil_img, 256 - level)
+
+
+def shear_x(pil_img, level, image_size):
+    level = float_parameter(sample_level(level), 0.3)
+    if np.random.uniform() > 0.5:
+        level = -level
+    return pil_img.transform(image_size,
+                             Image.AFFINE, (1, level, 0, 0, 1, 0),
+                             resample=Image.BILINEAR)
+
+
+def shear_y(pil_img, level, image_size):
+    level = float_parameter(sample_level(level), 0.3)
+    if np.random.uniform() > 0.5:
+        level = -level
+    return pil_img.transform(image_size,
+                             Image.AFFINE, (1, 0, 0, level, 1, 0),
+                             resample=Image.BILINEAR)
+
+
+def translate_x(pil_img, level, image_size):
+    level = int_parameter(sample_level(level), image_size[0] / 3)
+    if np.random.random() > 0.5:
+        level = -level
+    return pil_img.transform(image_size,
+                             Image.AFFINE, (1, 0, level, 0, 1, 0),
+                             resample=Image.BILINEAR)
+
+
+def translate_y(pil_img, level, image_size):
+    level = int_parameter(sample_level(level), image_size[1] / 3)
+    if np.random.random() > 0.5:
+        level = -level
+    return pil_img.transform(image_size,
+                             Image.AFFINE, (1, 0, 0, 0, 1, level),
+                             resample=Image.BILINEAR)
+
+
+# operation that overlaps with ImageNet-C's test set
+def color(pil_img, level, *args):
+    level = float_parameter(sample_level(level), 1.8) + 0.1
+    return ImageEnhance.Color(pil_img).enhance(level)
+
+
+# operation that overlaps with ImageNet-C's test set
+def contrast(pil_img, level, *args):
+    level = float_parameter(sample_level(level), 1.8) + 0.1
+    return ImageEnhance.Contrast(pil_img).enhance(level)
+
+
+# operation that overlaps with ImageNet-C's test set
+def brightness(pil_img, level, *args):
+    level = float_parameter(sample_level(level), 1.8) + 0.1
+    return ImageEnhance.Brightness(pil_img).enhance(level)
+
+
+# operation that overlaps with ImageNet-C's test set
+def sharpness(pil_img, level, *args):
+    level = float_parameter(sample_level(level), 1.8) + 0.1
+    return ImageEnhance.Sharpness(pil_img).enhance(level)
+
+
+augmentations_reid = [
+    autocontrast, equalize, posterize, shear_x, shear_y,
+    color, contrast, brightness, sharpness
+]
+
+augmentations = [
+    autocontrast, equalize, posterize, rotate, solarize, shear_x, shear_y,
+    translate_x, translate_y
+]
+
+augmentations_all = [
+    autocontrast, equalize, posterize, rotate, solarize, shear_x, shear_y,
+    translate_x, translate_y, color, contrast, brightness, sharpness
+]

fastreid/data/transforms/transforms.py

@@ -0,0 +1,204 @@
+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+
+__all__ = ['ToTensor', 'RandomErasing', 'RandomPatch', 'AugMix',]
+
+import math
+import random
+from collections import deque
+
+import numpy as np
+from PIL import Image
+
+from .functional import to_tensor, augmentations_reid
+
+
+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, 255.0]
+    if the PIL Image belongs to one of the modes (L, LA, P, I, F, RGB, YCbCr, RGBA, CMYK, 1)
+    or if the numpy.ndarray has dtype = np.uint8
+
+    In the other cases, tensors are returned without scaling.
+    """
+
+    def __call__(self, pic):
+        """
+        Args:
+            pic (PIL Image or numpy.ndarray): Image to be converted to tensor.
+
+        Returns:
+            Tensor: Converted image.
+        """
+        return to_tensor(pic)
+
+    def __repr__(self):
+        return self.__class__.__name__ + '()'
+
+
+class RandomErasing(object):
+    """ Randomly selects a rectangle region in an image and erases its pixels.
+        'Random Erasing Data Augmentation' by Zhong et al.
+        See https://arxiv.org/pdf/1708.04896.pdf
+    Args:
+        probability: The probability that the Random Erasing operation will be performed.
+        sl: Minimum proportion of erased area against input image.
+        sh: Maximum proportion of erased area against input image.
+        r1: Minimum aspect ratio of erased area.
+        mean: Erasing value.
+    """
+
+    def __init__(self, probability=0.5, sl=0.02, sh=0.4, r1=0.3, mean=255 * (0.49735, 0.4822, 0.4465)):
+        self.probability = probability
+        self.mean = mean
+        self.sl = sl
+        self.sh = sh
+        self.r1 = r1
+
+    def __call__(self, img):
+        img = np.asarray(img, dtype=np.float32).copy()
+        if random.uniform(0, 1) > self.probability:
+            return img
+
+        for attempt in range(100):
+            area = img.shape[0] * img.shape[1]
+            target_area = random.uniform(self.sl, self.sh) * area
+            aspect_ratio = random.uniform(self.r1, 1 / self.r1)
+
+            h = int(round(math.sqrt(target_area * aspect_ratio)))
+            w = int(round(math.sqrt(target_area / aspect_ratio)))
+
+            if w < img.shape[1] and h < img.shape[0]:
+                x1 = random.randint(0, img.shape[0] - h)
+                y1 = random.randint(0, img.shape[1] - w)
+                if img.shape[2] == 3:
+                    img[x1:x1 + h, y1:y1 + w, 0] = self.mean[0]
+                    img[x1:x1 + h, y1:y1 + w, 1] = self.mean[1]
+                    img[x1:x1 + h, y1:y1 + w, 2] = self.mean[2]
+                else:
+                    img[x1:x1 + h, y1:y1 + w, 0] = self.mean[0]
+                return img
+        return img
+
+
+class RandomPatch(object):
+    """Random patch data augmentation.
+    There is a patch pool that stores randomly extracted pathces from person images.
+    For each input image, RandomPatch
+        1) extracts a random patch and stores the patch in the patch pool;
+        2) randomly selects a patch from the patch pool and pastes it on the
+           input (at random position) to simulate occlusion.
+    Reference:
+        - Zhou et al. Omni-Scale Feature Learning for Person Re-Identification. ICCV, 2019.
+        - Zhou et al. Learning Generalisable Omni-Scale Representations
+          for Person Re-Identification. arXiv preprint, 2019.
+    """
+
+    def __init__(self, prob_happen=0.5, pool_capacity=50000, min_sample_size=100,
+                 patch_min_area=0.01, patch_max_area=0.5, patch_min_ratio=0.1,
+                 prob_rotate=0.5, prob_flip_leftright=0.5,
+                 ):
+        self.prob_happen = prob_happen
+
+        self.patch_min_area = patch_min_area
+        self.patch_max_area = patch_max_area
+        self.patch_min_ratio = patch_min_ratio
+
+        self.prob_rotate = prob_rotate
+        self.prob_flip_leftright = prob_flip_leftright
+
+        self.patchpool = deque(maxlen=pool_capacity)
+        self.min_sample_size = min_sample_size
+
+    def generate_wh(self, W, H):
+        area = W * H
+        for attempt in range(100):
+            target_area = random.uniform(self.patch_min_area, self.patch_max_area) * area
+            aspect_ratio = random.uniform(self.patch_min_ratio, 1. / self.patch_min_ratio)
+            h = int(round(math.sqrt(target_area * aspect_ratio)))
+            w = int(round(math.sqrt(target_area / aspect_ratio)))
+            if w < W and h < H:
+                return w, h
+        return None, None
+
+    def transform_patch(self, patch):
+        if random.uniform(0, 1) > self.prob_flip_leftright:
+            patch = patch.transpose(Image.FLIP_LEFT_RIGHT)
+        if random.uniform(0, 1) > self.prob_rotate:
+            patch = patch.rotate(random.randint(-10, 10))
+        return patch
+
+    def __call__(self, img):
+        if isinstance(img, np.ndarray):
+            img = Image.fromarray(img.astype(np.uint8))
+
+        W, H = img.size  # original image size
+
+        # collect new patch
+        w, h = self.generate_wh(W, H)
+        if w is not None and h is not None:
+            x1 = random.randint(0, W - w)
+            y1 = random.randint(0, H - h)
+            new_patch = img.crop((x1, y1, x1 + w, y1 + h))
+            self.patchpool.append(new_patch)
+
+        if len(self.patchpool) < self.min_sample_size:
+            return img
+
+        if random.uniform(0, 1) > self.prob_happen:
+            return img
+
+        # paste a randomly selected patch on a random position
+        patch = random.sample(self.patchpool, 1)[0]
+        patchW, patchH = patch.size
+        x1 = random.randint(0, W - patchW)
+        y1 = random.randint(0, H - patchH)
+        patch = self.transform_patch(patch)
+        img.paste(patch, (x1, y1))
+
+        return img
+
+
+class AugMix(object):
+    """ Perform AugMix augmentation and compute mixture.
+    Args:
+        aug_prob_coeff: Probability distribution coefficients.
+        mixture_width: Number of augmentation chains to mix per augmented example.
+        mixture_depth: Depth of augmentation chains. -1 denotes stochastic depth in [1, 3]'
+        severity: Severity of underlying augmentation operators (between 1 to 10).
+    """
+
+    def __init__(self, aug_prob_coeff=1, mixture_width=3, mixture_depth=-1, severity=1):
+        self.aug_prob_coeff = aug_prob_coeff
+        self.mixture_width = mixture_width
+        self.mixture_depth = mixture_depth
+        self.severity = severity
+        self.aug_list = augmentations_reid
+
+    def __call__(self, image):
+        """Perform AugMix augmentations and compute mixture.
+        Returns:
+          mixed: Augmented and mixed image.
+        """
+        ws = np.float32(
+            np.random.dirichlet([self.aug_prob_coeff] * self.mixture_width))
+        m = np.float32(np.random.beta(self.aug_prob_coeff, self.aug_prob_coeff))
+
+        image = np.asarray(image, dtype=np.float32).copy()
+        mix = np.zeros_like(image)
+        h, w = image.shape[0], image.shape[1]
+        for i in range(self.mixture_width):
+            image_aug = Image.fromarray(image.copy().astype(np.uint8))
+            depth = self.mixture_depth if self.mixture_depth > 0 else np.random.randint(1, 4)
+            for _ in range(depth):
+                op = np.random.choice(self.aug_list)
+                image_aug = op(image_aug, self.severity, (w, h))
+            mix += ws[i] * np.asarray(image_aug, dtype=np.float32)
+
+        mixed = (1 - m) * image + m * mix
+        return mixed

fastreid/demo.py

@@ -0,0 +1,165 @@
+import atexit
+import bisect
+from collections import deque
+
+import cv2
+import torch
+import torch.multiprocessing as mp
+
+from fastreid.engine import DefaultPredictor
+
+try:
+    mp.set_start_method('spawn')
+except RuntimeError:
+    pass
+
+
+class FeatureExtractionDemo(object):
+    def __init__(self, cfg, parallel=False):
+        """
+        Args:
+            cfg (CfgNode):
+            parallel (bool) whether to run the model in different processes from visualization.:
+                Useful since the visualization logic can be slow.
+        """
+        self.cfg = cfg
+        self.parallel = parallel
+
+        if parallel:
+            self.num_gpus = torch.cuda.device_count()
+            self.predictor = AsyncPredictor(cfg, self.num_gpus)
+        else:
+            self.predictor = DefaultPredictor(cfg)
+
+    def run_on_image(self, original_image):
+        """
+        Args:
+            original_image (np.ndarray): an image of shape (H, W, C) (in BGR order).
+                This is the format used by OpenCV.
+        Returns:
+            predictions (np.ndarray): normalized feature of the model.
+        """
+        # the model expects RGB inputs
+        original_image = original_image[:, :, ::-1]
+        # Apply pre-processing to image.
+        image = cv2.resize(original_image, tuple(self.cfg.INPUT.SIZE_TEST[::-1]), interpolation=cv2.INTER_CUBIC)
+        # Make shape with a new batch dimension which is adapted for
+        # network input
+        image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))[None]
+        predictions = self.predictor(image)
+        return predictions
+
+    def run_on_loader(self, data_loader):
+        if self.parallel:
+            buffer_size = self.predictor.default_buffer_size
+
+            batch_data = deque()
+
+            for cnt, batch in enumerate(data_loader):
+                batch_data.append(batch)
+                self.predictor.put(batch["images"])
+
+                if cnt >= buffer_size:
+                    batch = batch_data.popleft()
+                    predictions = self.predictor.get()
+                    yield predictions, batch["targets"].cpu().numpy(), batch["camids"].cpu().numpy()
+
+            while len(batch_data):
+                batch = batch_data.popleft()
+                predictions = self.predictor.get()
+                yield predictions, batch["targets"].cpu().numpy(), batch["camids"].cpu().numpy()
+        else:
+            for batch in data_loader:
+                predictions = self.predictor(batch["images"])
+                yield predictions, batch["targets"].cpu().numpy(), batch["camids"].cpu().numpy()
+
+
+class AsyncPredictor:
+    """
+    A predictor that runs the model asynchronously, possibly on >1 GPUs.
+    Because when the amount of data is large.
+    """
+
+    class _StopToken:
+        pass
+
+    class _PredictWorker(mp.Process):
+        def __init__(self, cfg, task_queue, result_queue):
+            self.cfg = cfg
+            self.task_queue = task_queue
+            self.result_queue = result_queue
+            super().__init__()
+
+        def run(self):
+            predictor = DefaultPredictor(self.cfg)
+
+            while True:
+                task = self.task_queue.get()
+                if isinstance(task, AsyncPredictor._StopToken):
+                    break
+                idx, data = task
+                result = predictor(data)
+                self.result_queue.put((idx, result))
+
+    def __init__(self, cfg, num_gpus: int = 1):
+        """
+        Args:
+            cfg (CfgNode):
+            num_gpus (int): if 0, will run on CPU
+        """
+        num_workers = max(num_gpus, 1)
+        self.task_queue = mp.Queue(maxsize=num_workers * 3)
+        self.result_queue = mp.Queue(maxsize=num_workers * 3)
+        self.procs = []
+        for gpuid in range(max(num_gpus, 1)):
+            cfg = cfg.clone()
+            cfg.defrost()
+            cfg.MODEL.DEVICE = "cuda:{}".format(gpuid) if num_gpus > 0 else "cpu"
+            self.procs.append(
+                AsyncPredictor._PredictWorker(cfg, self.task_queue, self.result_queue)
+            )
+
+        self.put_idx = 0
+        self.get_idx = 0
+        self.result_rank = []
+        self.result_data = []
+
+        for p in self.procs:
+            p.start()
+
+        atexit.register(self.shutdown)
+
+    def put(self, image):
+        self.put_idx += 1
+        self.task_queue.put((self.put_idx, image))
+
+    def get(self):
+        self.get_idx += 1
+        if len(self.result_rank) and self.result_rank[0] == self.get_idx:
+            res = self.result_data[0]
+            del self.result_data[0], self.result_rank[0]
+            return res
+
+        while True:
+            # Make sure the results are returned in the correct order
+            idx, res = self.result_queue.get()
+            if idx == self.get_idx:
+                return res
+            insert = bisect.bisect(self.result_rank, idx)
+            self.result_rank.insert(insert, idx)
+            self.result_data.insert(insert, res)
+
+    def __len__(self):
+        return self.put_idx - self.get_idx
+
+    def __call__(self, image):
+        self.put(image)
+        return self.get()
+
+    def shutdown(self):
+        for _ in self.procs:
+            self.task_queue.put(AsyncPredictor._StopToken())
+
+    @property
+    def default_buffer_size(self):
+        return len(self.procs) * 5

fastreid/engine/__init__.py

@@ -0,0 +1,15 @@
+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+from .train_loop import *
+
+__all__ = [k for k in globals().keys() if not k.startswith("_")]
+
+
+# prefer to let hooks and defaults live in separate namespaces (therefore not in __all__)
+# but still make them available here
+from .hooks import *
+from .defaults import *
+from .launch import *

fastreid/engine/defaults.py

@@ -0,0 +1,515 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+"""
+This file contains components with some default boilerplate logic user may need
+in training / testing. They will not work for everyone, but many users may find them useful.
+The behavior of functions/classes in this file is subject to change,
+since they are meant to represent the "common default behavior" people need in their projects.
+"""
+
+import argparse
+import logging
+import os
+import sys
+from collections import OrderedDict
+
+import torch
+import torch.nn.functional as F
+from torch.nn.parallel import DistributedDataParallel
+
+from fastreid.data import build_reid_test_loader, build_reid_train_loader
+from fastreid.evaluation import (DatasetEvaluator, ReidEvaluator,
+                                 inference_on_dataset, print_csv_format)
+from fastreid.modeling.meta_arch import build_model
+from fastreid.solver import build_lr_scheduler, build_optimizer
+from fastreid.utils import comm
+from fastreid.utils.checkpoint import Checkpointer
+from fastreid.utils.collect_env import collect_env_info
+from fastreid.utils.env import seed_all_rng
+from fastreid.utils.events import CommonMetricPrinter, JSONWriter, TensorboardXWriter
+from fastreid.utils.file_io import PathManager
+from fastreid.utils.logger import setup_logger
+from . import hooks
+from .train_loop import SimpleTrainer
+
+__all__ = ["default_argument_parser", "default_setup", "DefaultPredictor", "DefaultTrainer"]
+
+
+def default_argument_parser():
+    """
+    Create a parser with some common arguments used by fastreid users.
+    Returns:
+        argparse.ArgumentParser:
+    """
+    parser = argparse.ArgumentParser(description="fastreid Training")
+    parser.add_argument("--config-file", default="", metavar="FILE", help="path to config file")
+    parser.add_argument(
+        "--resume",
+        action="store_true",
+        help="whether to attempt to resume from the checkpoint directory",
+    )
+    parser.add_argument("--eval-only", action="store_true", help="perform evaluation only")
+    parser.add_argument("--num-gpus", type=int, default=1, help="number of gpus *per machine*")
+    parser.add_argument("--num-machines", type=int, default=1, help="total number of machines")
+    parser.add_argument(
+        "--machine-rank", type=int, default=0, help="the rank of this machine (unique per machine)"
+    )
+
+    # PyTorch still may leave orphan processes in multi-gpu training.
+    # Therefore we use a deterministic way to obtain port,
+    # so that users are aware of orphan processes by seeing the port occupied.
+    port = 2 ** 15 + 2 ** 14 + hash(os.getuid() if sys.platform != "win32" else 1) % 2 ** 14
+    parser.add_argument("--dist-url", default="tcp://127.0.0.1:{}".format(port))
+    parser.add_argument(
+        "opts",
+        help="Modify config options using the command-line",
+        default=None,
+        nargs=argparse.REMAINDER,
+    )
+    return parser
+
+
+def default_setup(cfg, args):
+    """
+    Perform some basic common setups at the beginning of a job, including:
+    1. Set up the detectron2 logger
+    2. Log basic information about environment, cmdline arguments, and config
+    3. Backup the config to the output directory
+    Args:
+        cfg (CfgNode): the full config to be used
+        args (argparse.NameSpace): the command line arguments to be logged
+    """
+    output_dir = cfg.OUTPUT_DIR
+    if comm.is_main_process() and output_dir:
+        PathManager.mkdirs(output_dir)
+
+    rank = comm.get_rank()
+    setup_logger(output_dir, distributed_rank=rank, name="fvcore")
+    logger = setup_logger(output_dir, distributed_rank=rank)
+
+    logger.info("Rank of current process: {}. World size: {}".format(rank, comm.get_world_size()))
+    logger.info("Environment info:\n" + collect_env_info())
+
+    logger.info("Command line arguments: " + str(args))
+    if hasattr(args, "config_file") and args.config_file != "":
+        logger.info(
+            "Contents of args.config_file={}:\n{}".format(
+                args.config_file, PathManager.open(args.config_file, "r").read()
+            )
+        )
+
+    logger.info("Running with full config:\n{}".format(cfg))
+    if comm.is_main_process() and output_dir:
+        # Note: some of our scripts may expect the existence of
+        # config.yaml in output directory
+        path = os.path.join(output_dir, "config.yaml")
+        with PathManager.open(path, "w") as f:
+            f.write(cfg.dump())
+        logger.info("Full config saved to {}".format(os.path.abspath(path)))
+
+    # make sure each worker has a different, yet deterministic seed if specified
+    seed_all_rng()
+
+    # cudnn benchmark has large overhead. It shouldn't be used considering the small size of
+    # typical validation set.
+    if not (hasattr(args, "eval_only") and args.eval_only):
+        torch.backends.cudnn.benchmark = cfg.CUDNN_BENCHMARK
+
+
+class DefaultPredictor:
+    """
+    Create a simple end-to-end predictor with the given config.
+    The predictor takes an BGR image, resizes it to the specified resolution,
+    runs the model and produces a dict of predictions.
+    This predictor takes care of model loading and input preprocessing for you.
+    If you'd like to do anything more fancy, please refer to its source code
+    as examples to build and use the model manually.
+    Attributes:
+    Examples:
+    .. code-block:: python
+        pred = DefaultPredictor(cfg)
+        inputs = cv2.imread("input.jpg")
+        outputs = pred(inputs)
+    """
+
+    def __init__(self, cfg):
+        self.cfg = cfg.clone()  # cfg can be modified by model
+        self.cfg.defrost()
+        self.cfg.MODEL.BACKBONE.PRETRAIN = False
+        self.model = build_model(self.cfg)
+        self.model.eval()
+
+        Checkpointer(self.model).load(cfg.MODEL.WEIGHTS)
+
+    def __call__(self, image):
+        """
+        Args:
+            image (torch.tensor): an image tensor of shape (B, C, H, W).
+        Returns:
+            predictions (torch.tensor): the output features of the model
+        """
+        inputs = {"images": image}
+        with torch.no_grad():  # https://github.com/sphinx-doc/sphinx/issues/4258
+            predictions = self.model(inputs)
+            # Normalize feature to compute cosine distance
+            features = F.normalize(predictions)
+            features = features.cpu().data
+            return features
+
+
+class DefaultTrainer(SimpleTrainer):
+    """
+    A trainer with default training logic. Compared to `SimpleTrainer`, it
+    contains the following logic in addition:
+    1. Create model, optimizer, scheduler, dataloader from the given config.
+    2. Load a checkpoint or `cfg.MODEL.WEIGHTS`, if exists.
+    3. Register a few common hooks.
+    It is created to simplify the **standard model training workflow** and reduce code boilerplate
+    for users who only need the standard training workflow, with standard features.
+    It means this class makes *many assumptions* about your training logic that
+    may easily become invalid in a new research. In fact, any assumptions beyond those made in the
+    :class:`SimpleTrainer` are too much for research.
+    The code of this class has been annotated about restrictive assumptions it mades.
+    When they do not work for you, you're encouraged to:
+    1. Overwrite methods of this class, OR:
+    2. Use :class:`SimpleTrainer`, which only does minimal SGD training and
+       nothing else. You can then add your own hooks if needed. OR:
+    3. Write your own training loop similar to `tools/plain_train_net.py`.
+    Also note that the behavior of this class, like other functions/classes in
+    this file, is not stable, since it is meant to represent the "common default behavior".
+    It is only guaranteed to work well with the standard models and training workflow in fastreid.
+    To obtain more stable behavior, write your own training logic with other public APIs.
+    Attributes:
+        scheduler:
+        checkpointer:
+        cfg (CfgNode):
+    Examples:
+    .. code-block:: python
+        trainer = DefaultTrainer(cfg)
+        trainer.resume_or_load()  # load last checkpoint or MODEL.WEIGHTS
+        trainer.train()
+    """
+
+    def __init__(self, cfg):
+        """
+        Args:
+            cfg (CfgNode):
+        """
+        logger = logging.getLogger("fastreid")
+        if not logger.isEnabledFor(logging.INFO):  # setup_logger is not called for fastreid
+            setup_logger()
+
+        # Assume these objects must be constructed in this order.
+        data_loader = self.build_train_loader(cfg)
+        cfg = self.auto_scale_hyperparams(cfg, data_loader)
+        model = self.build_model(cfg)
+        optimizer = self.build_optimizer(cfg, model)
+
+        # For training, wrap with DDP. But don't need this for inference.
+        if comm.get_world_size() > 1:
+            # ref to https://github.com/pytorch/pytorch/issues/22049 to set `find_unused_parameters=True`
+            # for part of the parameters is not updated.
+            model = DistributedDataParallel(
+                model, device_ids=[comm.get_local_rank()], broadcast_buffers=False
+            )
+
+        super().__init__(model, data_loader, optimizer, cfg.SOLVER.AMP_ENABLED)
+
+        self.scheduler = self.build_lr_scheduler(cfg, optimizer)
+        # Assume no other objects need to be checkpointed.
+        # We can later make it checkpoint the stateful hooks
+        self.checkpointer = Checkpointer(
+            # Assume you want to save checkpoints together with logs/statistics
+            model,
+            cfg.OUTPUT_DIR,
+            save_to_disk=comm.is_main_process(),
+            optimizer=optimizer,
+            scheduler=self.scheduler,
+        )
+        self.start_iter = 0
+        if cfg.SOLVER.SWA.ENABLED:
+            self.max_iter = cfg.SOLVER.MAX_ITER + cfg.SOLVER.SWA.ITER
+        else:
+            self.max_iter = cfg.SOLVER.MAX_ITER
+
+        self.cfg = cfg
+
+        self.register_hooks(self.build_hooks())
+
+    def resume_or_load(self, resume=True):
+        """
+        If `resume==True` and `cfg.OUTPUT_DIR` contains the last checkpoint (defined by
+        a `last_checkpoint` file), resume from the file. Resuming means loading all
+        available states (eg. optimizer and scheduler) and update iteration counter
+        from the checkpoint. ``cfg.MODEL.WEIGHTS`` will not be used.
+        Otherwise, this is considered as an independent training. The method will load model
+        weights from the file `cfg.MODEL.WEIGHTS` (but will not load other states) and start
+        from iteration 0.
+        Args:
+            resume (bool): whether to do resume or not
+        """
+        # The checkpoint stores the training iteration that just finished, thus we start
+        # at the next iteration (or iter zero if there's no checkpoint).
+        checkpoint = self.checkpointer.resume_or_load(self.cfg.MODEL.WEIGHTS, resume=resume)
+
+        if resume and self.checkpointer.has_checkpoint():
+            self.start_iter = checkpoint.get("iteration", -1) + 1
+            # The checkpoint stores the training iteration that just finished, thus we start
+            # at the next iteration (or iter zero if there's no checkpoint).
+
+    def build_hooks(self):
+        """
+        Build a list of default hooks, including timing, evaluation,
+        checkpointing, lr scheduling, precise BN, writing events.
+        Returns:
+            list[HookBase]:
+        """
+        logger = logging.getLogger(__name__)
+        cfg = self.cfg.clone()
+        cfg.defrost()
+        cfg.DATALOADER.NUM_WORKERS = 0  # save some memory and time for PreciseBN
+        cfg.DATASETS.NAMES = tuple([cfg.TEST.PRECISE_BN.DATASET])  # set dataset name for PreciseBN
+
+        ret = [
+            hooks.IterationTimer(),
+            hooks.LRScheduler(self.optimizer, self.scheduler),
+        ]
+
+        if cfg.SOLVER.SWA.ENABLED:
+            ret.append(
+                hooks.SWA(
+                    cfg.SOLVER.MAX_ITER,
+                    cfg.SOLVER.SWA.PERIOD,
+                    cfg.SOLVER.SWA.LR_FACTOR,
+                    cfg.SOLVER.SWA.ETA_MIN_LR,
+                    cfg.SOLVER.SWA.LR_SCHED,
+                )
+            )
+
+        if cfg.TEST.PRECISE_BN.ENABLED and hooks.get_bn_modules(self.model):
+            logger.info("Prepare precise BN dataset")
+            ret.append(hooks.PreciseBN(
+                # Run at the same freq as (but before) evaluation.
+                self.model,
+                # Build a new data loader to not affect training
+                self.build_train_loader(cfg),
+                cfg.TEST.PRECISE_BN.NUM_ITER,
+            ))
+
+        if cfg.MODEL.FREEZE_LAYERS != [''] and cfg.SOLVER.FREEZE_ITERS > 0:
+            freeze_layers = ",".join(cfg.MODEL.FREEZE_LAYERS)
+            logger.info(f'Freeze layer group "{freeze_layers}" training for {cfg.SOLVER.FREEZE_ITERS:d} iterations')
+            ret.append(hooks.FreezeLayer(
+                self.model,
+                self.optimizer,
+                cfg.MODEL.FREEZE_LAYERS,
+                cfg.SOLVER.FREEZE_ITERS,
+            ))
+        # Do PreciseBN before checkpointer, because it updates the model and need to
+        # be saved by checkpointer.
+        # This is not always the best: if checkpointing has a different frequency,
+        # some checkpoints may have more precise statistics than others.
+        if comm.is_main_process():
+            ret.append(hooks.PeriodicCheckpointer(self.checkpointer, cfg.SOLVER.CHECKPOINT_PERIOD))
+
+        def test_and_save_results():
+            self._last_eval_results = self.test(self.cfg, self.model)
+            return self._last_eval_results
+
+        # Do evaluation after checkpointer, because then if it fails,
+        # we can use the saved checkpoint to debug.
+        ret.append(hooks.EvalHook(cfg.TEST.EVAL_PERIOD, test_and_save_results))
+
+        if comm.is_main_process():
+            # run writers in the end, so that evaluation metrics are written
+            ret.append(hooks.PeriodicWriter(self.build_writers(), 200))
+
+        return ret
+
+    def build_writers(self):
+        """
+        Build a list of writers to be used. By default it contains
+        writers that write metrics to the screen,
+        a json file, and a tensorboard event file respectively.
+        If you'd like a different list of writers, you can overwrite it in
+        your trainer.
+        Returns:
+            list[EventWriter]: a list of :class:`EventWriter` objects.
+        It is now implemented by:
+        .. code-block:: python
+            return [
+                CommonMetricPrinter(self.max_iter),
+                JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
+                TensorboardXWriter(self.cfg.OUTPUT_DIR),
+            ]
+        """
+        # Assume the default print/log frequency.
+        return [
+            # It may not always print what you want to see, since it prints "common" metrics only.
+            CommonMetricPrinter(self.max_iter),
+            JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
+            TensorboardXWriter(self.cfg.OUTPUT_DIR),
+        ]
+
+    def train(self):
+        """
+        Run training.
+        Returns:
+            OrderedDict of results, if evaluation is enabled. Otherwise None.
+        """
+        super().train(self.start_iter, self.max_iter)
+        if comm.is_main_process():
+            assert hasattr(
+                self, "_last_eval_results"
+            ), "No evaluation results obtained during training!"
+            # verify_results(self.cfg, self._last_eval_results)
+            return self._last_eval_results
+
+    @classmethod
+    def build_model(cls, cfg):
+        """
+        Returns:
+            torch.nn.Module:
+        It now calls :func:`fastreid.modeling.build_model`.
+        Overwrite it if you'd like a different model.
+        """
+        model = build_model(cfg)
+        logger = logging.getLogger(__name__)
+        logger.info("Model:\n{}".format(model))
+        return model
+
+    @classmethod
+    def build_optimizer(cls, cfg, model):
+        """
+        Returns:
+            torch.optim.Optimizer:
+        It now calls :func:`fastreid.solver.build_optimizer`.
+        Overwrite it if you'd like a different optimizer.
+        """
+        return build_optimizer(cfg, model)
+
+    @classmethod
+    def build_lr_scheduler(cls, cfg, optimizer):
+        """
+        It now calls :func:`fastreid.solver.build_lr_scheduler`.
+        Overwrite it if you'd like a different scheduler.
+        """
+        return build_lr_scheduler(cfg, optimizer)
+
+    @classmethod
+    def build_train_loader(cls, cfg):
+        """
+        Returns:
+            iterable
+        It now calls :func:`fastreid.data.build_detection_train_loader`.
+        Overwrite it if you'd like a different data loader.
+        """
+        logger = logging.getLogger(__name__)
+        logger.info("Prepare training set")
+        return build_reid_train_loader(cfg)
+
+    @classmethod
+    def build_test_loader(cls, cfg, dataset_name):
+        """
+        Returns:
+            iterable
+        It now calls :func:`fastreid.data.build_detection_test_loader`.
+        Overwrite it if you'd like a different data loader.
+        """
+        return build_reid_test_loader(cfg, dataset_name)
+
+    @classmethod
+    def build_evaluator(cls, cfg, dataset_name, output_dir=None):
+        data_loader, num_query = cls.build_test_loader(cfg, dataset_name)
+        return data_loader, ReidEvaluator(cfg, num_query, output_dir)
+
+    @classmethod
+    def test(cls, cfg, model):
+        """
+        Args:
+            cfg (CfgNode):
+            model (nn.Module):
+        Returns:
+            dict: a dict of result metrics
+        """
+        logger = logging.getLogger(__name__)
+
+        results = OrderedDict()
+        for idx, dataset_name in enumerate(cfg.DATASETS.TESTS):
+            logger.info("Prepare testing set")
+            try:
+                data_loader, evaluator = cls.build_evaluator(cfg, dataset_name)
+            except NotImplementedError:
+                logger.warn(
+                    "No evaluator found. implement its `build_evaluator` method."
+                )
+                results[dataset_name] = {}
+                continue
+            results_i = inference_on_dataset(model, data_loader, evaluator)
+            results[dataset_name] = results_i
+
+        if comm.is_main_process():
+            assert isinstance(
+                results, dict
+            ), "Evaluator must return a dict on the main process. Got {} instead.".format(
+                results
+            )
+            print_csv_format(results)
+
+        if len(results) == 1: results = list(results.values())[0]
+
+        return results
+
+    @staticmethod
+    def auto_scale_hyperparams(cfg, data_loader):
+        r"""
+        This is used for auto-computation actual training iterations,
+        because some hyper-param, such as MAX_ITER, means training epochs rather than iters,
+        so we need to convert specific hyper-param to training iterations.
+        """
+        cfg = cfg.clone()
+        frozen = cfg.is_frozen()
+        cfg.defrost()
+
+        # If you don't hard-code the number of classes, it will compute the number automatically
+        if cfg.MODEL.HEADS.NUM_CLASSES == 0:
+            output_dir = cfg.OUTPUT_DIR
+            cfg.MODEL.HEADS.NUM_CLASSES = data_loader.dataset.num_classes
+            # Update the saved config file to make the number of classes valid
+            if comm.is_main_process() and output_dir:
+                # Note: some of our scripts may expect the existence of
+                # config.yaml in output directory
+                path = os.path.join(output_dir, "config.yaml")
+                with PathManager.open(path, "w") as f:
+                    f.write(cfg.dump())
+
+        iters_per_epoch = len(data_loader.dataset) // cfg.SOLVER.IMS_PER_BATCH
+        cfg.SOLVER.MAX_ITER *= iters_per_epoch
+        cfg.SOLVER.WARMUP_ITERS *= iters_per_epoch
+        cfg.SOLVER.FREEZE_ITERS *= iters_per_epoch
+        cfg.SOLVER.DELAY_ITERS *= iters_per_epoch
+        for i in range(len(cfg.SOLVER.STEPS)):
+            cfg.SOLVER.STEPS[i] *= iters_per_epoch
+        cfg.SOLVER.SWA.ITER *= iters_per_epoch
+        cfg.SOLVER.SWA.PERIOD *= iters_per_epoch
+
+        ckpt_multiple = cfg.SOLVER.CHECKPOINT_PERIOD / cfg.TEST.EVAL_PERIOD
+        # Evaluation period must be divided by 200 for writing into tensorboard.
+        eval_num_mod = (200 - cfg.TEST.EVAL_PERIOD * iters_per_epoch) % 200
+        cfg.TEST.EVAL_PERIOD = cfg.TEST.EVAL_PERIOD * iters_per_epoch + eval_num_mod
+        # Change checkpoint saving period consistent with evaluation period.
+        cfg.SOLVER.CHECKPOINT_PERIOD = int(cfg.TEST.EVAL_PERIOD * ckpt_multiple)
+
+        logger = logging.getLogger(__name__)
+        logger.info(
+            f"Auto-scaling the config to num_classes={cfg.MODEL.HEADS.NUM_CLASSES}, "
+            f"max_Iter={cfg.SOLVER.MAX_ITER}, wamrup_Iter={cfg.SOLVER.WARMUP_ITERS}, "
+            f"freeze_Iter={cfg.SOLVER.FREEZE_ITERS}, delay_Iter={cfg.SOLVER.DELAY_ITERS}, "
+            f"step_Iter={cfg.SOLVER.STEPS}, ckpt_Iter={cfg.SOLVER.CHECKPOINT_PERIOD}, "
+            f"eval_Iter={cfg.TEST.EVAL_PERIOD}."
+        )
+
+        if frozen: cfg.freeze()
+
+        return cfg

fastreid/engine/hooks.py

@@ -0,0 +1,503 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import datetime
+import itertools
+import logging
+import os
+import tempfile
+import time
+from collections import Counter
+
+import torch
+from torch import nn
+from torch.nn.parallel import DistributedDataParallel
+
+from fastreid.evaluation.testing import flatten_results_dict
+from fastreid.solver import optim
+from fastreid.utils import comm
+from fastreid.utils.checkpoint import PeriodicCheckpointer as _PeriodicCheckpointer
+from fastreid.utils.events import EventStorage, EventWriter
+from fastreid.utils.file_io import PathManager
+from fastreid.utils.precision_bn import update_bn_stats, get_bn_modules
+from fastreid.utils.timer import Timer
+from .train_loop import HookBase
+
+__all__ = [
+    "CallbackHook",
+    "IterationTimer",
+    "PeriodicWriter",
+    "PeriodicCheckpointer",
+    "LRScheduler",
+    "AutogradProfiler",
+    "EvalHook",
+    "PreciseBN",
+    "FreezeLayer",
+]
+
+"""
+Implement some common hooks.
+"""
+
+
+class CallbackHook(HookBase):
+    """
+    Create a hook using callback functions provided by the user.
+    """
+
+    def __init__(self, *, before_train=None, after_train=None, before_step=None, after_step=None):
+        """
+        Each argument is a function that takes one argument: the trainer.
+        """
+        self._before_train = before_train
+        self._before_step = before_step
+        self._after_step = after_step
+        self._after_train = after_train
+
+    def before_train(self):
+        if self._before_train:
+            self._before_train(self.trainer)
+
+    def after_train(self):
+        if self._after_train:
+            self._after_train(self.trainer)
+        # The functions may be closures that hold reference to the trainer
+        # Therefore, delete them to avoid circular reference.
+        del self._before_train, self._after_train
+        del self._before_step, self._after_step
+
+    def before_step(self):
+        if self._before_step:
+            self._before_step(self.trainer)
+
+    def after_step(self):
+        if self._after_step:
+            self._after_step(self.trainer)
+
+
+class IterationTimer(HookBase):
+    """
+    Track the time spent for each iteration (each run_step call in the trainer).
+    Print a summary in the end of training.
+    This hook uses the time between the call to its :meth:`before_step`
+    and :meth:`after_step` methods.
+    Under the convention that :meth:`before_step` of all hooks should only
+    take negligible amount of time, the :class:`IterationTimer` hook should be
+    placed at the beginning of the list of hooks to obtain accurate timing.
+    """
+
+    def __init__(self, warmup_iter=3):
+        """
+        Args:
+            warmup_iter (int): the number of iterations at the beginning to exclude
+                from timing.
+        """
+        self._warmup_iter = warmup_iter
+        self._step_timer = Timer()
+
+    def before_train(self):
+        self._start_time = time.perf_counter()
+        self._total_timer = Timer()
+        self._total_timer.pause()
+
+    def after_train(self):
+        logger = logging.getLogger(__name__)
+        total_time = time.perf_counter() - self._start_time
+        total_time_minus_hooks = self._total_timer.seconds()
+        hook_time = total_time - total_time_minus_hooks
+
+        num_iter = self.trainer.iter + 1 - self.trainer.start_iter - self._warmup_iter
+
+        if num_iter > 0 and total_time_minus_hooks > 0:
+            # Speed is meaningful only after warmup
+            # NOTE this format is parsed by grep in some scripts
+            logger.info(
+                "Overall training speed: {} iterations in {} ({:.4f} s / it)".format(
+                    num_iter,
+                    str(datetime.timedelta(seconds=int(total_time_minus_hooks))),
+                    total_time_minus_hooks / num_iter,
+                )
+            )
+
+        logger.info(
+            "Total training time: {} ({} on hooks)".format(
+                str(datetime.timedelta(seconds=int(total_time))),
+                str(datetime.timedelta(seconds=int(hook_time))),
+            )
+        )
+
+    def before_step(self):
+        self._step_timer.reset()
+        self._total_timer.resume()
+
+    def after_step(self):
+        # +1 because we're in after_step
+        iter_done = self.trainer.iter - self.trainer.start_iter + 1
+        if iter_done >= self._warmup_iter:
+            sec = self._step_timer.seconds()
+            self.trainer.storage.put_scalars(time=sec)
+        else:
+            self._start_time = time.perf_counter()
+            self._total_timer.reset()
+
+        self._total_timer.pause()
+
+
+class PeriodicWriter(HookBase):
+    """
+    Write events to EventStorage periodically.
+    It is executed every ``period`` iterations and after the last iteration.
+    """
+
+    def __init__(self, writers, period=20):
+        """
+        Args:
+            writers (list[EventWriter]): a list of EventWriter objects
+            period (int):
+        """
+        self._writers = writers
+        for w in writers:
+            assert isinstance(w, EventWriter), w
+        self._period = period
+
+    def after_step(self):
+        if (self.trainer.iter + 1) % self._period == 0 or (
+                self.trainer.iter == self.trainer.max_iter - 1
+        ):
+            for writer in self._writers:
+                writer.write()
+
+    def after_train(self):
+        for writer in self._writers:
+            writer.close()
+
+
+class PeriodicCheckpointer(_PeriodicCheckpointer, HookBase):
+    """
+    Same as :class:`fastreid.utils.checkpoint.PeriodicCheckpointer`, but as a hook.
+    Note that when used as a hook,
+    it is unable to save additional data other than what's defined
+    by the given `checkpointer`.
+    It is executed every ``period`` iterations and after the last iteration.
+    """
+
+    def before_train(self):
+        self.max_iter = self.trainer.max_iter
+
+    def after_step(self):
+        # No way to use **kwargs
+        self.step(self.trainer.iter)
+
+
+class LRScheduler(HookBase):
+    """
+    A hook which executes a torch builtin LR scheduler and summarizes the LR.
+    It is executed after every iteration.
+    """
+
+    def __init__(self, optimizer, scheduler):
+        """
+        Args:
+            optimizer (torch.optim.Optimizer):
+            scheduler (torch.optim._LRScheduler)
+        """
+        self._optimizer = optimizer
+        self._scheduler = scheduler
+
+        # NOTE: some heuristics on what LR to summarize
+        # summarize the param group with most parameters
+        largest_group = max(len(g["params"]) for g in optimizer.param_groups)
+
+        if largest_group == 1:
+            # If all groups have one parameter,
+            # then find the most common initial LR, and use it for summary
+            lr_count = Counter([g["lr"] for g in optimizer.param_groups])
+            lr = lr_count.most_common()[0][0]
+            for i, g in enumerate(optimizer.param_groups):
+                if g["lr"] == lr:
+                    self._best_param_group_id = i
+                    break
+        else:
+            for i, g in enumerate(optimizer.param_groups):
+                if len(g["params"]) == largest_group:
+                    self._best_param_group_id = i
+                    break
+
+    def after_step(self):
+        lr = self._optimizer.param_groups[self._best_param_group_id]["lr"]
+        self.trainer.storage.put_scalar("lr", lr, smoothing_hint=False)
+        self._scheduler.step()
+
+
+class AutogradProfiler(HookBase):
+    """
+    A hook which runs `torch.autograd.profiler.profile`.
+    Examples:
+    .. code-block:: python
+        hooks.AutogradProfiler(
+             lambda trainer: trainer.iter > 10 and trainer.iter < 20, self.cfg.OUTPUT_DIR
+        )
+    The above example will run the profiler for iteration 10~20 and dump
+    results to ``OUTPUT_DIR``. We did not profile the first few iterations
+    because they are typically slower than the rest.
+    The result files can be loaded in the ``chrome://tracing`` page in chrome browser.
+    Note:
+        When used together with NCCL on older version of GPUs,
+        autograd profiler may cause deadlock because it unnecessarily allocates
+        memory on every device it sees. The memory management calls, if
+        interleaved with NCCL calls, lead to deadlock on GPUs that do not
+        support `cudaLaunchCooperativeKernelMultiDevice`.
+    """
+
+    def __init__(self, enable_predicate, output_dir, *, use_cuda=True):
+        """
+        Args:
+            enable_predicate (callable[trainer -> bool]): a function which takes a trainer,
+                and returns whether to enable the profiler.
+                It will be called once every step, and can be used to select which steps to profile.
+            output_dir (str): the output directory to dump tracing files.
+            use_cuda (bool): same as in `torch.autograd.profiler.profile`.
+        """
+        self._enable_predicate = enable_predicate
+        self._use_cuda = use_cuda
+        self._output_dir = output_dir
+
+    def before_step(self):
+        if self._enable_predicate(self.trainer):
+            self._profiler = torch.autograd.profiler.profile(use_cuda=self._use_cuda)
+            self._profiler.__enter__()
+        else:
+            self._profiler = None
+
+    def after_step(self):
+        if self._profiler is None:
+            return
+        self._profiler.__exit__(None, None, None)
+        out_file = os.path.join(
+            self._output_dir, "profiler-trace-iter{}.json".format(self.trainer.iter)
+        )
+        if "://" not in out_file:
+            self._profiler.export_chrome_trace(out_file)
+        else:
+            # Support non-posix filesystems
+            with tempfile.TemporaryDirectory(prefix="fastreid_profiler") as d:
+                tmp_file = os.path.join(d, "tmp.json")
+                self._profiler.export_chrome_trace(tmp_file)
+                with open(tmp_file) as f:
+                    content = f.read()
+            with PathManager.open(out_file, "w") as f:
+                f.write(content)
+
+
+class EvalHook(HookBase):
+    """
+    Run an evaluation function periodically, and at the end of training.
+    It is executed every ``eval_period`` iterations and after the last iteration.
+    """
+
+    def __init__(self, eval_period, eval_function):
+        """
+        Args:
+            eval_period (int): the period to run `eval_function`.
+            eval_function (callable): a function which takes no arguments, and
+                returns a nested dict of evaluation metrics.
+        Note:
+            This hook must be enabled in all or none workers.
+            If you would like only certain workers to perform evaluation,
+            give other workers a no-op function (`eval_function=lambda: None`).
+        """
+        self._period = eval_period
+        self._func = eval_function
+
+    def _do_eval(self):
+        results = self._func()
+
+        if results:
+            assert isinstance(
+                results, dict
+            ), "Eval function must return a dict. Got {} instead.".format(results)
+
+            flattened_results = flatten_results_dict(results)
+            for k, v in flattened_results.items():
+                try:
+                    v = float(v)
+                except Exception:
+                    raise ValueError(
+                        "[EvalHook] eval_function should return a nested dict of float. "
+                        "Got '{}: {}' instead.".format(k, v)
+                    )
+            self.trainer.storage.put_scalars(**flattened_results, smoothing_hint=False)
+
+        # Remove extra memory cache of main process due to evaluation
+        torch.cuda.empty_cache()
+
+    def after_step(self):
+        next_iter = self.trainer.iter + 1
+        is_final = next_iter == self.trainer.max_iter
+        if is_final or (self._period > 0 and next_iter % self._period == 0):
+            self._do_eval()
+        # Evaluation may take different time among workers.
+        # A barrier make them start the next iteration together.
+        comm.synchronize()
+
+    def after_train(self):
+        # func is likely a closure that holds reference to the trainer
+        # therefore we clean it to avoid circular reference in the end
+        del self._func
+
+
+class PreciseBN(HookBase):
+    """
+    The standard implementation of BatchNorm uses EMA in inference, which is
+    sometimes suboptimal.
+    This class computes the true average of statistics rather than the moving average,
+    and put true averages to every BN layer in the given model.
+    It is executed after the last iteration.
+    """
+
+    def __init__(self, model, data_loader, num_iter):
+        """
+        Args:
+            model (nn.Module): a module whose all BN layers in training mode will be
+                updated by precise BN.
+                Note that user is responsible for ensuring the BN layers to be
+                updated are in training mode when this hook is triggered.
+            data_loader (iterable): it will produce data to be run by `model(data)`.
+            num_iter (int): number of iterations used to compute the precise
+                statistics.
+        """
+        self._logger = logging.getLogger(__name__)
+        if len(get_bn_modules(model)) == 0:
+            self._logger.info(
+                "PreciseBN is disabled because model does not contain BN layers in training mode."
+            )
+            self._disabled = True
+            return
+
+        self._model = model
+        self._data_loader = data_loader
+        self._num_iter = num_iter
+        self._disabled = False
+
+        self._data_iter = None
+
+    def after_step(self):
+        next_iter = self.trainer.iter + 1
+        is_final = next_iter == self.trainer.max_iter
+        if is_final:
+            self.update_stats()
+
+    def update_stats(self):
+        """
+        Update the model with precise statistics. Users can manually call this method.
+        """
+        if self._disabled:
+            return
+
+        if self._data_iter is None:
+            self._data_iter = iter(self._data_loader)
+
+        def data_loader():
+            for num_iter in itertools.count(1):
+                if num_iter % 100 == 0:
+                    self._logger.info(
+                        "Running precise-BN ... {}/{} iterations.".format(num_iter, self._num_iter)
+                    )
+                # This way we can reuse the same iterator
+                yield next(self._data_iter)
+
+        with EventStorage():  # capture events in a new storage to discard them
+            self._logger.info(
+                "Running precise-BN for {} iterations...  ".format(self._num_iter)
+                + "Note that this could produce different statistics every time."
+            )
+            update_bn_stats(self._model, data_loader(), self._num_iter)
+
+
+class FreezeLayer(HookBase):
+    def __init__(self, model, optimizer, freeze_layers, freeze_iters):
+        self._logger = logging.getLogger(__name__)
+
+        if isinstance(model, DistributedDataParallel):
+            model = model.module
+        self.model = model
+        self.optimizer = optimizer
+
+        self.freeze_layers = freeze_layers
+        self.freeze_iters = freeze_iters
+
+        # Previous parameters freeze status
+        param_freeze = {}
+        for param_group in self.optimizer.param_groups:
+            param_name = param_group['name']
+            param_freeze[param_name] = param_group['freeze']
+        self.param_freeze = param_freeze
+
+        self.is_frozen = False
+
+    def before_step(self):
+        # Freeze specific layers
+        if self.trainer.iter <= self.freeze_iters and not self.is_frozen:
+            self.freeze_specific_layer()
+
+        # Recover original layers status
+        if self.trainer.iter > self.freeze_iters and self.is_frozen:
+            self.open_all_layer()
+
+    def freeze_specific_layer(self):
+        for layer in self.freeze_layers:
+            if not hasattr(self.model, layer):
+                self._logger.info(f'{layer} is not an attribute of the model, will skip this layer')
+
+        for param_group in self.optimizer.param_groups:
+            param_name = param_group['name']
+            if param_name.split('.')[0] in self.freeze_layers:
+                param_group['freeze'] = True
+
+        # Change BN in freeze layers to eval mode
+        for name, module in self.model.named_children():
+            if name in self.freeze_layers: module.eval()
+
+        self.is_frozen = True
+
+    def open_all_layer(self):
+        self.model.train()
+        for param_group in self.optimizer.param_groups:
+            param_name = param_group['name']
+            param_group['freeze'] = self.param_freeze[param_name]
+
+        self.is_frozen = False
+
+
+class SWA(HookBase):
+    def __init__(self, swa_start: int, swa_freq: int, swa_lr_factor: float, eta_min: float, lr_sched=False, ):
+        self.swa_start = swa_start
+        self.swa_freq = swa_freq
+        self.swa_lr_factor = swa_lr_factor
+        self.eta_min = eta_min
+        self.lr_sched = lr_sched
+
+    def before_step(self):
+        is_swa = self.trainer.iter == self.swa_start
+        if is_swa:
+            # Wrapper optimizer with SWA
+            self.trainer.optimizer = optim.SWA(self.trainer.optimizer, self.swa_freq, self.swa_lr_factor)
+            self.trainer.optimizer.reset_lr_to_swa()
+
+            if self.lr_sched:
+                self.scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
+                    optimizer=self.trainer.optimizer,
+                    T_0=self.swa_freq,
+                    eta_min=self.eta_min,
+                )
+
+    def after_step(self):
+        next_iter = self.trainer.iter + 1
+
+        # Use Cyclic learning rate scheduler
+        if next_iter > self.swa_start and self.lr_sched:
+            self.scheduler.step()
+
+        is_final = next_iter == self.trainer.max_iter
+        if is_final:
+            self.trainer.optimizer.swap_swa_param()