Move rectangle creation

page_xml/output_pageXML.py

@@ -56,7 +56,7 @@ def __init__(
         region_type: Optional[list[str]] = None,
         cfg: Optional[CfgNode] = None,
         whitelist: Optional[Iterable[str]] = None,
-        rectangle_regions: Optional[list[str]] = [],
+        rectangle_regions: Optional[Iterable[str]] = None,
         min_region_size: int = 10,
     ) -> None:
         """
@@ -81,8 +81,6 @@ def __init__(
         """
         super().__init__(mode, line_width, regions, merge_regions, region_type)
 
-        self.min_region_size = min_region_size
-        self.rectangle_regions = rectangle_regions
         self.logger = logging.getLogger(get_logger_name())
 
         self.output_dir = None
@@ -96,6 +94,8 @@ def __init__(
         self.cfg = cfg
 
         self.whitelist = set() if whitelist is None else set(whitelist)
+        self.min_region_size = min_region_size
+        self.rectangle_regions = set() if rectangle_regions is None else set(rectangle_regions)
 
     def set_output_dir(self, output_dir: str | Path):
         if isinstance(output_dir, str):
@@ -165,7 +165,6 @@ def generate_single_page(
         height, width = sem_seg.shape[-2:]
 
         scaling = np.asarray([old_width, old_height] / np.asarray([width, height]))
-        # scaling = np.asarray((1,1))
 
         page = PageData(xml_output_path)
         page.new_page(image_path.name, str(old_height), str(old_width))
@@ -189,30 +188,30 @@ def generate_single_page(
                 contours, hierarchy = cv2.findContours(binary_region_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
 
                 for cnt in contours:
-                    # --- remove small objects
+                    # remove small objects
                     if cnt.shape[0] < 4:
                         continue
                     if cv2.contourArea(cnt) < self.min_region_size:
                         continue
 
                     region_id += 1
 
-                    # --- soft a bit the region to prevent spikes
-                    epsilon = 0.0005 * cv2.arcLength(cnt, True)
-                    approx_poly = cv2.approxPolyDP(cnt, epsilon, True)
-
-                    approx_poly = np.round((approx_poly * scaling)).astype(np.int32)
-
                     region_coords = ""
                     if region in self.rectangle_regions:
                         # find bounding box
-                        rect = cv2.minAreaRect(approx_poly)
-                        rect = cv2.boxPoints(rect)
-                        for coords in rect:
-                            region_coords = region_coords + f" {round(coords[0])},{round(coords[1])}"
+                        rect = cv2.minAreaRect(cnt)
+                        poly = cv2.boxPoints(rect) * scaling
                     else:
-                        for coords in approx_poly.reshape(-1, 2):
-                            region_coords = region_coords + f" {coords[0]},{coords[1]}"
+                        # soft a bit the region to prevent spikes
+                        epsilon = 0.0005 * cv2.arcLength(cnt, True)
+                        approx_poly = cv2.approxPolyDP(cnt, epsilon, True)
+
+                        approx_poly = np.round((approx_poly * scaling)).astype(np.int32)
+
+                        poly = approx_poly.reshape(-1, 2)
+
+                    for coords in poly:
+                        region_coords = region_coords + f" {round(coords[0])},{round(coords[1])}"
 
                     region_coords = region_coords.strip()