def load_model(num_classes):
print('loading model...')
dir_weight = os.path.join(FLAGS.weights) #later replace for FLAGS.weigths
prior_box = {
'strides': [8, 16, 32, 64, 128],
'sizes': [3] * 5,
'aspects': [[1, 2, 4, 8]] * 5,
'scales': [[2**0, 2**(1 / 3), 2**(2 / 3)]] * 5,
'old_version': False
}
cfg = {
'prior_box': prior_box,
'num_classes': num_classes,
'extra': 2,
'conf_thresh': FLAGS.conf_thresh,
'nms_thresh': FLAGS.nms_thresh
}
backbone = resnet.resnet101
model = RDD(backbone(fetch_feature=True), cfg)
model.build_pipe(shape=[2, 3, FLAGS.image_size, FLAGS.image_size])
model.restore(dir_weight)
device_ids = [int(gpu_id) for gpu_id in FLAGS.GPUS.split(',')]
if len(device_ids) > 1:
model = CustomDetDataParallel(model, device_ids)
model.cuda()
model.eval()
print('model restored')
return model
def main():
dir_weight = os.path.join(dir_save, 'weight')
dir_log = os.path.join(dir_save, 'log')
os.makedirs(dir_weight, exist_ok=True)
writer = SummaryWriter(dir_log)
indexes = [
int(os.path.splitext(path)[0]) for path in os.listdir(dir_weight)
]
current_step = max(indexes) if indexes else 0
image_size = 768
lr = 1e-3
batch_size = 12
num_workers = 4
max_step = 250000
lr_cfg = [[100000, lr], [200000, lr / 10], [max_step, lr / 50]]
warm_up = [1000, lr / 50, lr]
save_interval = 1000
aug = Compose([
ops.ToFloat(),
ops.PhotometricDistort(),
ops.RandomHFlip(),
ops.RandomVFlip(),
ops.RandomRotate90(),
ops.ResizeJitter([0.8, 1.2]),
ops.PadSquare(),
ops.Resize(image_size),
ops.BBoxFilter(24 * 24 * 0.4)
])
dataset = DOTA(dir_dataset, ['train', 'val'], aug)
loader = DataLoader(dataset,
batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=True,
drop_last=True,
collate_fn=dataset.collate)
num_classes = len(dataset.names)
prior_box = {
'strides': [8, 16, 32, 64, 128],
'sizes': [3] * 5,
'aspects': [[1, 2, 4, 8]] * 5,
'scales': [[2**0, 2**(1 / 3), 2**(2 / 3)]] * 5,
}
cfg = {
'prior_box': prior_box,
'num_classes': num_classes,
'extra': 2,
}
model = RDD(backbone(fetch_feature=True), cfg)
model.build_pipe(shape=[2, 3, image_size, image_size])
if current_step:
model.restore(os.path.join(dir_weight, '%d.pth' % current_step))
else:
model.init()
if len(device_ids) > 1:
model = convert_model(model)
model = CustomDetDataParallel(model, device_ids)
model.cuda()
optimizer = optim.SGD(model.parameters(),
lr=lr,
momentum=0.9,
weight_decay=5e-4)
training = True
while training and current_step < max_step:
tqdm_loader = tqdm.tqdm(loader)
for images, targets, infos in tqdm_loader:
current_step += 1
adjust_lr_multi_step(optimizer, current_step, lr_cfg, warm_up)
images = images.cuda() / 255
losses = model(images, targets)
loss = sum(losses.values())
loss.backward()
optimizer.step()
optimizer.zero_grad()
for key, val in list(losses.items()):
losses[key] = val.item()
writer.add_scalar(key, val, global_step=current_step)
writer.flush()
tqdm_loader.set_postfix(losses)
tqdm_loader.set_description(f'<{current_step}/{max_step}>')
if current_step % save_interval == 0:
save_path = os.path.join(dir_weight, '%d.pth' % current_step)
state_dict = model.state_dict() if len(
device_ids) == 1 else model.module.state_dict()
torch.save(state_dict, save_path)
cache_file = os.path.join(
dir_weight, '%d.pth' % (current_step - save_interval))
if os.path.exists(cache_file):
os.remove(cache_file)
if current_step >= max_step:
training = False
writer.close()
break
def main():
global checkpoint
if checkpoint is None:
dir_weight = os.path.join(dir_save, 'weight')
indexes = [
int(os.path.splitext(path)[0]) for path in os.listdir(dir_weight)
]
current_step = max(indexes)
checkpoint = os.path.join(dir_weight, '%d.pth' % current_step)
batch_size = 32
num_workers = 4
image_size = 768
aug = Compose([ops.PadSquare(), ops.Resize(image_size)])
dataset = DOTA(dir_dataset, image_set, aug)
loader = DataLoader(dataset,
batch_size,
num_workers=num_workers,
pin_memory=True,
collate_fn=dataset.collate)
num_classes = len(dataset.names)
prior_box = {
'strides': [8, 16, 32, 64, 128],
'sizes': [3] * 5,
'aspects': [[1, 2, 4, 8]] * 5,
'scales': [[2**0, 2**(1 / 3), 2**(2 / 3)]] * 5,
'old_version': old_version
}
conf_thresh = 0.01
nms_thresh = 0.45
cfg = {
'prior_box': prior_box,
'num_classes': num_classes,
'extra': 2,
'conf_thresh': conf_thresh,
'nms_thresh': nms_thresh,
}
model = RDD(backbone(fetch_feature=True), cfg)
model.build_pipe(shape=[2, 3, image_size, image_size])
model.restore(checkpoint)
if len(device_ids) > 1:
model = CustomDetDataParallel(model, device_ids)
model.cuda()
model.eval()
ret_raw = defaultdict(list)
for images, targets, infos in tqdm.tqdm(loader):
images = images.cuda() / 255
dets = model(images)
for (det, info) in zip(dets, infos):
if det:
bboxes, scores, labels = det
bboxes = bboxes.cpu().numpy()
scores = scores.cpu().numpy()
labels = labels.cpu().numpy()
fname, x, y, w, h = os.path.splitext(
os.path.basename(info['img_path']))[0].split('-')[:5]
x, y, w, h = int(x), int(y), int(w), int(h)
long_edge = max(w, h)
pad_x, pad_y = (long_edge - w) // 2, (long_edge - h) // 2
bboxes = np.stack([xywha2xy4(bbox) for bbox in bboxes])
bboxes *= long_edge / image_size
bboxes -= [pad_x, pad_y]
bboxes += [x, y]
bboxes = np.stack([xy42xywha(bbox) for bbox in bboxes])
ret_raw[fname].append([bboxes, scores, labels])
print('merging results...')
ret = []
for fname, dets in ret_raw.items():
bboxes, scores, labels = zip(*dets)
bboxes = np.concatenate(list(bboxes))
scores = np.concatenate(list(scores))
labels = np.concatenate(list(labels))
keeps = rbbox_batched_nms(bboxes, scores, labels, nms_thresh)
ret.append([fname, [bboxes[keeps], scores[keeps], labels[keeps]]])
print('converting to submission format...')
ret_save = defaultdict(list)
for fname, (bboxes, scores, labels) in ret:
for bbox, score, label in zip(bboxes, scores, labels):
bbox = xywha2xy4(bbox).ravel()
line = '%s %.12f %.1f %.1f %.1f %.1f %.1f %.1f %.1f %.1f' % (
fname, score, *bbox)
ret_save[dataset.label2name[label]].append(line)
print('saving...')
os.makedirs(os.path.join(dir_save, 'submission'), exist_ok=True)
for name, dets in ret_save.items():
with open(
os.path.join(dir_save, 'submission',
'Task%d_%s.txt' % (1, name)), 'wt') as f:
f.write('\n'.join(dets))
print('finished')
def main():
global checkpoint
if checkpoint is None:
dir_weight = os.path.join(dir_save, 'weight')
indexes = [int(os.path.splitext(path)[0]) for path in os.listdir(dir_weight)]
current_step = max(indexes)
checkpoint = os.path.join(dir_weight, '%d.pth' % current_step)
image_size = 768
batch_size = 32
num_workers = 4
aug = ops.Resize(image_size)
dataset = HRSC2016(dir_dataset, 'test', aug)
loader = DataLoader(dataset, batch_size, num_workers=num_workers, pin_memory=True, collate_fn=dataset.collate)
num_classes = len(dataset.names)
prior_box = {
'strides': [8, 16, 32, 64, 128],
'sizes': [3] * 5,
'aspects': [[1.5, 3, 5, 8]] * 5,
'scales': [[2 ** 0, 2 ** (1 / 3), 2 ** (2 / 3)]] * 5,
'old_version': old_version
}
conf_thresh = 0.01
nms_thresh = 0.45
cfg = {
'prior_box': prior_box,
'num_classes': num_classes,
'extra': 2,
'conf_thresh': conf_thresh,
'nms_thresh': nms_thresh,
}
model = RDD(backbone(fetch_feature=True), cfg)
model.build_pipe(shape=[2, 3, image_size, image_size])
model.restore(checkpoint)
if len(device_ids) > 1:
model = CustomDetDataParallel(model, device_ids)
model.cuda()
model.eval()
count = 0
gt_list, det_list = [], []
for images, targets, infos in tqdm.tqdm(loader):
images = images.cuda() / 255
dets = model(images)
for target, det, info in zip(targets, dets, infos):
if target:
bboxes = np.stack([xy42xywha(bbox) for bbox in info['objs']['bboxes']])
labels = info['objs']['labels']
gt_list.extend([count, bbox, 1, label] for bbox, label in zip(bboxes, labels))
if det:
ih, iw = info['shape'][:2]
bboxes, scores, labels = list(map(lambda x: x.cpu().numpy(), det))
bboxes = np.stack([xywha2xy4(bbox) for bbox in bboxes])
bboxes_ = bboxes * [iw / image_size, ih / image_size]
# bboxes = np.stack([xy42xywha(bbox) for bbox in bboxes_])
bboxes = []
for bbox in bboxes_.astype(np.float32):
(x, y), (w, h), a = cv.minAreaRect(bbox)
bboxes.append([x, y, w, h, a])
bboxes = np.array(bboxes)
det_list.extend([count, bbox, score, label] for bbox, score, label in zip(bboxes, scores, labels))
count += 1
APs = get_det_aps(det_list, gt_list, num_classes, use_07_metric=use_07_metric)
mAP = sum(APs) / len(APs)
print('AP')
for label in range(num_classes):
print(f'{dataset.label2name[label]}: {APs[label]}')
print(f'mAP: {mAP}')