align box in multiple function

modules/toXML.py

@@ -16,131 +16,145 @@ def find_position(pool_index, BPMN_id):
 
     return position
 
-def align_boxes(pred, size):
-    modified_pred = copy.deepcopy(pred)  # Make a deep copy of the prediction
-
-    # Step 1: Calculate the center of each bounding box and group them by pool
+# Calculate the center of each bounding box and group them by pool
+def calculate_centers_and_group_by_pool(pred, class_dict):
     pool_groups = {}
     for pool_index, element_indices in pred['pool_dict'].items():
         pool_groups[pool_index] = []
         for i in element_indices:
-            if i >= len(modified_pred['labels']):
-                continue 
-            if class_dict[modified_pred['labels'][i]] != 'dataObject' or class_dict[modified_pred['labels'][i]] != 'dataStore':
-                x1, y1, x2, y2 = modified_pred['boxes'][i]
+            if i >= len(pred['labels']):
+                continue
+            if class_dict[pred['labels'][i]] not in ['dataObject', 'dataStore']:
+                x1, y1, x2, y2 = pred['boxes'][i]
                 center = [(x1 + x2) / 2, (y1 + y2) / 2]
                 pool_groups[pool_index].append((center, i))
-                
-    # Function to group centers within a specified range
-    def group_centers(centers, axis, range_=50):
-        groups = []
-        while centers:
-            center, idx = centers.pop(0)
-            group = [(center, idx)]
-            for other_center, other_idx in centers[:]:
-                if abs(center[axis] - other_center[axis]) <= range_:
-                    group.append((other_center, other_idx))
-                    centers.remove((other_center, other_idx))
-            groups.append(group)
-        return groups
-
-    # Step 2: Align the elements within each pool
+    return pool_groups
+
+# Group centers within a specified range
+def group_centers(centers, axis, range_=50):
+    groups = []
+    while centers:
+        center, idx = centers.pop(0)
+        group = [(center, idx)]
+        for other_center, other_idx in centers[:]:
+            if abs(center[axis] - other_center[axis]) <= range_:
+                group.append((other_center, other_idx))
+                centers.remove((other_center, other_idx))
+        groups.append(group)
+    return groups
+
+# Align the elements within each pool
+def align_elements_within_pool(modified_pred, pool_groups, class_dict, size):
     for pool_index, centers in pool_groups.items():
-        # Group bounding boxes by checking if their centers are within ±50 pixels on the y-axis
         y_groups = group_centers(centers.copy(), axis=1)
+        align_y_coordinates(modified_pred, y_groups, class_dict, size)
+        
+        centers = recalculate_centers(modified_pred, y_groups)
+        x_groups = group_centers(centers.copy(), axis=0)
+        align_x_coordinates(modified_pred, x_groups, class_dict, size)
+
+# Align the y-coordinates of the centers of grouped bounding boxes
+def align_y_coordinates(modified_pred, y_groups, class_dict, size):
+    for group in y_groups:
+        avg_y = sum([c[0][1] for c in group]) / len(group)
+        for (center, idx) in group:
+            label = class_dict[modified_pred['labels'][idx]]
+            if label in size:
+                new_center = (center[0], avg_y)
+                modified_pred['boxes'][idx] = [
+                    new_center[0] - size[label][0] / 2, 
+                    new_center[1] - size[label][1] / 2, 
+                    new_center[0] + size[label][0] / 2, 
+                    new_center[1] + size[label][1] / 2
+                ]
+
+# Recalculate centers after alignment
+def recalculate_centers(modified_pred, groups):
+    centers = []
+    for group in groups:
+        for center, idx in group:
+            x1, y1, x2, y2 = modified_pred['boxes'][idx]
+            center = [(x1 + x2) / 2, (y1 + y2) / 2]
+            centers.append((center, idx))
+    return centers
+
+# Align the x-coordinates of the centers of grouped bounding boxes
+def align_x_coordinates(modified_pred, x_groups, class_dict, size):
+    for group in x_groups:
+        avg_x = sum([c[0][0] for c in group]) / len(group)
+        for (center, idx) in group:
+            label = class_dict[modified_pred['labels'][idx]]
+            if label in size:
+                new_center = (avg_x, center[1])
+                modified_pred['boxes'][idx] = [
+                    new_center[0] - size[label][0] / 2, 
+                    modified_pred['boxes'][idx][1], 
+                    new_center[0] + size[label][0] / 2, 
+                    modified_pred['boxes'][idx][3]
+                ]
+
+# Expand the pool bounding boxes to fit the aligned elements
+def expand_pool_bounding_boxes(modified_pred, pred, size_elements):
+    for idx, (pool_index, keep_elements) in enumerate(modified_pred['pool_dict'].items()):
+        if len(keep_elements) != 0:
+            marge = 50
+        else:
+            marge = 0
 
-        # Align the y-coordinates of the centers of grouped bounding boxes
-        for group in y_groups:
-            avg_y = sum([c[0][1] for c in group]) / len(group)  # Calculate the average y-coordinate
-            for (center, idx) in group:
-                label = class_dict[modified_pred['labels'][idx]]
-                if label in size:
-                    new_center = (center[0], avg_y)  # Align the y-coordinate
-                    modified_pred['boxes'][idx] = [
-                        new_center[0] - size[label][0] / 2, 
-                        new_center[1] - size[label][1] / 2, 
-                        new_center[0] + size[label][0] / 2, 
-                        new_center[1] + size[label][1] / 2
-                    ]
-
-        # Recalculate centers after vertical alignment
-        centers = []
-        for group in y_groups:
-            for center, idx in group:
-                x1, y1, x2, y2 = modified_pred['boxes'][idx]
-                center = [(x1 + x2) / 2, (y1 + y2) / 2]
-                centers.append((center, idx))
+        position = find_position(pool_index, modified_pred['BPMN_id'])
 
-        # Group bounding boxes by checking if their centers are within ±50 pixels on the x-axis
-        x_groups = group_centers(centers.copy(), axis=0)
+        if keep_elements == [] or position is None:
+            min_x, min_y, max_x, max_y = modified_pred['boxes'][position]
+        else:
+            min_x, min_y, max_x, max_y = calculate_pool_bounds(modified_pred['boxes'], modified_pred['labels'], keep_elements, size_elements)
 
-        # Align the x-coordinates of the centers of grouped bounding boxes
-        for group in x_groups:
-            avg_x = sum([c[0][0] for c in group]) / len(group)  # Calculate the average x-coordinate
-            for (center, idx) in group:
-                label = class_dict[modified_pred['labels'][idx]]
-                if label in size:
-                    new_center = (avg_x, center[1])  # Align the x-coordinate
-                    modified_pred['boxes'][idx] = [
-                        new_center[0] - size[label][0] / 2, 
-                        modified_pred['boxes'][idx][1], 
-                        new_center[0] + size[label][0] / 2, 
-                        modified_pred['boxes'][idx][3]
-                    ]
-
-    # Step 3: Expand the pool bounding boxes to fit the aligned elements
-    if len(pred['pool_dict'])>1:
-        for idx, (pool_index, keep_elements) in enumerate(modified_pred['pool_dict'].items()):
-            size_elements = get_size_elements(st.session_state.size_scale)
-            if len(keep_elements) != 0:
-                marge = 50
-            else:
-                marge = 0     
-
-            position = find_position(pool_index, modified_pred['BPMN_id'])
-            
-            if keep_elements == [] or position is None:
-                min_x, min_y, max_x, max_y = modified_pred['boxes'][position]
-            else:
-                size_elements = get_size_elements(st.session_state.size_scale)
-                min_x, min_y, max_x, max_y = calculate_pool_bounds(modified_pred['boxes'], modified_pred['labels'], keep_elements, size_elements)
-                
-            pool_width = max_x - min_x  
-            pool_height = max_y - min_y  
-            if pool_width < 300 or pool_height < 30:
-                error("The pool is maybe too small, please add more elements or increase the scale by zooming on the image.")
-                continue
+        pool_width = max_x - min_x
+        pool_height = max_y - min_y
+        if pool_width < 300 or pool_height < 30:
+            error("The pool is maybe too small, please add more elements or increase the scale by zooming on the image.")
+            continue
+
+        modified_pred['boxes'][position] = [min_x - marge, min_y - marge // 2, min_x + pool_width + marge, min_y + pool_height + marge // 2]
+
+# Adjust left and right boundaries of all pools
+def adjust_pool_boundaries(modified_pred, pred):
+    min_left, max_right = 0, 0
+    for pool_index, element_indices in pred['pool_dict'].items():
+        position = find_position(pool_index, modified_pred['BPMN_id'])
+        if position >= len(modified_pred['boxes']):
+            continue
+        x1, y1, x2, y2 = modified_pred['boxes'][position]
+        left = x1
+        right = x2
+        if left < min_left:
+            min_left = left
+        if right > max_right:
+            max_right = right
 
-            modified_pred['boxes'][position] = [min_x -marge, min_y-marge//2, min_x+pool_width+marge, min_y+pool_height+marge//2]
-
-        min_left,max_right = 0, 0
-        for pool_index, element_indices in pred['pool_dict'].items():
-                position = find_position(pool_index, modified_pred['BPMN_id'])
-                if position >= len(modified_pred['boxes']):
-                    print(f"Problem with the index {pool_index} with a length of {len(modified_pred['boxes'])}")
-                    continue
-                x1, y1, x2, y2 = modified_pred['boxes'][position]
-                left = x1
-                right = x2
-                if left < min_left:
-                    min_left = left
-                if right > max_right:
-                    max_right = right
-
-        for pool_index, element_indices in pred['pool_dict'].items():
-                position = find_position(pool_index, modified_pred['BPMN_id'])
-                if position >= len(modified_pred['boxes']):
-                    #print(f"Problem with the index {pool_index} with a length of {len(modified_pred['boxes'])}")
-                    continue
-                x1, y1, x2, y2 = modified_pred['boxes'][position]
-                if x1 > min_left:
-                    x1 = min_left
-                if x2 < max_right:
-                    x2 = max_right
-                modified_pred['boxes'][position] = [x1, y1, x2, y2]
+    for pool_index, element_indices in pred['pool_dict'].items():
+        position = find_position(pool_index, modified_pred['BPMN_id'])
+        if position >= len(modified_pred['boxes']):
+            continue
+        x1, y1, x2, y2 = modified_pred['boxes'][position]
+        if x1 > min_left:
+            x1 = min_left
+        if x2 < max_right:
+            x2 = max_right
+        modified_pred['boxes'][position] = [x1, y1, x2, y2]
+
+# Main function to align boxes
+def align_boxes(pred, size, class_dict):
+    modified_pred = copy.deepcopy(pred)
+    pool_groups = calculate_centers_and_group_by_pool(pred, class_dict)
+    align_elements_within_pool(modified_pred, pool_groups, class_dict, size)
+    
+    if len(pred['pool_dict']) > 1:
+        expand_pool_bounding_boxes(modified_pred, pred, size)
+        adjust_pool_boundaries(modified_pred, pred)
 
     return modified_pred['boxes']
 
+
 # Function to create a BPMN XML file from prediction results
 def create_XML(full_pred, text_mapping, size_scale, scale):
     namespaces = {
@@ -184,7 +198,7 @@ def create_XML(full_pred, text_mapping, size_scale, scale):
     bpmnplane = ET.SubElement(bpmndi, 'bpmndi:BPMNPlane', id='BPMNPlane_1', bpmnElement='collaboration_1')
 
     full_pred['boxes'] = rescale_boxes(scale, old_boxes['boxes'])
-    full_pred['boxes'] = align_boxes(full_pred, size_elements)
+    full_pred['boxes'] = align_boxes(full_pred, size_elements, class_dict)
          
 
     # Add diagram elements for each pool