|
import torch |
|
import numpy as np |
|
import cv2 |
|
|
|
|
|
def building_metric(logits, label): |
|
preds = torch.argmax(logits, dim=-1) |
|
true_ids = torch.where(label==1) |
|
num_true = true_ids[0].shape[0] |
|
tp = (preds[true_ids] == 1).sum().double() |
|
recall = tp / num_true |
|
prec = tp / (preds == 1).sum() |
|
fscore = 2 * recall * prec / (prec + recall) |
|
return recall, prec, fscore |
|
|
|
|
|
def edge_acc(logits, label, lengths, gt_values): |
|
""" |
|
edge f1-score for training/validation logging |
|
""" |
|
all_acc = list() |
|
for i in range(logits.shape[0]): |
|
length = lengths[i] |
|
gt_value = gt_values[i, :length] |
|
pred_idx = torch.where(gt_value == 2) |
|
if len(pred_idx[0]) == 0: |
|
continue |
|
else: |
|
preds = torch.argmax(logits[i, :, :length][:, pred_idx[0]], dim=0) |
|
gts = label[i, :length][pred_idx[0]] |
|
pos_ids = torch.where(gts == 1) |
|
correct = (preds[pos_ids] == gts[pos_ids]).sum().float() |
|
num_pos_gt = len(pos_ids[0]) |
|
recall = correct / num_pos_gt if num_pos_gt > 0 else torch.tensor(0) |
|
num_pos_pred = (preds == 1).sum().float() |
|
prec = correct / num_pos_pred if num_pos_pred > 0 else torch.tensor(0) |
|
f_score = 2.0 * prec * recall / (recall + prec + 1e-8) |
|
f_score = f_score.cpu() |
|
all_acc.append(f_score) |
|
if len(all_acc) > 1: |
|
all_acc = torch.stack(all_acc, 0) |
|
avg_acc = all_acc.mean() |
|
else: |
|
avg_acc = all_acc[0] |
|
return avg_acc |
|
|
|
|
|
def corner_eval(targets, outputs): |
|
assert isinstance(targets, np.ndarray) |
|
assert isinstance(outputs, np.ndarray) |
|
output_to_gt = dict() |
|
gt_to_output = dict() |
|
for target_i, target in enumerate(targets): |
|
dist = (outputs - target) ** 2 |
|
dist = np.sqrt(dist.sum(axis=-1)) |
|
min_dist = dist.min() |
|
min_idx = dist.argmin() |
|
if min_dist < 5 and min_idx not in output_to_gt: |
|
output_to_gt[min_idx] = target_i |
|
gt_to_output[target_i] = min_idx |
|
tp = len(output_to_gt) |
|
prec = tp / len(outputs) |
|
recall = tp / len(targets) |
|
return prec, recall |
|
|
|
|
|
def rectify_data(image, annot): |
|
rows, cols, ch = image.shape |
|
bins = [0 for _ in range(180)] |
|
|
|
|
|
gauss_weights = [0.1, 0.2, 0.5, 1, 0.5, 0.2, 0.1] |
|
|
|
for src, connections in annot.items(): |
|
for end in connections: |
|
edge = [(end[0] - src[0]), -(end[1] - src[1])] |
|
edge_len = np.sqrt(edge[0] ** 2 + edge[1] ** 2) |
|
if edge_len <= 10: |
|
continue |
|
if edge[0] == 0: |
|
bin_id = 90 |
|
else: |
|
theta = np.arctan(edge[1] / edge[0]) / np.pi * 180 |
|
if edge[0] * edge[1] < 0: |
|
theta += 180 |
|
bin_id = int(theta.round()) |
|
if bin_id == 180: |
|
bin_id = 0 |
|
for offset in range(-3, 4): |
|
bin_idx = bin_id + offset |
|
if bin_idx >= 180: |
|
bin_idx -= 180 |
|
bins[bin_idx] += np.sqrt(edge[1] ** 2 + edge[0] ** 2) * gauss_weights[offset + 2] |
|
|
|
bins = np.array(bins) |
|
sorted_ids = np.argsort(bins)[::-1] |
|
bin_1 = sorted_ids[0] |
|
remained_ids = [idx for idx in sorted_ids if angle_dist(bin_1, idx) >= 30] |
|
bin_2 = remained_ids[0] |
|
if bin_1 < bin_2: |
|
bin_1, bin_2 = bin_2, bin_1 |
|
|
|
dir_1, dir_2 = bin_1, bin_2 |
|
|
|
origin = [127, 127] |
|
p1_old = [127 + 100 * np.cos(dir_1 / 180 * np.pi), 127 - 100 * np.sin(dir_1 / 180 * np.pi)] |
|
p2_old = [127 + 100 * np.cos(dir_2 / 180 * np.pi), 127 - 100 * np.sin(dir_2 / 180 * np.pi)] |
|
pts1 = np.array([origin, p1_old, p2_old]).astype(np.float32) |
|
p1_new = [127, 27] |
|
p2_new = [227, 127] |
|
pts2 = np.array([origin, p1_new, p2_new]).astype(np.float32) |
|
|
|
M1 = cv2.getAffineTransform(pts1, pts2) |
|
|
|
all_corners = list(annot.keys()) |
|
all_corners_ = np.array(all_corners) |
|
ones = np.ones([all_corners_.shape[0], 1]) |
|
all_corners_ = np.concatenate([all_corners_, ones], axis=-1) |
|
new_corners = np.matmul(M1, all_corners_.T).T |
|
|
|
M = np.concatenate([M1, np.array([[0, 0, 1]])], axis=0) |
|
|
|
x_max = new_corners[:, 0].max() |
|
x_min = new_corners[:, 0].min() |
|
y_max = new_corners[:, 1].max() |
|
y_min = new_corners[:, 1].min() |
|
|
|
side_x = (x_max - x_min) * 0.1 |
|
side_y = (y_max - y_min) * 0.1 |
|
right_border = x_max + side_x |
|
left_border = x_min - side_x |
|
bot_border = y_max + side_y |
|
top_border = y_min - side_y |
|
pts1 = np.array([[left_border, top_border], [right_border, top_border], [right_border, bot_border]]).astype( |
|
np.float32) |
|
pts2 = np.array([[5, 5], [250, 5], [250, 250]]).astype(np.float32) |
|
M_scale = cv2.getAffineTransform(pts1, pts2) |
|
|
|
M = np.matmul(np.concatenate([M_scale, np.array([[0, 0, 1]])], axis=0), M) |
|
|
|
new_image = cv2.warpAffine(image, M[:2, :], (cols, rows), borderValue=(255, 255, 255)) |
|
all_corners_ = np.concatenate([all_corners, ones], axis=-1) |
|
new_corners = np.matmul(M[:2, :], all_corners_.T).T |
|
|
|
corner_mapping = dict() |
|
for idx, corner in enumerate(all_corners): |
|
corner_mapping[corner] = new_corners[idx] |
|
|
|
new_annot = dict() |
|
for corner, connections in annot.items(): |
|
new_corner = corner_mapping[corner] |
|
tuple_new_corner = tuple(new_corner) |
|
new_annot[tuple_new_corner] = list() |
|
for to_corner in connections: |
|
new_annot[tuple_new_corner].append(corner_mapping[tuple(to_corner)]) |
|
|
|
|
|
return new_image, new_annot, M |
|
|
|
|
|
def angle_dist(a1, a2): |
|
if a1 > a2: |
|
a1, a2 = a2, a1 |
|
d1 = a2 - a1 |
|
d2 = a1 + 180 - a2 |
|
dist = min(d1, d2) |
|
return dist |
|
|