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}')