change query id show syntax

add expandable ReID framework

.gitignore

@@ -4,4 +4,7 @@
 *.jpg
 *.mp4
 *.t7
-*.pth
+*.pth
+*label*
+result.txt
+rank10_detection*.txt

deep_sort/deep/feature_extractor.py

@@ -1,13 +1,11 @@
 import torch
+from torch import nn
 import torchvision.transforms as transforms
 import numpy as np
 import cv2
 import logging
 
-from .model import Net
-from fastreid.config import get_cfg
-from fastreid.engine import DefaultTrainer
-from fastreid.utils.checkpoint import Checkpointer
+from .model import Net,MyNet,RestNet18
 
 class Extractor(object):
     def __init__(self, model_path, use_cuda=True):
@@ -49,36 +47,40 @@ class Extractor(object):
             features = self.net(im_batch)
         return features.cpu().numpy()
 
-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)
+    
+class MyExtractor(object):
+    def __init__(self, model_path, use_cuda=True):
         self.device = "cuda" if torch.cuda.is_available() and use_cuda else "cpu"
-
-        Checkpointer(self.net).load(model_path)
+        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)
         logger = logging.getLogger("root.tracker")
         logger.info("Loading weights from {}... Done!".format(model_path))
-        self.net.to(self.device)
-        self.net.eval()
-        height, width = cfg.INPUT.SIZE_TEST
-        self.size = (width, height)
+        self.raw_transformer = transforms.Compose([
+        transforms.ToPILImage(),
+        transforms.Resize((128,64)),# hxw
+        transforms.ToTensor(),
+        ])   
         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):
-        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()
+        """
+        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)
         return im_batch
 
-
     def __call__(self, im_crops):
         im_batch = self._preprocess(im_crops)
         with torch.no_grad():
@@ -87,10 +89,3 @@ class FastReIDExtractor(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,12 +93,136 @@ 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)
-    import ipdb; ipdb.set_trace()
+    
+#########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
 
 

deep_sort/deep_sort.py

@@ -1,7 +1,7 @@
 import numpy as np
 import torch
 
-from .deep.feature_extractor import Extractor, FastReIDExtractor
+from .deep.feature_extractor import Extractor,MyExtractor
 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=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):
+    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):
         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 = FastReIDExtractor(model_config, model_path, use_cuda=use_cuda)
+            self.extractor = MyExtractor(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,19 +42,21 @@ 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,features
+        return outputs,output_features
 
 
     """

deep_sort/sort/iou_matching.py

@@ -23,6 +23,8 @@ 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:]
@@ -36,7 +38,14 @@ def iou(bbox, candidates):
     area_intersection = wh.prod(axis=1)
     area_bbox = bbox[2:].prod()
     area_candidates = candidates[:, 2:].prod(axis=1)
-    return area_intersection / (area_bbox + area_candidates - area_intersection)
+    
+    # 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)
 
 
 def iou_cost(tracks, detections, track_indices=None,

deep_sort/sort/tracker.py

@@ -86,7 +86,8 @@ 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)
 
@@ -99,7 +100,6 @@ 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,10 +1,7 @@
-import cv2
 import torch
-import numpy as np
 
-from torchvision import transforms
 from models.experimental import attempt_load
-from utils.general import check_img_size,non_max_suppression
+from utils.general import non_max_suppression
 
 import param
 
@@ -13,7 +10,7 @@ class YOLOv7(object):
     def __init__(self,option:param.Parameters) -> None:
         # load FP32 model
         self.opt = option
-        self.model = attempt_load(weights=self.opt.weights, device=self.opt.device)
+        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]

example_deepsort.py

@@ -1,84 +0,0 @@
-import torch
-from torchvision import transforms
-import numpy as np
-import cv2
-
-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 opencv_box_plot(img:cv2.Mat,pred_img:np.array):
-    pred_img = pred_img.astype(np.uint)
-    print(type(img))
-    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)
-    print(type(img))
-    cv2.imshow("test",img)
-    cv2.waitKey(0)
-    cv2.imwrite("result.png",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)
-    cv2.imshow("test",img)
-    cv2.waitKey(1)
-    return img
-
-
-
-# example        
-if __name__ == '__main__':
-    dsort_path = "./deep_sort/deep/checkpoint/ckpt.t7"
-    my_option = param.Parameters()
-    model = YOLOv7(my_option)                       # load model in head
-    deepsort = dsort.DeepSort(model_path=dsort_path,model_config=None,use_cuda=(torch.device("cuda:0") == my_option.device))
-    with torch.no_grad():
-        sources = cv2.VideoCapture("test.mp4")
-        target = cv2.VideoWriter("output.mp4",cv2.VideoWriter_fourcc('M', 'P', '4', '2'),24,(640,640))
-        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)     
-            result = model.detect(img_tensor)                  # get the sequence of result   
-            result = result[0].detach().cpu().numpy()   # the single img is index 0
-            bbox_xywhs = []
-            confs      = []
-            for xyxycc in result: 
-                xywh   =  deepsort._xyxy_to_xywh(xyxycc[0:4])
-                conf   =  xyxycc[4]
-                clas   =  int(xyxycc[5])
-                bbox_xywhs.append(xywh[:])
-                confs.append(conf)
-            dpsort = deepsort.update(bbox_xywh=np.array(bbox_xywhs),confidences=np.array(confs),ori_img=np.array(img))
-            print(f"result:{result}")
-            print(f"dpsort:{dpsort}")
-            frame_t = opencv_sort_plot(np.array(img),result,dpsort)
-            target.write(frame_t)
-        target.release()
-            # img_src = cv2.imread("test.png")            # re-read for imshow
-            # img_src = cv2.resize(img_src,(640,360))
-            # opencv_box_plot(img_src,result)             # show the result in picture

example_reid

@@ -1,97 +0,0 @@
-import torch
-from torchvision import transforms
-import numpy as np
-import cv2
-
-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 opencv_box_plot(img:cv2.Mat,pred_img:np.array):
-    pred_img = pred_img.astype(np.uint)
-    print(type(img))
-    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)
-    print(type(img))
-    cv2.imshow("test",img)
-    cv2.waitKey(0)
-    cv2.imwrite("result.png",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)
-    cv2.imshow("test",img)
-    cv2.waitKey(1)
-    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__':
-    dsort_path = "./deep_sort/deep/checkpoint/market_bot_R50-ibn.pth"
-    cfg_path = "fastreid/cfgs/Market1501/bagtricks_R50-ibn.yml"
-    my_option = param.Parameters()
-    model = YOLOv7(my_option)                       # load model in head
-    deepsort = dsort.DeepSort(model_path=dsort_path,model_config=cfg_path,use_cuda=(torch.device("cuda:0") == my_option.device))
-    id_try = extractidfeature("cam1/1_25.jpg",deepsort.extractor)
-    print(f"id_try:{id_try}")
-    with torch.no_grad():
-        sources = cv2.VideoCapture("test.mp4")
-        target = cv2.VideoWriter("output.mp4",cv2.VideoWriter_fourcc('m', 'p', '4', '2'),24,(640,640))
-        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)     
-            result = model.detect(img_tensor)                  # get the sequence of result   
-            result = result[0].detach().cpu().numpy()   # the single img is index 0
-            bbox_xywhs = []
-            confs      = []
-            for xyxycc in result: 
-                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(img))
-            img_drawing = np.array(img)
-            for i in range(len(features)):
-                similarity = np.sum(np.abs(features[i]-id_try))
-                print(similarity)
-                cx = int((result[i][0]+result[i][2])/2)
-                cy = int((result[i][1]+result[i][3])/2)
-                cv2.putText(img_drawing,f"{similarity*10:>.2f}",(cx,cy-10), cv2.FONT_HERSHEY_PLAIN, 1.0, (255,0,255), 2)
-            print(f"result:{result}")
-            print(f"dpsort:{dpsort}")
-            frame_t = opencv_sort_plot(img_drawing,result,dpsort)
-            target.write(frame_t)
-        target.release()
-            # img_src = cv2.imread("test.png")            # re-read for imshow
-            # img_src = cv2.resize(img_src,(640,360))
-            # opencv_box_plot(img_src,result)             # show the result in picture

example_reid.py

@@ -0,0 +1,185 @@
+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")

example_yolov7.py

@@ -1,40 +0,0 @@
-import torch
-from torchvision import transforms
-import numpy as np
-import cv2
-
-import param
-from detect import YOLOv7
-
-def get_img_test(path_src:str):
-    raw_img = cv2.imread(path_src)
-    raw_transform = transforms.Compose([transforms.ToPILImage(),
-                                        transforms.Resize((360,640)),
-                                        transforms.Pad((0,(640-360)//2)),
-                                        transforms.ToTensor()])
-    return raw_transform(raw_img).unsqueeze(dim=0)
-
-def opencv_box_plot(img:cv2.Mat,pred_img:np.array):
-    pred_img = pred_img.astype(np.uint)
-    print(type(img))
-    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)
-    print(type(img))
-    cv2.imshow("test",img)
-    cv2.waitKey(0)
-    cv2.imwrite("result.png",img)
-
-# example        
-if __name__ == '__main__':
-    my_option = param.Parameters()
-    model = YOLOv7(my_option)                       # load model in head
-    with torch.no_grad():
-        img = get_img_test("test.png")              # get image as torch.Tensor with size of [1,1,640,640]
-        print(f"img size is {img.size()}")        
-        result = model.detect(img)                  # get the sequence of result
-        print(f"result is :{result}")
-        result = result[0].detach().cpu().numpy()   # the single img is index 0
-        img_src = cv2.imread("test.png")            # re-read for imshow
-        img_src = cv2.resize(img_src,(640,360))
-        opencv_box_plot(img_src,result)             # show the result in picture

fastreid/__init__.py

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

fastreid/cfgs/Base-AGW.yml

@@ -1,19 +0,0 @@
-_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

@@ -1,25 +0,0 @@
-_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

@@ -1,59 +0,0 @@
-_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

@@ -1,73 +0,0 @@
-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

@@ -1,12 +0,0 @@
-_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

@@ -1,11 +0,0 @@
-_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

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

fastreid/cfgs/DukeMTMC/AGW_S50.yml

@@ -1,11 +0,0 @@
-_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

@@ -1,12 +0,0 @@
-_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

@@ -1,11 +0,0 @@
-_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

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

fastreid/cfgs/DukeMTMC/bagtricks_S50.yml

@@ -1,11 +0,0 @@
-_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

@@ -1,11 +0,0 @@
-_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

@@ -1,12 +0,0 @@
-_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

@@ -1,11 +0,0 @@
-_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

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

fastreid/cfgs/DukeMTMC/sbs_S50.yml

@@ -1,11 +0,0 @@
-_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

@@ -1,12 +0,0 @@
-_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

@@ -1,11 +0,0 @@
-_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

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

fastreid/cfgs/MSMT17/AGW_S50.yml

@@ -1,11 +0,0 @@
-_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

@@ -1,13 +0,0 @@
-_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

@@ -1,12 +0,0 @@
-_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

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

fastreid/cfgs/MSMT17/bagtricks_S50.yml

@@ -1,12 +0,0 @@
-_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

@@ -1,11 +0,0 @@
-_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

@@ -1,12 +0,0 @@
-_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

@@ -1,11 +0,0 @@
-_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

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

fastreid/cfgs/MSMT17/sbs_S50.yml

@@ -1,11 +0,0 @@
-_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

@@ -1,12 +0,0 @@
-_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

@@ -1,11 +0,0 @@
-_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

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

fastreid/cfgs/Market1501/AGW_S50.yml

@@ -1,11 +0,0 @@
-_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

@@ -1,12 +0,0 @@
-_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

@@ -1,11 +0,0 @@
-_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

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

fastreid/cfgs/Market1501/bagtricks_S50.yml

@@ -1,11 +0,0 @@
-_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

@@ -1,11 +0,0 @@
-_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

@@ -1,12 +0,0 @@
-_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

@@ -1,11 +0,0 @@
-_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

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

fastreid/cfgs/Market1501/sbs_S50.yml

@@ -1,11 +0,0 @@
-_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

@@ -1,33 +0,0 @@
-_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

@@ -1,32 +0,0 @@
-_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

@@ -1,35 +0,0 @@
-_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

@@ -1,8 +0,0 @@
-# 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

@@ -1,159 +0,0 @@
-# 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

@@ -1,273 +0,0 @@
-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

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

fastreid/data/build.py

@@ -1,115 +0,0 @@
-# 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

@@ -1,55 +0,0 @@
-# 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

@@ -1,45 +0,0 @@
-# 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

@@ -1,46 +0,0 @@
-# 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

@@ -1,41 +0,0 @@
-# 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

@@ -1,172 +0,0 @@
-# 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

@@ -1,46 +0,0 @@
-# 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

@@ -1,274 +0,0 @@
-# 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

@@ -1,58 +0,0 @@
-# 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

@@ -1,70 +0,0 @@
-# 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

@@ -1,43 +0,0 @@
-# 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

@@ -1,47 +0,0 @@
-# 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

@@ -1,90 +0,0 @@
-# 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

@@ -1,114 +0,0 @@
-# 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

@@ -1,46 +0,0 @@
-# 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

@@ -1,42 +0,0 @@
-# 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

@@ -1,44 +0,0 @@
-# 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

@@ -1,45 +0,0 @@
-# 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

@@ -1,46 +0,0 @@
-# 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

@@ -1,48 +0,0 @@
-# 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

@@ -1,45 +0,0 @@
-# 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

@@ -1,124 +0,0 @@
-# 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

@@ -1,68 +0,0 @@
-# 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

@@ -1,138 +0,0 @@
-# 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

@@ -1,45 +0,0 @@
-# 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

@@ -1,59 +0,0 @@
-# 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

@@ -1,8 +0,0 @@
-# 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

@@ -1,85 +0,0 @@
-# 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

@@ -1,169 +0,0 @@
-# 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

@@ -1,10 +0,0 @@
-# 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

@@ -1,812 +0,0 @@
-# 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

@@ -1,71 +0,0 @@
-# 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

@@ -1,190 +0,0 @@
-# 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

@@ -1,204 +0,0 @@
-# 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/engine/__init__.py

@@ -1,15 +0,0 @@
-# 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

@@ -1,515 +0,0 @@
-# -*- 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

@@ -1,503 +0,0 @@
-# -*- 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()

fastreid/engine/launch.py

@@ -1,103 +0,0 @@
-# encoding: utf-8
-"""
-@author:  xingyu liao
-@contact: sherlockliao01@gmail.com
-"""
-
-# based on:
-# https://github.com/facebookresearch/detectron2/blob/master/detectron2/engine/launch.py
-
-
-import logging
-
-import torch
-import torch.distributed as dist
-import torch.multiprocessing as mp
-
-from fastreid.utils import comm
-
-__all__ = ["launch"]
-
-
-def _find_free_port():
-    import socket
-
-    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-    # Binding to port 0 will cause the OS to find an available port for us
-    sock.bind(("", 0))
-    port = sock.getsockname()[1]
-    sock.close()
-    # NOTE: there is still a chance the port could be taken by other processes.
-    return port
-
-
-def launch(main_func, num_gpus_per_machine, num_machines=1, machine_rank=0, dist_url=None, args=()):
-    """
-    Launch multi-gpu or distributed training.
-    This function must be called on all machines involved in the training.
-    It will spawn child processes (defined by ``num_gpus_per_machine`) on each machine.
-    Args:
-        main_func: a function that will be called by `main_func(*args)`
-        num_gpus_per_machine (int): number of GPUs per machine
-        num_machines (int): the total number of machines
-        machine_rank (int): the rank of this machine
-        dist_url (str): url to connect to for distributed jobs, including protocol
-                       e.g. "tcp://127.0.0.1:8686".
-                       Can be set to "auto" to automatically select a free port on localhost
-        args (tuple): arguments passed to main_func
-    """
-    world_size = num_machines * num_gpus_per_machine
-    if world_size > 1:
-        # https://github.com/pytorch/pytorch/pull/14391
-        # TODO prctl in spawned processes
-
-        if dist_url == "auto":
-            assert num_machines == 1, "dist_url=auto not supported in multi-machine jobs."
-            port = _find_free_port()
-            dist_url = f"tcp://127.0.0.1:{port}"
-        if num_machines > 1 and dist_url.startswith("file://"):
-            logger = logging.getLogger(__name__)
-            logger.warning(
-                "file:// is not a reliable init_method in multi-machine jobs. Prefer tcp://"
-            )
-
-        mp.spawn(
-            _distributed_worker,
-            nprocs=num_gpus_per_machine,
-            args=(main_func, world_size, num_gpus_per_machine, machine_rank, dist_url, args),
-            daemon=False,
-        )
-    else:
-        main_func(*args)
-
-
-def _distributed_worker(
-        local_rank, main_func, world_size, num_gpus_per_machine, machine_rank, dist_url, args
-):
-    assert torch.cuda.is_available(), "cuda is not available. Please check your installation."
-    global_rank = machine_rank * num_gpus_per_machine + local_rank
-    try:
-        dist.init_process_group(
-            backend="NCCL", init_method=dist_url, world_size=world_size, rank=global_rank
-        )
-    except Exception as e:
-        logger = logging.getLogger(__name__)
-        logger.error("Process group URL: {}".format(dist_url))
-        raise e
-    # synchronize is needed here to prevent a possible timeout after calling init_process_group
-    # See: https://github.com/facebookresearch/maskrcnn-benchmark/issues/172
-    comm.synchronize()
-
-    assert num_gpus_per_machine <= torch.cuda.device_count()
-    torch.cuda.set_device(local_rank)
-
-    # Setup the local process group (which contains ranks within the same machine)
-    assert comm._LOCAL_PROCESS_GROUP is None
-    num_machines = world_size // num_gpus_per_machine
-    for i in range(num_machines):
-        ranks_on_i = list(range(i * num_gpus_per_machine, (i + 1) * num_gpus_per_machine))
-        pg = dist.new_group(ranks_on_i)
-        if i == machine_rank:
-            comm._LOCAL_PROCESS_GROUP = pg
-
-    main_func(*args)