OCR / util.py
onipot
first
b025645
classes = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
class Detection(object):
def __init__(self, id: int, xmin: int, ymin: int, xmax:int, ymax:int, conf: float, class_id:int, class_name:str, orig_img_sz: "tuple[int]") -> None:
self.id = id
self.xmin = xmin
self.ymin = ymin
self.xmax = xmax
self.ymax = ymax
self.w = self.xmax - self.xmin
self.h = self.ymax - self.ymin
self.conf = conf
self.class_id = class_id
self.class_name = class_name
self.orig_img_h = orig_img_sz[1]
self.orig_img_w = orig_img_sz[0]
def get_hw_ratio(self):
return self.h / self.w
def get_height_proportion(self):
return self.h / self.orig_img_h
def get_width_proportion(self):
return self.w / self.orig_img_w
def contains(self, detection2: "Detection"):
if self.xmin <= detection2.xmin and self.xmax >= detection2.xmax and \
self.ymin <= detection2.ymin and self.ymax >= detection2.ymax:
return True
return False
def get_iou(self, detection2: "Detection"):
"""
Calculate the Intersection over Union (IoU) of two bounding boxes.
Returns
-------
float
in [0, 1]
"""
assert self.xmin < self.xmax
assert self.ymin < self.ymax
assert detection2.xmin < detection2.xmax
assert detection2.ymin < detection2.ymax
# determine the coordinates of the intersection rectangle
x_left = max(self.xmin, detection2.xmin)
y_top = max(self.ymin, detection2.ymin)
x_right = min(self.xmax, detection2.xmax)
y_bottom = min(self.ymax, detection2.ymax)
if x_right < x_left or y_bottom < y_top:
return 0.0
# The intersection of two axis-aligned bounding boxes is always an
# axis-aligned bounding box
intersection_area = (x_right - x_left) * (y_bottom - y_top)
# compute the area of both AABBs
bb1_area = (self.xmax - self.xmin) * (self.ymax - self.ymin)
bb2_area = (detection2.xmax - detection2.xmin) * (detection2.ymax - detection2.ymin)
# compute the intersection over union by taking the intersection
# area and dividing it by the sum of prediction + ground-truth
# areas - the interesection area
iou = intersection_area / float(bb1_area + bb2_area - intersection_area)
return iou
def __str__(self) -> str:
return f"[{self.xmin}, {self.ymin}, {self.xmax}, {self.ymax}]"