Merge pull request #3 from stefanklut/actions/black

Format Python code with psf/black push
stefanklut · Nov 13, 2023 · 622ff43 · 622ff43
2 parents 52c239f + 24978af
commit 622ff43
Showing 1 changed file with 64 additions and 55 deletions.
diff --git a/utils/vector_utils.py b/utils/vector_utils.py
@@ -1,6 +1,7 @@
 import numpy as np
 import cv2
 
+
 def point_line_segment_distance(line_segments: np.ndarray, points: np.ndarray) -> np.ndarray:
     """
     Calculate the distance between points and line segments.
@@ -22,24 +23,25 @@ def point_line_segment_distance(line_segments: np.ndarray, points: np.ndarray) -
     end_to_points_vector = points[:, np.newaxis, :] - line_ends
     end_to_points_norm = np.linalg.norm(end_to_points_vector, axis=2)
 
-    cross_product =  start_to_points_vector[:, :, 0] * line_vector[:, 1] - start_to_points_vector[:, :, 1] * line_vector[:, 0]
+    cross_product = start_to_points_vector[:, :, 0] * line_vector[:, 1] - start_to_points_vector[:, :, 1] * line_vector[:, 0]
     line_to_points_norm = np.abs(cross_product) / line_norm
-    
+
     dot_product = start_to_points_vector[:, :, 0] * line_vector[:, 0] + start_to_points_vector[:, :, 1] * line_vector[:, 1]
     dot_product_projection = dot_product / (line_norm * line_norm)
-    
+
     mask_less_0 = dot_product_projection <= 0
     mask_more_1 = dot_product_projection >= 1
     mask_between_0_1 = np.logical_not(np.logical_or(mask_less_0, mask_more_1))
-    
+
     distances = np.zeros((len(points), len(line_segments) - 1))
-    
+
     distances[mask_less_0] = start_to_points_norm[mask_less_0]
     distances[mask_more_1] = end_to_points_norm[mask_more_1]
     distances[mask_between_0_1] = line_to_points_norm[mask_between_0_1]
-    
+
     return distances
 
+
 def consecutive_booleans(array: np.ndarray) -> np.ndarray:
     """
     Find consecutive True values in a boolean array.
@@ -53,6 +55,7 @@ def consecutive_booleans(array: np.ndarray) -> np.ndarray:
     overlap = np.logical_and(array[1:], array[:-1])
     return overlap
 
+
 def point_line_segment_assignment(line_segments: np.ndarray, points: np.ndarray) -> np.ndarray:
     """
     Assign points to line segments based on minimum distance.
@@ -74,53 +77,56 @@ def point_line_segment_assignment(line_segments: np.ndarray, points: np.ndarray)
     end_to_points_vector = points[:, np.newaxis, :] - line_ends
     end_to_points_norm = np.linalg.norm(end_to_points_vector, axis=2)
 
-    cross_product =  start_to_points_vector[:, :, 0] * line_vector[:, 1] - start_to_points_vector[:, :, 1] * line_vector[:, 0]
+    cross_product = start_to_points_vector[:, :, 0] * line_vector[:, 1] - start_to_points_vector[:, :, 1] * line_vector[:, 0]
     line_to_points_norm = np.abs(cross_product) / line_norm
-    
+
     dot_product = start_to_points_vector[:, :, 0] * line_vector[:, 0] + start_to_points_vector[:, :, 1] * line_vector[:, 1]
     dot_product_projection = dot_product / (line_norm * line_norm)
-    
+
     mask_less_0 = dot_product_projection <= 0
     mask_more_1 = dot_product_projection >= 1
     mask_between_0_1 = np.logical_not(np.logical_or(mask_less_0, mask_more_1))
-    
+
     distances = np.zeros((len(points), len(line_segments) - 1))
-    
+
     distances[mask_less_0] = start_to_points_norm[mask_less_0]
     distances[mask_more_1] = end_to_points_norm[mask_more_1]
     distances[mask_between_0_1] = line_to_points_norm[mask_between_0_1]
-    
+
     min_values = np.min(distances, axis=1)
     is_min_value = distances == min_values[:, np.newaxis]
-    
+
     min_count = np.sum(is_min_value, axis=1)
     min_occurs_more_than_once = min_count > 1
-    
+
     label = np.zeros(min_occurs_more_than_once.shape)
-    
+
     for i, check in enumerate(min_occurs_more_than_once):
         if check:
             overlap_bool = consecutive_booleans(is_min_value[i])
             overlap_location = np.argwhere(overlap_bool)
             for j in overlap_location:
-                middle_vector = line_vector[j] / line_norm[j] - line_vector[j+1] / line_norm[j+1]
-
-                cross_product_1 =  end_to_points_vector[i, j, 0] * middle_vector[:, 1] - end_to_points_vector[i, j, 1] * middle_vector[:, 0]
+                middle_vector = line_vector[j] / line_norm[j] - line_vector[j + 1] / line_norm[j + 1]
+
+                cross_product_1 = (
+                    end_to_points_vector[i, j, 0] * middle_vector[:, 1] - end_to_points_vector[i, j, 1] * middle_vector[:, 0]
+                )
                 cross_product_2 = -line_vector[j, 0] * middle_vector[:, 1] + line_vector[j, 1] * middle_vector[:, 0]
-                
+
                 if cross_product_1 == 0 or cross_product_2 == 0:
-                    label[i] = np.random.choice([j, j+1])
+                    label[i] = np.random.choice([j, j + 1])
                 elif np.sign(cross_product_1) == np.sign(cross_product_2):
                     label[i] = j
                 else:
-                    label[i] = j+1
+                    label[i] = j + 1
 
             # label[i] = 5
         else:
             label[i] = np.argwhere(is_min_value[i])
-            
+
     return label
 
+
 def point_line_segment_side(line_segments: np.ndarray, points: np.ndarray) -> tuple:
     """
     Determine which side of the line segments points are located.
@@ -138,14 +144,14 @@ def point_line_segment_side(line_segments: np.ndarray, points: np.ndarray) -> tu
     line_vector = line_ends - line_starts
     start_to_points_vector = points[:, np.newaxis, :] - line_starts
 
+    cross_product = start_to_points_vector[:, :, 0] * line_vector[:, 1] - start_to_points_vector[:, :, 1] * line_vector[:, 0]
 
-    cross_product =  start_to_points_vector[:, :, 0] * line_vector[:, 1] - start_to_points_vector[:, :, 1] * line_vector[:, 0]
-
     left_of_line_vector = cross_product > 0
     right_of_line_vector = cross_product < 0
     on_line_vector = cross_product == 0
     return (left_of_line_vector, right_of_line_vector, on_line_vector)
 
+
 def point_top_bottom_assignment(line_segments: np.ndarray, points: np.ndarray) -> np.ndarray:
     """
     Assign points to top or bottom sides of line segments based on minimum distance.
@@ -162,57 +168,56 @@ def point_top_bottom_assignment(line_segments: np.ndarray, points: np.ndarray) -
 
     line_vector = line_ends - line_starts
     line_norm = np.linalg.norm(line_vector, axis=1)
-    
+
     non_zero = line_norm != 0
-    
-    #HACK If baseline has lenght 0, assume top bottom from the pages perspective
+
+    # HACK If baseline has lenght 0, assume top bottom from the pages perspective
     if not np.any(non_zero):
         line_starts = line_segments[0:1]
-        line_ends = line_starts[0:1] + np.asarray([1,0])
+        line_ends = line_starts[0:1] + np.asarray([1, 0])
         line_vector = line_ends - line_starts
         line_norm = np.linalg.norm(line_vector, axis=1)
     else:
         line_vector = line_vector[non_zero]
         line_ends = line_ends[non_zero]
         line_starts = line_starts[non_zero]
         line_norm = line_norm[non_zero]
-    
+
     start_to_points_vector = points[:, np.newaxis, :] - line_starts
     start_to_points_norm = np.linalg.norm(start_to_points_vector, axis=2)
     end_to_points_vector = points[:, np.newaxis, :] - line_ends
     end_to_points_norm = np.linalg.norm(end_to_points_vector, axis=2)
 
-    cross_product =  start_to_points_vector[:, :, 0] * line_vector[:, 1] - start_to_points_vector[:, :, 1] * line_vector[:, 0]
+    cross_product = start_to_points_vector[:, :, 0] * line_vector[:, 1] - start_to_points_vector[:, :, 1] * line_vector[:, 0]
     line_to_points_norm = np.abs(cross_product) / line_norm
-    
+
     dot_product = start_to_points_vector[:, :, 0] * line_vector[:, 0] + start_to_points_vector[:, :, 1] * line_vector[:, 1]
     dot_product_projection = dot_product / (line_norm * line_norm)
-    
+
     mask_less_0 = dot_product_projection <= 0
     mask_more_1 = dot_product_projection >= 1
     mask_between_0_1 = np.logical_not(np.logical_or(mask_less_0, mask_more_1))
-    
+
     distances = np.zeros((len(points), len(line_vector)))
-    
+
     distances[mask_less_0] = start_to_points_norm[mask_less_0]
     distances[mask_more_1] = end_to_points_norm[mask_more_1]
     distances[mask_between_0_1] = line_to_points_norm[mask_between_0_1]
-    
+
     min_values = np.min(distances, axis=1)
     is_min_value = distances == min_values[:, np.newaxis]
-    
+
     min_count = np.sum(is_min_value, axis=1)
     min_occurs_more_than_once = min_count > 1
-    
+
     label = np.zeros(min_occurs_more_than_once.shape)
-    
+
     for i, check in enumerate(min_occurs_more_than_once):
         if check:
             overlap_bool = consecutive_booleans(is_min_value[i])
             if not np.any(overlap_bool):
                 min_lines = np.nonzero(is_min_value[i])[0]
-
-
+
                 # import matplotlib.pyplot as plt
                 # print(is_min_value[i], distances[i])
                 # line = np.vstack((line_starts, line_ends[-1]))
@@ -228,20 +233,22 @@ def point_top_bottom_assignment(line_segments: np.ndarray, points: np.ndarray) -
             else:
                 cross_product_i = 0
                 overlap_location = np.argwhere(overlap_bool)[0]
-                
+
                 for j in overlap_location:
-                    middle_vector = line_vector[j] / line_norm[j] - line_vector[j+1] / line_norm[j+1]
-
-                    cross_product_1 =  end_to_points_vector[i, j, 0] * middle_vector[1] - end_to_points_vector[i, j, 1] * middle_vector[0]
+                    middle_vector = line_vector[j] / line_norm[j] - line_vector[j + 1] / line_norm[j + 1]
+
+                    cross_product_1 = (
+                        end_to_points_vector[i, j, 0] * middle_vector[1] - end_to_points_vector[i, j, 1] * middle_vector[0]
+                    )
                     cross_product_2 = -line_vector[j, 0] * middle_vector[1] + line_vector[j, 1] * middle_vector[0]
-                    
+
                     if cross_product_1 == 0 or cross_product_2 == 0:
-                        cross_product_i = cross_product[i, np.random.choice([j, j+1])]
+                        cross_product_i = cross_product[i, np.random.choice([j, j + 1])]
                     elif np.sign(cross_product_1) == np.sign(cross_product_2):
                         cross_product_i = cross_product[i, j]
                     else:
-                        cross_product_i = cross_product[i, j+1]
-                        
+                        cross_product_i = cross_product[i, j + 1]
+
             if cross_product_i == 0:
                 label[i] = np.random.choice([0, 1])
             elif cross_product_i > 0:
@@ -256,9 +263,10 @@ def point_top_bottom_assignment(line_segments: np.ndarray, points: np.ndarray) -
                 label[i] = 1
             else:
                 label[i] = 0
-            
+
     return label
 
+
 def draw_lines(lines: np.ndarray, thickness: int) -> np.ndarray:
     """
     Draw lines with specified thickness and assign point labels.
@@ -279,12 +287,13 @@ def draw_lines(lines: np.ndarray, thickness: int) -> np.ndarray:
         line_pixel_coords = np.column_stack(np.where(empty_mask == 1))[:, ::-1]
         print(line_pixel_coords.shape)
         empty_mask.fill(0)
-        mask[line_pixel_coords[:, 1], line_pixel_coords[:, 0]] = (point_top_bottom_assignment(line, line_pixel_coords)+1) * 100
+        mask[line_pixel_coords[:, 1], line_pixel_coords[:, 0]] = (
+            point_top_bottom_assignment(line, line_pixel_coords) + 1
+        ) * 100
 
     return mask
 
 
-
 if __name__ == "__main__":
     # coordinates = np.array([(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)])
     # coordinates = np.array([(10, 10)])
@@ -297,7 +306,6 @@ def draw_lines(lines: np.ndarray, thickness: int) -> np.ndarray:
     # import sys
     # sys.exit()
     # result = point_line_segment_distance(line_segments, coordinates)
-
 
     # labels = np.argmin(result, axis=1)
     # print(labels.shape)
@@ -318,6 +326,7 @@ def draw_lines(lines: np.ndarray, thickness: int) -> np.ndarray:
     # plt.show()
 
     import matplotlib.pyplot as plt
-    im = draw_lines(np.array([[(10,10), (10,10)]]), 10)
-    plt.imshow(im, cmap='gray')
-    plt.show()
+
+    im = draw_lines(np.array([[(10, 10), (10, 10)]]), 10)
+    plt.imshow(im, cmap="gray")
+    plt.show()