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)
