84 lines
2.4 KiB
Python
84 lines
2.4 KiB
Python
import numpy as np
|
|
from cv_analysis.utils.structures import Rectangle
|
|
|
|
|
|
def xyxy_from_object(box_object):
|
|
try:
|
|
x1, y1, x2, y2 = box_object.xyxy()
|
|
except:
|
|
try:
|
|
x1 = box_object["x"]
|
|
y1 = box_object["y"]
|
|
x2 = x1 + box_object["width"]
|
|
y2 = y1 + box_object["height"]
|
|
except:
|
|
x1, y1, x2, y2 = box_object
|
|
return x1, y1, x2, y2
|
|
|
|
|
|
def xywh_from_object(box_object):
|
|
try:
|
|
x, y, w, h = box_object.xywh()
|
|
except:
|
|
try:
|
|
x = box_object["x"]
|
|
y = box_object["y"]
|
|
w = box_object["width"]
|
|
h = box_object["height"]
|
|
except:
|
|
x, y, w, h = box_object
|
|
return x, y, w, h
|
|
|
|
|
|
def compute_iou_from_boxes(box1: Rectangle, box2: list):
|
|
"""
|
|
Each box of the form (x1, y1, delx, dely)
|
|
"""
|
|
ax1, ay1, aw, ah = xywh_from_object(box1)
|
|
bx1, by1, bw, bh = xywh_from_object(box2)
|
|
ax2, ay2, bx2, by2 = ax1 + aw, ay1 + ah, bx1 + bw, by1 + bh
|
|
if (ax1 > bx2) or (bx1 > ax2) or (ay1 > by2) or (by1 > ay2):
|
|
return 0
|
|
intersection = (min(ax2, bx2) - max(ax1, bx1)) * (min(ay2, by2) - max(ay1, by1))
|
|
area_a = (ax2 - ax1) * (ay2 - ay1)
|
|
area_b = (bx2 - bx1) * (by2 - by1)
|
|
union = area_a + area_b - intersection
|
|
return intersection / union
|
|
|
|
|
|
def find_max_overlap(box, box_list):
|
|
best_candidate = max(box_list, key=lambda x: compute_iou_from_boxes(box, x))
|
|
iou = compute_iou_from_boxes(box, best_candidate)
|
|
return best_candidate, iou
|
|
|
|
|
|
def compute_page_iou(results_box_list, gt_box_list):
|
|
results = results_box_list.copy()
|
|
gt = gt_box_list.copy()
|
|
if (not results) or (not gt):
|
|
return 0
|
|
iou_sum = 0
|
|
denominator = max(len(results), len(gt))
|
|
while gt and results:
|
|
gt_box = gt.pop()
|
|
best_match, best_iou = find_max_overlap(gt_box, results)
|
|
results.remove(best_match)
|
|
iou_sum += best_iou
|
|
score = iou_sum / denominator
|
|
return score
|
|
|
|
|
|
def compute_document_score(results_dict, annotation_dict):
|
|
|
|
page_weights = np.array([len(page["cells"]) for page in annotation_dict["pages"]])
|
|
page_weights = page_weights / sum(page_weights)
|
|
|
|
scores = []
|
|
for i in range(len(annotation_dict["pages"])):
|
|
scores.append(compute_page_iou(results_dict["pages"][i]["cells"], annotation_dict["pages"][i]["cells"]))
|
|
scores = np.array(scores)
|
|
|
|
doc_score = np.average(scores, weights=page_weights)
|
|
|
|
return doc_score
|