def convert_objs(objs): target = dict() if objs: # Limit the angle between -45° and 45° by set flag=2 target['bboxes'] = torch.from_numpy(np.stack([xy42xywha(bbox, flag=2) for bbox in objs['bboxes']])).float() target['labels'] = torch.from_numpy(objs['labels']).long() return target
def __call__(self, img, anno=None): ih, iw = img.shape[:2] polygons = [] if anno: for bbox in anno['bboxes']: x, y, w, h, a = xy42xywha(bbox) polygons.append(((x, y), (w, h), a)) for count in range(self.max_try): if isinstance(self.size, int): nh = nw = min(ih, iw, self.size) else: if self.max_aspect == 1: nh = nw = random.randint(min(ih, iw, self.size[0]), min(ih, iw, self.size[1])) else: nh = random.randint(min(ih, self.size[0]), min(ih, self.size[1])) nw = random.randint(min(iw, self.size[0]), min(iw, self.size[1])) if max(nh / nw, nw / nh) > self.max_aspect: continue oh = random.randint(0, ih - nh) ow = random.randint(0, iw - nw) a = np.random.uniform(0, 360) src = xywha2xy4([ow + nw / 2, oh + nh / 2, nw, nh, a]) dst = np.array([[0, 0], [nw, 0], [nw, nh]], dtype=np.float32) m = cv.getAffineTransform(src.astype(np.float32)[:3], dst) if anno: bound = (ow + nw / 2, oh + nh / 2), (nw, nh), a iou, intersections = [], [] for polygon in polygons: inter_points = cv.rotatedRectangleIntersection(bound, polygon)[1] if inter_points is None: iou.append(0) intersections.append(None) else: order_points = cv.convexHull(inter_points, returnPoints=True) inter_area = cv.contourArea(order_points) iou.append(inter_area / (polygon[1][0] * polygon[1][1])) intersections.append(cv.boxPoints(cv.minAreaRect(order_points))) iou = np.array(iou) if isinstance(self.iou_thresh, float): mask = iou >= self.iou_thresh else: mask = (iou > self.iou_thresh[0]) & (iou < self.iou_thresh[1]) if np.any(mask): continue mask = iou >= self.iou_thresh[1] if np.any(mask): bboxes = np.array([inter for inter, m in zip(intersections, mask) if m]) bboxes = np.concatenate([bboxes, np.ones_like(bboxes[:, :, [0]])], axis=-1) bboxes = np.matmul(m, bboxes.transpose([0, 2, 1])).transpose([0, 2, 1]) anno['bboxes'] = bboxes anno['labels'] = anno['labels'][mask] else: if self.nonempty: continue else: anno = None img = cv.warpAffine(img, m, (nw, nh)) break return img, anno
def __call__(self, img, anno=None): if anno: wh = np.stack([xy42xywha(bbox)[2:4] for bbox in anno['bboxes']]) area = wh[:, 0] * wh[:, 1] mask = area >= self.min_area if np.any(mask): anno['bboxes'] = anno['bboxes'][mask] anno['labels'] = anno['labels'][mask] else: anno.clear() return img, anno
def __call__(self, img, anno=None): nh, nw = ih, iw = img.shape[:2] center = ((iw - 1) / 2, (ih - 1) / 2) m = cv.getRotationMatrix2D(center, angle=-self.angle, scale=self.scale) if self.expand or self.shift: corner = np.array([[0, 0, 1], [iw - 1, 0, 1], [iw - 1, ih - 1, 1], [0, ih - 1, 1]], dtype=np.float32) corner = np.matmul(m, corner.T).T left, top = np.min(corner, axis=0) right, bottom = np.max(corner, axis=0) dx = (right - left - iw) / 2 dy = (bottom - top - ih) / 2 if self.expand: nw = int(np.ceil(right - left)) nh = int(np.ceil(bottom - top)) shiftX = dx shiftY = dy else: shiftX = np.random.uniform(-dx, dx) if dx > 0 else 0 shiftY = np.random.uniform(-dy, dy) if dy > 0 else 0 m[0, 2] += shiftX m[1, 2] += shiftY if anno: bound = (nw / 2, nh / 2), (nw, nh), 0 bboxes, labels = [], [] for bbox, label in zip(anno['bboxes'], anno['labels']): corner = np.matmul(m, np.c_[bbox, np.ones((4, 1))].T).T if not self.expand: x, y, w, h, a = xy42xywha(corner) inter_points = cv.rotatedRectangleIntersection( bound, ((x, y), (w, h), a))[1] if inter_points is not None: order_points = cv.convexHull(inter_points, returnPoints=True) inter_area = cv.contourArea(order_points) iou = inter_area / (w * h) if iou >= 0.5: corner = cv.boxPoints(cv.minAreaRect(order_points)) else: continue bboxes.append(corner) labels.append(label) if bboxes: anno['bboxes'] = np.stack(bboxes) anno['labels'] = np.stack(labels) else: anno = None img = cv.warpAffine(img, m, (nw, nh)) return img, anno
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}')
def predict(model, data_loader, output_csv): print('predicting...') count_predicted, count_not_predicted = 0, 0 ret_raw = defaultdict(list) with torch.no_grad(): for images, targets, infos in tqdm.tqdm(data_loader): images = images.cuda() / 255 dets = model(images) for (det, info) in zip(dets, infos): fname = ntpath.basename(info['img_path']) if det: bboxes, scores, labels = det bboxes = bboxes.cpu().numpy() scores = scores.cpu().numpy() labels = labels.cpu().numpy() x, y, w, h = 0, 0, info['shape'][1], info['shape'][0] 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 / FLAGS.image_size bboxes -= [pad_x, pad_y] bboxes += [x, y] bboxes = np.stack( [xy42xywha(bbox, flag=1) for bbox in bboxes]) ret_raw[fname].append([bboxes, scores, labels]) count_predicted += 1 else: count_not_predicted += 1 ret_raw[fname].append( [np.zeros((1, 5)), np.zeros((1, )), np.zeros((1, ))]) print( f'{count_predicted} jpgs with at least one object found, {count_not_predicted} with no object found' ) print('merging results...') ret = [] for fname, dets in tqdm.tqdm(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, FLAGS.nms_thresh) ret.append([fname, [bboxes, scores[keeps], labels[keeps]]]) columns = [ 'img_name', 'x1', 'y1', 'x2', 'y2', 'x3', 'y3', 'x4', 'y4', 'angle_class', 'angle_class_degrees', 'angle_class_degrees_prob', 'cx', 'cy', 'w', 'h', 'bb_angle' ] df = pd.DataFrame(columns=columns) idx = 0 print('exporting csv...') for fname, (bboxes, scores, labels) in tqdm.tqdm(ret): for bbox, score, label in zip(bboxes, scores, labels): angle_step = 5 xy4 = xywha2xy4(bbox).ravel() row = [ fname, xy4[0], xy4[1], xy4[2], xy4[3], xy4[4], xy4[5], xy4[6], xy4[7], label, label * angle_step, score, bbox[0], bbox[1], bbox[2], bbox[3], bbox[4], ] df.loc[idx] = row idx += 1 df.to_csv(output_csv, sep=';') print('prediction results saved to', output_csv)