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