def test_subtract_rectangle_with_shared_boundaries():
rect1 = Rectangle(0, 0, 20, 20)
rect2 = Rectangle(10, 0, 10, 10)
diff_rects = list(subtract(rect1, rect2))
assert len(diff_rects) == 2
assert Rectangle(0, 10, 20, 10) in diff_rects
assert Rectangle(0, 0, 10, 10) in diff_rects
def test_union_rectangles():
rects = [Rectangle(0, 10, 20, 20), Rectangle(10, 0, 20, 20)]
union_rects = list(union(rects))
assert len(union_rects) == 3
assert Rectangle(0, 10, 20, 20) in union_rects
assert Rectangle(20, 10, 10, 10) in union_rects
assert Rectangle(10, 0, 20, 10) in union_rects
def test_subtract_rectangle_iterable_from_rectangle():
rect = Rectangle(10, -5, 20, 20)
other_rects = [Rectangle(0, 5, 20, 20), Rectangle(15, 0, 20, 20)]
diff_rects = list(subtract_multiple(rect, other_rects))
assert len(diff_rects) == 2
assert Rectangle(10, 0, 5, 5) in diff_rects
assert Rectangle(10, -5, 20, 5) in diff_rects
def test_compute_page_iou_for_rectangle_iterables():
# There's a 10px-wide overlap between rect1 and rect2; the algorithm for IOU should use the
# union of the areas of the input rects.
rects = [Rectangle(0, 0, 20, 20), Rectangle(10, 0, 30, 20)]
other_rects = [Rectangle(10, 0, 20, 20), Rectangle(35, 0, 20, 20)]
# Intersection area = 10 x 20 + 10 x 20 + 5 x 20 = 500
# Union area = 55 x 20 => 1100
assert iou(rects, other_rects) == float(500) / 1100
def test_rectangle_precision_recall():
expected = [fs(Rectangle(0, 0, 20, 20)), fs(Rectangle(10, 0, 30, 20))]
actual = [Rectangle(10, 0, 20, 20), Rectangle(35, 0, 20, 20)]
# The threshold value is set to the level where rectangle 1 in 'expected' does not have a
# a match in 'actual', and rectangle 2 in 'expected' only has a match if you consider
# its overlap with *all* rectangles in 'actual'.
precision, recall = compute_accuracy(expected, actual, minimum_iou=0.5)
assert precision == 0.5
assert recall == 0.5
def test_compute_iou_per_rectangle_set():
# There's a 10px-wide overlap between rect1 and rect2; the algorithm for IOU should use the
# union of the areas of the input rects.
rects = [
fs(Rectangle(0, 0, 20, 20)),
fs(Rectangle(10, 0, 30, 20)),
fs(Rectangle(40, 10, 10, 10), Rectangle(40, 0, 10, 10)),
]
other_rects = [Rectangle(10, 0, 20, 20), Rectangle(35, 0, 20, 20)]
ious = iou_per_rectangle(rects, other_rects)
assert ious[fs(Rectangle(0, 0, 20, 20))] == float(10) / 30
assert ious[fs(Rectangle(10, 0, 30, 20))] == float(25) / 45
assert (ious[fs(Rectangle(40, 10, 10, 10),
Rectangle(40, 0, 10, 10))] == float(10) / 20)
def find_boxes_with_color(
image: np.ndarray,
hue: float,
tolerance: float = 0.01,
masks: Optional[Iterable[Rectangle]] = None,
) -> List[Rectangle]:
"""
Arguments:
- 'hue': is a floating point number between 0 and 1
- 'tolerance': is the amount of difference from 'hue' (from 0-to-1) still considered that hue.
- 'masks': a set of masks to apply to the image, one at a time. Bounding boxes are extracted
from within each of those boxes. Masks should be in pixel coordinates.
"""
height, width, _ = image.shape
if masks is None:
masks = (Rectangle(left=0, top=0, width=width, height=height),)
CV2_MAXIMUM_HUE = 180
SATURATION_THRESHOLD = 50 # out of 255
cv2_hue = hue * CV2_MAXIMUM_HUE
cv2_tolerance = tolerance * CV2_MAXIMUM_HUE
img_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
saturated_pixels = img_hsv[:, :, 1] > SATURATION_THRESHOLD
hues = img_hsv[:, :, 0]
distance_to_hue = np.abs(hues.astype(np.int16) - cv2_hue)
abs_distance_to_hue = np.minimum(distance_to_hue, CV2_MAXIMUM_HUE - distance_to_hue)
boxes = []
for mask in masks:
masked_distances = np.full(abs_distance_to_hue.shape, np.inf)
right = mask.left + mask.width
bottom = mask.top + mask.height
masked_distances[mask.top : bottom, mask.left : right] = abs_distance_to_hue[
mask.top : bottom, mask.left : right
]
# To determine which pixels have a color, we look for those that:
# 1. Match the hue
# 2. Are heavily saturated (i.e. aren't white---white pixels could be detected as having
# any hue, with no saturation.)
matching_pixels = np.where(
(masked_distances <= cv2_tolerance) & saturated_pixels
)
matching_pixels_list: List[Point] = []
for i in range(len(matching_pixels[0])):
matching_pixels_list.append(
Point(matching_pixels[1][i], matching_pixels[0][i])
)
boxes.extend(list(PixelMerger().merge_pixels(matching_pixels_list)))
return boxes
def test_intersect_rectangle_iterables():
rects = [Rectangle(0, 0, 20, 20), Rectangle(20, 0, 20, 20)]
other_rects = [Rectangle(10, 0, 20, 20), Rectangle(35, 0, 20, 20)]
intersection_rects = list(intersect(rects, other_rects))
assert Rectangle(10, 0, 10, 20) in intersection_rects
assert Rectangle(20, 0, 10, 20) in intersection_rects
assert Rectangle(35, 0, 5, 20) in intersection_rects
def _create_rectangle(self) -> Rectangle:
assert self.min_x is not None
assert self.max_x is not None
assert self.top_y is not None
assert self.bottom_y is not None
return Rectangle(
left=self.min_x,
top=self.top_y,
width=self.max_x - self.min_x + 1,
height=self.bottom_y - self.top_y + 1,
)
def test_another_union():
rects = [
Rectangle(0, 0, 20, 20),
Rectangle(20, 0, 20, 20),
Rectangle(10, 0, 20, 20),
Rectangle(35, 0, 20, 20),
]
union_rects = list(union(rects))
assert len(union_rects) == 3
assert Rectangle(0, 0, 20, 20) in union_rects
assert Rectangle(20, 0, 20, 20) in union_rects
assert Rectangle(40, 0, 15, 20) in union_rects
def test_subtract_rectangle_inside_another():
outer = Rectangle(0, 0, 20, 20)
inner = Rectangle(5, 5, 10, 10)
diff_rects = list(subtract(outer, inner))
assert len(diff_rects) == 4
assert Rectangle(0, 0, 20, 5) in diff_rects
assert Rectangle(0, 5, 5, 10) in diff_rects
assert Rectangle(15, 5, 5, 10) in diff_rects
assert Rectangle(0, 0, 20, 5) in diff_rects
def test_compute_iou_per_rectangle_set():
regions = [
fs(Rectangle(0, 0, 20, 20)),
fs(Rectangle(10, 0, 30, 20)),
fs(Rectangle(40, 0, 10, 10)),
]
other_regions = [
fs(Rectangle(0, 0, 10, 10), Rectangle(10, 0, 10, 10)), # overlaps 1
fs(Rectangle(10, 0, 20, 20)), # overlaps both 1 and 2
]
ious = iou_per_region(regions, other_regions, minimum_iou=0.25)
assert len(ious) == 2
assert ious[(regions[0], other_regions[0])] == float(10) / 20
assert ious[(regions[1], other_regions[1])] == float(20) / 30
def extract_bounding_boxes(
diff_image: np.ndarray,
page_number: int,
hue: float,
masks: Optional[Iterable[FloatRectangle]] = None,
) -> List[BoundingBox]:
"""
See 'PixelMerger' for description of how bounding boxes are extracted.
Masks are assumed to be non-intersecting. Masks should be expressed as ratios relative to the
page's width and height instead of pixel values---left, top, width, and height all have values
in the range 0..1).
"""
image_height, image_width, _ = diff_image.shape
pixel_masks = None
if masks is not None:
pixel_masks = [
Rectangle(
left=round(m.left * image_width),
top=round(m.top * image_height),
width=round(m.width * image_width),
height=round(m.height * image_height),
) for m in masks
]
pixel_boxes = list(
find_boxes_with_color(diff_image, hue, masks=pixel_masks))
boxes = []
for box in pixel_boxes:
left_ratio = float(box.left) / image_width
top_ratio = float(box.top) / image_height
width_ratio = float(box.width) / image_width
height_ratio = float(box.height) / image_height
boxes.append(
BoundingBox(left_ratio, top_ratio, width_ratio, height_ratio,
page_number))
return boxes
def test_subtract_rectangle_iterable_from_rectangle_iterable():
rects = [Rectangle(0, 0, 20, 20), Rectangle(20, 0, 20, 20)]
other_rects = [Rectangle(10, 0, 20, 20), Rectangle(35, 0, 20, 20)]
diff_rects = list(subtract_multiple_from_multiple(rects, other_rects))
assert Rectangle(0, 0, 10, 20) in diff_rects
assert Rectangle(30, 0, 5, 20) in diff_rects
def test_subtract_outer_rectangle_from_inner_rectangle():
outer = Rectangle(0, 0, 20, 20)
inner = Rectangle(5, 5, 10, 10)
assert len(list(subtract(inner, outer))) == 0
def test_subtract_nonintersecting_rectangle():
rect1 = Rectangle(0, 0, 20, 20)
rect2 = Rectangle(30, 0, 20, 20)
diff_rects = list(subtract(rect1, rect2))
assert len(diff_rects) == 1
assert diff_rects == [rect1]
def test_another_subtract():
rect1 = Rectangle(20, 0, 10, 10)
rect2 = Rectangle(15, -5, 20, 20)
diff_rects = list(subtract(rect1, rect2))
assert len(diff_rects) == 0
def test_subtract_rectangle_from_itself():
rect1 = Rectangle(0, 0, 20, 20)
rect2 = Rectangle(0, 0, 20, 20)
assert len(list(subtract(rect1, rect2))) == 0
def load_hues(self, arxiv_id: ArxivId,
iteration: str) -> List[HueSearchRegion]:
equation_boxes_path = os.path.join(
directories.arxiv_subdir("hue-locations-for-equations", arxiv_id),
"hue_locations.csv",
)
bounding_boxes: Dict[EquationId, BoundingBoxesByFile] = {}
for location_info in file_utils.load_from_csv(equation_boxes_path,
HueLocationInfo):
equation_id = EquationId(
tex_path=location_info.tex_path,
equation_index=int(location_info.entity_id),
)
if equation_id not in bounding_boxes:
bounding_boxes[equation_id] = {}
file_path = location_info.relative_file_path
if file_path not in bounding_boxes[equation_id]:
bounding_boxes[equation_id][file_path] = []
box = BoundingBox(
page=location_info.page,
left=location_info.left,
top=location_info.top,
width=location_info.width,
height=location_info.height,
)
bounding_boxes[equation_id][file_path].append(box)
token_records_by_equation: Dict[EquationId, Dict[
int, EquationTokenColorizationRecord]] = {}
token_hues_path = os.path.join(
directories.iteration(
"sources-with-colorized-equation-tokens",
arxiv_id,
iteration,
),
"entity_hues.csv",
)
for record in file_utils.load_from_csv(
token_hues_path, EquationTokenColorizationRecord):
equation_id = EquationId(tex_path=record.tex_path,
equation_index=record.equation_index)
token_index = int(record.token_index)
if equation_id not in token_records_by_equation:
token_records_by_equation[equation_id] = {}
token_records_by_equation[equation_id][token_index] = record
hue_searches = []
for equation_id, boxes_by_file in bounding_boxes.items():
for file_path, boxes in boxes_by_file.items():
masks_by_page: MasksForPages = {}
for box in boxes:
if box.page not in masks_by_page:
masks_by_page[box.page] = []
masks_by_page[box.page].append(
Rectangle(box.left, box.top, box.width, box.height))
if equation_id in token_records_by_equation:
for token_index, record in token_records_by_equation[
equation_id].items():
hue_searches.append(
HueSearchRegion(
hue=record.hue,
record=record,
relative_file_path=file_path,
masks=masks_by_page,
))
return hue_searches
