Create utils.py

utils.py

@@ -0,0 +1,179 @@
+from reading_order import GraphBasedOrdering
+from shapely.geometry import Polygon
+import glob
+import os
+
+from shapely.validation import make_valid
+
+order_poly = GraphBasedOrdering()
+
+def load_image_paths(input_folder, extensions=None):
+    """
+    Load image files from a folder.
+    
+    Args:
+        input_folder: Path to the folder containing images
+        extensions: List of file extensions to include (default: common image formats)
+    
+    Returns:
+        List of image file paths
+    """
+    if extensions is None:
+        extensions = ['jpg', 'jpeg', 'png', 'gif', 'bmp', 'tiff', 'tif', 'webp']
+    
+    images = []
+    for ext in extensions:
+        # Add both lowercase and uppercase versions
+        images.extend(glob.glob(os.path.join(input_folder, f'*.{ext}')))
+        images.extend(glob.glob(os.path.join(input_folder, f'*.{ext.upper()}')))
+    
+    return sorted(images)
+
+def get_default_region(image_shape):
+    """Function for creating a default region if no regions are detected."""
+    w, h = image_shape
+    region = {'coords': [[0.0, 0.0], [w, 0.0], [w, h], [0.0, h]], 
+            'max_min': [w, 0.0, h, 0.0], 
+            'name': "paragraph", 
+            'id': '0', 
+            'img_shape': (h, w)}
+    return [region]
+
+def get_dist(line_polygon, regions):
+    """Function for finding the closest region to the text line."""
+    dist, reg_id = 1000000, None
+    line_polygon = validate_polygon(line_polygon)
+    if line_polygon:
+        for region in regions:
+            # Calculates dictance between line and regions polygons
+            region_polygon = validate_polygon(region['coords'])
+            if region_polygon:
+                line_reg_dist = line_polygon.distance(region_polygon)
+                if line_reg_dist < dist:
+                    dist = line_reg_dist
+                    reg_id = region['id']
+    return reg_id
+
+def validate_polygon(polygon):
+    """"Function for testing and correcting the validity of polygons."""
+    if len(polygon) > 2:
+        polygon = Polygon(polygon)
+        if not polygon.is_valid:
+            polygon = make_valid(polygon)
+        return polygon
+    else:
+        return None
+
+def get_iou(poly1, poly2):
+    """Function for calculating Intersection over Union (IoU) values."""
+    # If the polygons don't intersect, IoU is 0
+    iou = 0
+    poly1 = validate_polygon(poly1)
+    poly2 = validate_polygon(poly2)
+    if poly1 and poly2:
+        if poly1.intersects(poly2):
+            # Calculates intersection of the 2 polygons
+            intersect = poly1.intersection(poly2).area
+            # Calculates union of the 2 polygons
+            uni = poly1.union(poly2)
+            # Calculates intersection over union
+            iou = intersect / uni.area
+    return iou
+
+def get_line_regions(lines, regions):
+    """Function for connecting each text line to one region.
+    
+    Docstring generated with Claude
+    Connect each text line to a region based on intersection or distance.
+
+    Args:
+        lines (dict): Dictionary containing text line information with keys:
+            - 'coords' (list): List of line polygons (coordinates). E.g. [[x1,y1], ..., [xn,yn]]
+            - 'max_min' (list): List of bounding box coordinates for each line
+        regions (list): List of region dictionaries, each containing:
+            - 'coords': Region polygon coordinates. E.g. [[x1,y1], ..., [xn,yn]]
+            - 'id': Region identifier
+
+    Returns:
+        list: List of dictionaries, each representing a line with keys:
+            - 'polygon': Line polygon coordinates
+            - 'reg_id': ID of the region the line belongs to
+            - 'max_min': Bounding box coordinates [x_min, y_min, x_max, y_max]
+
+    """
+    lines_list = []
+    for i in range(len(lines['coords'])):
+        iou, reg_id = 0, ''
+        max_min = [0.0, 0.0, 0.0, 0.0]
+        polygon = lines['coords'][i]
+        for region in regions:
+            line_reg_iou = get_iou(polygon, region['coords']) 
+            if line_reg_iou > iou:
+                iou = line_reg_iou
+                reg_id = region['id']
+        # If line polygon does not intersect with any region, a distance metric is used for defining 
+        # the region that the line belongs to
+        if iou == 0:
+            reg_id = get_dist(polygon, regions)
+
+        if (len(lines['max_min']) - 1) >= i:
+            max_min = lines['max_min'][i]
+
+        new_line = {'polygon': polygon, 'reg_id': reg_id, 'max_min': max_min}
+        lines_list.append(new_line)
+    return lines_list
+
+def order_regions_lines(lines, regions):
+    """Function for ordering line predictions inside each region.
+    
+    Docstring generated with Claude
+    Order text lines within each region and order the regions themselves.
+
+    Args:
+        lines (list): List of line dictionaries, each containing:
+            - 'reg_id': ID of the region the line belongs to
+            - 'max_min': Bounding box coordinates for the line
+            - 'polygon': Line polygon coordinates
+        regions (list): List of region dictionaries, each containing:
+            - 'id': Region identifier
+            - 'coords': Region polygon coordinates
+            - 'name': Region name
+            - 'max_min': Bounding box coordinates for the region
+            - 'img_shape': Shape of the source image
+
+    Returns:
+        list: List of ordered region dictionaries, each containing:
+            - 'region_coords': Region polygon coordinates
+            - 'region_name': Region name
+            - 'lines': Ordered list of line polygons within the region
+            - 'img_shape': Shape of the source image
+        
+        Note: Only regions containing at least one line are included in the output.
+        Both lines within regions and regions themselves are ordered by their spatial positions.
+    """
+    regions_with_rows = []
+    region_max_mins = []
+    for i, region in enumerate(regions):
+        line_max_mins = []
+        line_polygons = []
+        for line in lines:
+            if line['reg_id'] == region['id']:
+                line_max_mins.append(line['max_min'])
+                line_polygons.append(line['polygon'])
+        if line_polygons:
+            # If one or more lines are connected to a region, line order inside the region is defined
+            # and the predicted text lines are joined in the same python dict
+            line_order = order_poly.order(line_max_mins)
+            line_polygons = [line_polygons[i] for i in line_order]
+            new_region = {'region_coords': region['coords'], 
+                        'region_name': region['name'], 
+                        'lines': line_polygons, 
+                        'img_shape': region['img_shape']}
+            region_max_mins.append(region['max_min'])
+            regions_with_rows.append(new_region)
+        else:
+            continue
+    # Creates an ordering of the detected regions based on their polygon coordinates
+    region_order = order_poly.order(region_max_mins)
+    regions_with_rows = [regions_with_rows[i] for i in region_order]
+    return regions_with_rows