train Customer data; will support mult-object tracking soon

Author: zhangming8

.gitignore

@@ -1,4 +1,6 @@
 
+#!demo/*.jpg
+
 *.so
 *.o
 .DS_Store

README.md

@@ -1,31 +1,54 @@
-# a Pytorch easy re-implement of "YOLOX: Exceeding YOLO Series in 2021"
+## a Pytorch easy re-implement of "YOLOX: Exceeding YOLO Series in 2021"
 
+## Notes
+    1. this is a Pytorch easy re-implement of "YOLOX: Exceeding YOLO Series in 2021"
+    2. the repo is still under development
+    3. we needn't install apex, Pytorch(version >= 1.7.0) has supported it
 
-# environment
-    pytorch>=1.7.0, python>=3.6
+## Environment
+    pytorch>=1.7.0, python>=3.6, Ubuntu/Windows, see more in 'requirements.txt'
+
+## Dataset
+    put COCO dataset in following folders:
 
-# dataset
-    put your COCO dataset in fllowing folder:
-    
     /path/to/dataset/annotations/instances_train2017.json
     /path/to/dataset/annotations/instances_val2017.json
     /path/to/dataset/images/train2017/*.jpg
     /path/to/dataset/images/val2017/*.jpg
 
-    modify 'config.py'
-    opt.dataset_path = "/path/to/dataset"
+    change opt.dataset_path = "/path/to/dataset" in 'config.py'
 
-# train
+## Train
     sh train.sh
 
-# evaluate
+## Evaluate
     sh evaluate.sh
 
-# predict/inference/demo
+## Predict/Inference/Demo
     sh predict.sh
-   
 
-# reference
+## Train Customer Dataset(VOC format)
+    
+    1. put your annotations(.xml) and images(.jpg) into:
+        /path/to/voc_data/images/train2017/*.jpg  # train images
+        /path/to/voc_data/images/train2017/*.xml  # train xml annotations
+        /path/to/voc_data/images/val2017/*.jpg  # val images
+        /path/to/voc_data/images/val2017/*.xml  # val xml annotations
+
+    2. change opt.label_name = ['your', 'dataset', 'label'] in 'config.py'
+       change opt.dataset_path = '/path/to/voc_data' in 'config.py'
+
+    3. python tools/voc_to_coco.py
+       Converted COCO format annotation will be saved into:
+        /path/to/voc_data/annotations/instances_train2017.json
+        /path/to/voc_data/annotations/instances_val2017.json
+    
+    4. (Optional) you can visualize the converted annotations by:
+        python tools/show_coco_anns.py
+    
+    5. run train.sh, evaluate.sh, predict.sh (are the same as COCO)
+
+## Reference
     https://github.com/Megvii-BaseDetection/YOLOX
     https://github.com/PaddlePaddle/PaddleDetection
     https://github.com/open-mmlab/mmdetection

config.py

@@ -34,9 +34,25 @@ def update_nano_tiny(cfg):
 opt.random_size = (14, 26)  # None
 opt.accumulate = 1  # real batch size = accumulate * batch_size
 
+# coco 80 classes
+opt.label_name = [
+    'person', 'bicycle', 'car', 'motorcycle', 'airplane',
+    'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
+    'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
+    'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack',
+    'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
+    'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
+    'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass',
+    'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
+    'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake',
+    'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv',
+    'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
+    'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase',
+    'scissors', 'teddy bear', 'hair drier', 'toothbrush']
+# opt.label_name = ['person']
 # TODO: support MOT(multi-object tracking) like FairMot/JDE when reid_dim > 0
 opt.reid_dim = 0  # 128
-opt.id_num = None  # tracking id number in train dataset
+opt.tracking_id_nums = None  # tracking id number in train dataset
 
 opt.warmup_lr = 0
 opt.basic_lr_per_img = 0.01 / 64.0
@@ -72,20 +88,6 @@ def update_nano_tiny(cfg):
 opt.cuda_benchmark = True
 opt.nms_thresh = 0.65
 
-opt.label_name = [
-    'person', 'bicycle', 'car', 'motorcycle', 'airplane',
-    'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
-    'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
-    'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack',
-    'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
-    'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
-    'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass',
-    'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
-    'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake',
-    'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv',
-    'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
-    'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase',
-    'scissors', 'teddy bear', 'hair drier', 'toothbrush']
 opt.rgb_means = [0.485, 0.456, 0.406]
 opt.std = [0.229, 0.224, 0.225]
 

data/__init__.py

@@ -2,7 +2,7 @@
 # -*- coding:utf-8 -*-
 # Copyright (c) Megvii, Inc. and its affiliates.
 
-from .data_augment import TrainTransform, ValTransform
+from .data_augment import TrainTransform
 from .data_prefetcher import DataPrefetcher
 from .dataloading import DataLoader, get_yolox_datadir
 from .datasets import *

data/coco_dataset.py

@@ -9,28 +9,30 @@
 import sys
 
 sys.path.append(".")
-from data import (COCODataset, TrainTransform, YoloBatchSampler, DataLoader, InfiniteSampler, MosaicDetection,
-                  ValTransform)
+from data import (COCODataset, TrainTransform, YoloBatchSampler, DataLoader, InfiniteSampler, MosaicDetection)
 
 
 def get_dataloader(opt, no_aug=False):
     # train
+    do_tracking = opt.reid_dim > 0
     train_dataset = COCODataset(data_dir=opt.data_dir,
                                 json_file=opt.train_ann,
                                 img_size=opt.input_size,
-                                preproc=TrainTransform(rgb_means=opt.rgb_means, std=opt.std, max_labels=50),
+                                tracking=do_tracking,
+                                preproc=TrainTransform(rgb_means=opt.rgb_means, std=opt.std, tracking=do_tracking),
                                 )
     train_dataset = MosaicDetection(
         train_dataset,
         mosaic=not no_aug,
         img_size=opt.input_size,
-        preproc=TrainTransform(rgb_means=opt.rgb_means, std=opt.std, max_labels=120),
+        preproc=TrainTransform(rgb_means=opt.rgb_means, std=opt.std, max_labels=120, tracking=do_tracking),
         degrees=opt.degrees,
         translate=opt.translate,
         scale=opt.scale,
         shear=opt.shear,
         perspective=opt.perspective,
         enable_mixup=opt.enable_mixup,
+        tracking=do_tracking,
     )
     train_sampler = InfiniteSampler(len(train_dataset), seed=opt.seed)
     batch_sampler = YoloBatchSampler(
@@ -49,7 +51,9 @@ def get_dataloader(opt, no_aug=False):
         json_file=opt.val_ann,
         name="val2017",
         img_size=opt.test_size,
-        preproc=ValTransform(rgb_means=opt.rgb_means, std=opt.std))
+        tracking=do_tracking,
+        preproc=TrainTransform(rgb_means=opt.rgb_means, std=opt.std, max_labels=120, tracking=do_tracking,
+                               augment=False))
     val_sampler = torch.utils.data.SequentialSampler(val_dataset)
     val_kwargs = {"num_workers": opt.data_num_workers, "pin_memory": True, "sampler": val_sampler,
                   "batch_size": opt.batch_size}
@@ -68,24 +72,33 @@ def vis_inputs(inputs, targets, opt):
         img = (((inp.transpose((1, 2, 0)) * opt.std) + opt.rgb_means) * 255).astype(np.uint8)
         img = img[:, :, ::-1]
         img = np.ascontiguousarray(img)
+        gt_n = 0
         for t in target:
             if t.sum() > 0:
-                cls, c_x, c_y, w, h = [int(i) for i in t]
+                if len(t) == 5:
+                    cls, c_x, c_y, w, h = [int(i) for i in t]
+                    tracking_id = None
+                elif len(t) == 6:
+                    cls, c_x, c_y, w, h, tracking_id = [int(i) for i in t]
+                else:
+                    raise ValueError("target shape != 5 or 6")
                 bbox = [c_x - w // 2, c_y - h // 2, c_x + w // 2, c_y + h // 2]
                 label = opt.label_name[cls]
                 # print(label, bbox)
                 color = label_color[cls]
                 # show box
                 cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 2)
                 # show label and conf
-                txt = '{}'.format(label)
+                txt = '{}-{}'.format(label, tracking_id) if tracking_id is not None else '{}'.format(label)
                 font = cv2.FONT_HERSHEY_SIMPLEX
                 txt_size = cv2.getTextSize(txt, font, 0.5, 2)[0]
                 cv2.rectangle(img, (bbox[0], bbox[1] - txt_size[1] - 2), (bbox[0] + txt_size[0], bbox[1] - 2), color,
                               -1)
                 cv2.putText(img, txt, (bbox[0], bbox[1] - 2), font, 0.5, (255, 255, 255), thickness=1,
                             lineType=cv2.LINE_AA)
+                gt_n += 1
 
+        print("img {}/{} gt number: {}".format(b_i, len(inputs), gt_n))
         cv2.namedWindow("input", 0)
         cv2.imshow("input", img)
         key = cv2.waitKey(0)
@@ -99,7 +112,7 @@ def run_epoch(data_iter, loader, total_iter, e, phase, opt):
         batch = next(data_iter)
         inps, targets, img_info, ind = batch
         print("------------ epoch {} batch {}/{} ---------------".format(e, batch_i, total_iter))
-        print(inps.shape, targets.shape)
+        print("batch img shape {}, target shape {}".format(inps.shape, targets.shape))
         vis_inputs(inps, targets, opt)
         if batch_i == 0:
             print(ind)
@@ -119,8 +132,10 @@ def main():
     from config import opt
 
     opt.input_size = (640, 640)
+    opt.test_size = (640, 640)
     opt.batch_size = 2
     opt.data_num_workers = 0
+    opt.reid_dim = 0  # 128
     print(opt)
     train_loader, val_loader = get_dataloader(opt, no_aug=False)
 
@@ -132,8 +147,8 @@ def main():
     total_iter = len(loader)
     data_iter = iter(loader)
     for e in range(100):
-        # train_loader.dataset.enable_mosaic = False
         # train_loader.dataset.enable_mixup = False
+        # train_loader.dataset.enable_mosaic = False
         # train_loader.close_mosaic()
         # print(train_loader.batch_sampler.mosaic)
         run_epoch(data_iter, loader, total_iter, e, phase, opt)