# encoding: utf-8 import os import random import torch import torch.nn as nn import torch.distributed as dist from yolox.exp import MetaExp as MyMetaExp from yolox.data import get_yolox_datadir from os import listdir from os.path import isfile, join class Exp(MyMetaExp): def __init__(self): super(Exp, self).__init__() self.num_classes = 1 self.depth = 1.33 self.width = 1.25 self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] self.train_dir = '/home/abdollahpour.ce.sharif/ByteTrackData/MOT17/annotations' onlyfiles = [f for f in listdir(self.train_dir) if isfile(join(self.train_dir, f))] self.train_anns = [file for file in onlyfiles if file.__contains__('train') and file.__contains__('FRCNN')] # # TODO: remove # self.train_anns = self.train_anns[3:] self.val_dir = '/home/abdollahpour.ce.sharif/ByteTrackData/MOT20/annotations' onlyfiles = [f for f in listdir(self.val_dir) if isfile(join(self.val_dir, f))] self.val_anns = [file for file in onlyfiles if file.__contains__('train') and file.__contains__( 'MOT20')] # self.val_anns = self.val_anns[-1:] print('train_anns', self.train_anns) print('val_anns', self.val_anns) self.input_size = (800, 1440) # TODO: try this self.test_size = (800, 1440) # self.test_size = (896, 1600) # self.test_size = (736, 1920) self.random_size = (20, 36) self.max_epoch = 80 self.print_interval = 250 self.eval_interval = 10 self.test_conf = 0.001 self.nmsthre = 0.7 self.no_aug_epochs = 10 # self.basic_lr_per_img = 0.001 / 64.0 self.basic_lr_per_img = 0.0001 / 64.0 self.warmup_epochs = 1 def get_data_loaders(self, batch_size, is_distributed, no_aug=False): from yolox.data import ( MOTDataset, TrainTransform, YoloBatchSampler, DataLoader, InfiniteSampler, MosaicDetection, ) train_loaders = [] for train_ann in self.train_anns: dataset = MOTDataset( data_dir=os.path.join(get_yolox_datadir(), "MOT17"), json_file=train_ann, name='train', img_size=self.input_size, preproc=TrainTransform( rgb_means=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_labels=500, ), ) dataset = MosaicDetection( dataset, mosaic=not no_aug, img_size=self.input_size, preproc=TrainTransform( rgb_means=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_labels=1000, ), degrees=self.degrees, translate=self.translate, scale=self.scale, shear=self.shear, perspective=self.perspective, enable_mixup=self.enable_mixup, ) self.dataset = dataset if is_distributed: batch_size = batch_size // dist.get_world_size() sampler = InfiniteSampler( len(self.dataset), seed=self.seed if self.seed else 0 ) batch_sampler = YoloBatchSampler( sampler=sampler, batch_size=batch_size, drop_last=False, input_dimension=self.input_size, mosaic=not no_aug, ) dataloader_kwargs = {"num_workers": self.data_num_workers, "pin_memory": True} dataloader_kwargs["batch_sampler"] = batch_sampler train_loader = DataLoader(self.dataset, **dataloader_kwargs) train_loaders.append(train_loader) return train_loaders def get_eval_loaders(self, batch_size, is_distributed, testdev=False): from yolox.data import MOTDataset, ValTransform, ValTransformWithPseudo, TrainTransform val_loaders = [] for val_ann in self.val_anns: valdataset = MOTDataset( data_dir=os.path.join(get_yolox_datadir(), "MOT20"), json_file=val_ann, img_size=self.test_size, name='train', # change to train when running on training set preproc=ValTransformWithPseudo( rgb_means=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_labels=500, ), load_weak=False ) if is_distributed: batch_size = batch_size // dist.get_world_size() sampler = torch.utils.data.distributed.DistributedSampler( valdataset, shuffle=False ) else: sampler = torch.utils.data.SequentialSampler(valdataset) dataloader_kwargs = { "num_workers": self.data_num_workers, "pin_memory": True, "sampler": sampler, } dataloader_kwargs["batch_size"] = batch_size val_loader = torch.utils.data.DataLoader(valdataset, **dataloader_kwargs) val_loaders.append(val_loader) return val_loaders def get_evaluators(self, batch_size, is_distributed, testdev=False): from yolox.evaluators import COCOEvaluator val_loaders = self.get_eval_loaders(batch_size, is_distributed, testdev=testdev) evaluators = [] for val_loader in val_loaders: evaluator = COCOEvaluator( dataloader=val_loader, img_size=self.test_size, confthre=self.test_conf, nmsthre=self.nmsthre, num_classes=self.num_classes, testdev=testdev, ) evaluators.append(evaluator) return evaluators