if random.random() < 0.5:
img, boxes = augment.random_rotation(img, boxes, degree=5)
if random.random() < 0.5:
img, boxes = augment.random_crop_resize(img, boxes, size=512,
crop_scale_min=0.2, aspect_ratio=[3./4, 4./3], remain_min=0.7,
attempt_max=10)
return img, boxes
if cfg['augmentation']:
augmentation = aug_func
else:
augmentation = None
# Get train/eval dataset and dataloader
dataset_train = Dataset_CSV(cfg['root_train'], cfg['list_train'], cfg['name_file'],
size=net.module.view_size, train=True, normalize=True,
boxarea_th = cfg['boxarea_th'],
img_scale_min = cfg['img_scale_min'], augmentation=augmentation)
dataset_eval = Dataset_CSV(cfg['root_eval'], cfg['list_eval'], cfg['name_file'],
size=net.module.view_size, train=False, normalize=True)
loader_train = torch.utils.data.DataLoader(dataset_train, batch_size=cfg['nbatch_train'],
shuffle=True, num_workers=cfg['num_workers'], collate_fn=dataset_train.collate_fn)
loader_eval = torch.utils.data.DataLoader(dataset_eval, batch_size=cfg['nbatch_eval'],
shuffle=False, num_workers=0, collate_fn=dataset_eval.collate_fn)
# Prepare optimizer
lr_base = cfg['lr_base']
lr_gamma = cfg['lr_gamma']
lr_schedule = cfg['lr_schedule']
# opt = torch.optim.SGD(net.parameters(), lr=lr_base,
# momentum=cfg['momentum'], weight_decay=cfg['weight_decay'])
# Prepare the network
device_out = 'cuda:%d' % (cfg['device'][0])
net.load_state_dict(torch.load('net.pkl', map_location=device_out))
net = torch.nn.DataParallel(net, device_ids=cfg['device'])
net = net.cuda(cfg['device'][0])
net.eval()
# Get eval dataset and dataloader
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
dataset_eval = Dataset_CSV(cfg['root_eval'],
cfg['list_eval'],
cfg['name_file'],
size=net.module.view_size,
train=False,
transform=transform)
loader_eval = torch.utils.data.DataLoader(dataset_eval,
batch_size=cfg['nbatch_eval'],
shuffle=False,
num_workers=0,
collate_fn=dataset_eval.collate_fn)
# Eval
with torch.no_grad():
net.eval()
pred_bboxes = []
pred_labels = []
pred_scores = []
gt_bboxes = []
coco_table = json.load(load_f)
with open('train.json', 'r') as load_f:
cfg = json.load(load_f)
torch.cuda.set_device(cfg['device'][0])
# Prepare the network
net = Detector(pretrained=False)
device_out = 'cuda:%d' % (cfg['device'][0])
net.load_state_dict(torch.load('net.pkl', map_location=device_out))
net = net.cuda(cfg['device'][0])
net.eval()
# Get eval dataset
dataset_eval = Dataset_CSV(cfg['root_eval'],
cfg['list_eval'],
cfg['name_file'],
size=net.view_size,
train=False,
normalize=True)
# Prepare API structure
inferencer = api.COCOEvaluator(net, dataset_eval, coco_table['val_image_ids'],
coco_table['coco_labels'])
# Eval
inferencer.step_epoch()
coco = COCO(
os.path.join(coco_anno_root, 'annotations',
'instances_' + set_name + '.json'))
coco_pred = coco.loadRes('coco_bbox_results.json')
coco_eval = COCOeval(coco, coco_pred, 'bbox')
coco_eval.params.imgIds = coco_table['val_image_ids']
if random.random() < 0.5:
img, boxes = augment.random_rotation(img, boxes, degree=5)
if random.random() < 0.5:
img, boxes = augment.random_crop_resize(img, boxes, size=512,
crop_scale_min=0.2, aspect_ratio=[3./4, 4./3], remain_min=0.7,
attempt_max=10)
return img, boxes
if cfg['augmentation']:
augmentation = aug_func
else:
augmentation = None
# Get train/eval dataset and dataloader
dataset_train = Dataset_CSV(cfg['root_train'], cfg['list_train'], cfg['name_file'], cfg['frame_name_file'],
size=513, train=True, normalize=True,
boxarea_th = cfg['boxarea_th'],
img_scale_min = cfg['img_scale_min'], augmentation=augmentation, mosaic=cfg['mosaic'])
dataset_eval = Dataset_CSV(cfg['root_eval'], cfg['list_eval'], cfg['name_file'], cfg['frame_name_file'],
size=513, train=False, normalize=True, mosaic=cfg['mosaic'])
loader_train = torch.utils.data.DataLoader(dataset_train, batch_size=cfg['nbatch_train'],
shuffle=True, num_workers=cfg['num_workers'], collate_fn=dataset_train.collate_fn)
loader_eval = torch.utils.data.DataLoader(dataset_eval, batch_size=cfg['nbatch_eval'],
shuffle=False, num_workers=cfg['num_workers'], collate_fn=dataset_eval.collate_fn)
# Prepare the network and read log
net = Detector(frame_weights=dataset_train.weights, pretrained=cfg['pretrain'])
log = []
device_out = 'cuda:%d' % (cfg['device'][0])
if cfg['load']:
net.load_state_dict(torch.load('net.pkl', map_location=device_out))
if cfg['load']:
net.load_state_dict(torch.load('net.pkl', map_location=device_out))
log = list(np.load('log.npy'))
net = torch.nn.DataParallel(net, device_ids=cfg['device'])
net = net.cuda(cfg['device'][0])
net.train()
# Get train/eval dataset and dataloader
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485,0.456,0.406), (0.229,0.224,0.225))])
dataset_train = Dataset_CSV(cfg['root_train'], cfg['list_train'], cfg['name_file'],
size=net.module.view_size, train=True, transform=transform,
boxarea_th = cfg['boxarea_th'],
img_scale_min = cfg['img_scale_min'],
crop_scale_min = cfg['crop_scale_min'],
aspect_ratio = cfg['aspect_ratio'],
remain_min = cfg['remain_min'],
augmentation = cfg['augmentation'])
dataset_eval = Dataset_CSV(cfg['root_eval'], cfg['list_eval'], cfg['name_file'],
size=net.module.view_size, train=False, transform=transform)
loader_train = torch.utils.data.DataLoader(dataset_train, batch_size=cfg['nbatch_train'],
shuffle=True, num_workers=cfg['num_workers'], collate_fn=dataset_train.collate_fn)
loader_eval = torch.utils.data.DataLoader(dataset_eval, batch_size=cfg['nbatch_eval'],
shuffle=False, num_workers=0, collate_fn=dataset_eval.collate_fn)
# Prepare optimizer
lr = cfg['lr']
lr_decay = cfg['lr_decay']
opt = torch.optim.SGD(net.parameters(), lr=lr,
