correction of infinit loop bug

app.py

@@ -7,6 +7,7 @@ from modules.streamlit_utils import *
 
 
 def main():
+    st.session_state.first_run = True
     is_mobile, screen_width = configure_page()
     display_banner(is_mobile)
     display_title(is_mobile)

modules/eval.py

@@ -173,20 +173,6 @@ def mix_predictions(objects_pred, arrow_pred):
 
 
 def regroup_elements_by_pool(boxes, labels, scores, keypoints, class_dict, iou_threshold=0.3):
-    """
-    Regroups elements by the pool they belong to, and creates a single new pool for elements that are not in any existing pool.
-    Filters out pools that have an IoU greater than the specified threshold.
-
-    Parameters:
-    - boxes (list): List of bounding boxes.
-    - labels (list): List of labels corresponding to each bounding box.
-    - class_dict (dict): Dictionary mapping class indices to class names.
-    - iou_threshold (float): IoU threshold for filtering pools.
-
-    Returns:
-    - dict: A dictionary where each key is a pool's index and the value is a list of elements within that pool.
-    """
-    # Initialize a dictionary to hold the elements in each pool
     pool_dict = {}
 
     # Filter out pools with IoU greater than the threshold
@@ -197,32 +183,25 @@ def regroup_elements_by_pool(boxes, labels, scores, keypoints, class_dict, iou_t
                 if iou(np.array(boxes[i]), np.array(boxes[j])) > iou_threshold:
                     to_delete.append(j)
 
-   
     boxes = np.delete(boxes, to_delete, axis=0)
     labels = np.delete(labels, to_delete)
     scores = np.delete(scores, to_delete)
     keypoints = np.delete(keypoints, to_delete, axis=0)
 
-    # Identify the bounding boxes of the pools
-    pool_indices = [i for i, label in enumerate(labels) if (class_dict[label.item()] == 'pool')]
+    pool_indices = [i for i, label in enumerate(labels) if class_dict[label.item()] == 'pool']
     pool_boxes = [boxes[i] for i in pool_indices]
 
-
     if pool_indices:
-        # Initialize each pool index with an empty list
         for pool_index in pool_indices:
             pool_dict[pool_index] = []
 
-        # Initialize a list for elements not in any pool
         elements_not_in_pool = []
 
-        # Iterate over all elements
         for i, box in enumerate(boxes):
-            if i in pool_indices or class_dict[labels[i]] == 'messageFlow' or class_dict[labels[i]] == 'pool':
-                continue  # Skip pool boxes themselves and messageFlow elements
             assigned_to_pool = False
+            if i in pool_indices or class_dict[labels[i]] in ['messageFlow', 'pool']:
+                continue
             for j, pool_box in enumerate(pool_boxes):
-                # Check if the element is within the pool's bounding box
                 if (box[0] >= pool_box[0] and box[1] >= pool_box[1] and 
                     box[2] <= pool_box[2] and box[3] <= pool_box[3]):
                     pool_index = pool_indices[j]
@@ -230,44 +209,44 @@ def regroup_elements_by_pool(boxes, labels, scores, keypoints, class_dict, iou_t
                     assigned_to_pool = True
                     break
             if not assigned_to_pool:
-                if class_dict[labels[i]] != 'messageFlow' and class_dict[labels[i]] != 'lane' or class_dict[labels[i]] != 'pool':
+                if class_dict[labels[i]] not in ['messageFlow', 'lane', 'pool']:
                     elements_not_in_pool.append(i)
 
-        if elements_not_in_pool:
-            elements_not_in_pool_to_delete = []
-            #find the messageflow,pool and lane in the elements_not_in_pool
-            for i in elements_not_in_pool:
-                if class_dict[labels[i]] == 'messageFlow' or class_dict[labels[i]] == 'lane' or class_dict[labels[i]] == 'pool':
-                    elements_not_in_pool_to_delete.append(i)
-            #delete the messageflow from the elements_not_in_pool
-            new_elements_not_in_pool = [i for i in elements_not_in_pool if i not in elements_not_in_pool_to_delete]
-            count = 0
-            for i in elements_not_in_pool:
-                if labels[i] != list(class_dict.values()).index('sequenceFlow') or labels[i] != list(class_dict.values()).index('messageFlow'):
-                    count += 1
-            #check if there is only sequenceFlow or messageFlow in the new pool
-            if all([labels[i] in [list(class_dict.values()).index('sequenceFlow'),
-                                    list(class_dict.values()).index('messageFlow'),
-                                    list(class_dict.values()).index('dataAssociation')] for i in new_elements_not_in_pool]):
-                print('The new pool contains only sequenceFlow or messageFlow') 
-            elif len(new_elements_not_in_pool) > 1 and count > 1:
+        if len(elements_not_in_pool) > 1:
+            new_elements_not_in_pool = [i for i in elements_not_in_pool if class_dict[labels[i]] not in ['messageFlow', 'lane', 'pool']]
+            
+            # Indices of relevant classes
+            sequence_flow_index = list(class_dict.values()).index('sequenceFlow')
+            message_flow_index = list(class_dict.values()).index('messageFlow')
+            data_association_index = list(class_dict.values()).index('dataAssociation')
+
+            if all(labels[i] in {sequence_flow_index, message_flow_index, data_association_index} for i in new_elements_not_in_pool):
+                print('The new pool contains only sequenceFlow, messageFlow, or dataAssociation') 
+
+            elif len(new_elements_not_in_pool) > 1:
                 new_pool_index = len(labels)
-                size_elements = get_size_elements(1)
-                box = calculate_pool_bounds(boxes,labels, new_elements_not_in_pool, size_elements)
+                box = calculate_pool_bounds(boxes, labels, new_elements_not_in_pool, None)
                 boxes = np.append(boxes, [box], axis=0)
                 labels = np.append(labels, list(class_dict.values()).index('pool'))
                 scores = np.append(scores, 1.0)
                 keypoints = np.append(keypoints, np.zeros((1, 2, 3)), axis=0)
                 pool_dict[new_pool_index] = new_elements_not_in_pool
                 print(f"Created a new pool index {new_pool_index} with elements: {new_elements_not_in_pool}")
+    else:
+        all_elements = [i for i in range(len(boxes))]
+        new_pool_index = len(labels)
+        box = calculate_pool_bounds(boxes, labels, all_elements, None)
+        boxes = np.append(boxes, [box], axis=0)
+        labels = np.append(labels, list(class_dict.values()).index('pool'))
+        scores = np.append(scores, 1.0)
+        keypoints = np.append(keypoints, np.zeros((1, 2, 3)), axis=0)
+        pool_dict[new_pool_index] = all_elements
+        print(f"Created a super pool {new_pool_index} with elements: {all_elements}")
        
-    # Separate empty pools
     non_empty_pools = {k: v for k, v in pool_dict.items() if v}
     empty_pools = {k: v for k, v in pool_dict.items() if not v}
-
-    # Merge non-empty pools followed by empty pools
     pool_dict = {**non_empty_pools, **empty_pools}
-
+    
     return pool_dict, boxes, labels, scores, keypoints
 
 
@@ -449,8 +428,6 @@ def develop_prediction(boxes, labels, scores, keypoints, class_dict, correction=
 
     bpmn_id, pool_dict = create_BPMN_id(labels,pool_dict)
 
-    print('Pool dict:', pool_dict)
-
     # Create links between elements
     flow_links, best_points = create_links(keypoints, boxes, labels, class_dict)
     

modules/streamlit_utils.py

@@ -315,20 +315,22 @@ def launch_prediction(cropped_image, score_threshold, is_mobile, screen_width):
 def modify_results(percentage_text_dist_thresh=0.5):
     with st.expander("Method and Style modification (beta version)"):
         label_list = list(object_dict.values())
-        bboxes = [[int(coord) for coord in box] for box in st.session_state.prediction['boxes']]
+        if st.session_state.prediction['labels'][-1] == 6:
+            bboxes = [[int(coord) for coord in box] for box in st.session_state.prediction['boxes'][:-1]]
+            labels = [int(label) for label in st.session_state.prediction['labels'][:-1]]
+        else:
+            bboxes = [[int(coord) for coord in box] for box in st.session_state.prediction['boxes']]
+            labels = [int(label) for label in st.session_state.prediction['labels']]
         for i in range(len(bboxes)):
             bboxes[i][2] = bboxes[i][2] - bboxes[i][0]
             bboxes[i][3] = bboxes[i][3] - bboxes[i][1]
-        labels = [int(label) for label in st.session_state.prediction['labels']]
-
-        #st.session_state.arrow_predic = [boxes, labels, scores, keypoints]
 
         arrow_bboxes = st.session_state.arrow_pred['boxes']
         arrow_labels = st.session_state.arrow_pred['labels']
         arrow_score = st.session_state.arrow_pred['scores']
         arrow_keypoints = st.session_state.arrow_pred['keypoints']
 
-        # Filter boxes and labels where label is less than 12
+        # Filter boxes and labels where label is less than 12 to only have objects
         object_bboxes = []
         object_labels = []     
         for i in range(len(bboxes)):
@@ -338,22 +340,58 @@ def modify_results(percentage_text_dist_thresh=0.5):
 
         uploaded_image = prepare_image(st.session_state.crop_image, new_size=(1333, 1333), pad=False)
 
-        new_labels = detection(
+        new_data = detection(
             image=uploaded_image, bboxes=object_bboxes, labels=object_labels, 
             label_list=label_list, line_width=3, width=2000, use_space=False
         )
 
-        if new_labels is not None:
-            new_lab = np.array([label['label_id'] for label in new_labels])  
+        if new_data is not None:
+            changes = False
+            new_lab = np.array([data['label_id'] for data in new_data])  
             # Convert back to original format
-            bboxes = np.array([label['bbox'] for label in new_labels])
+            bboxes = np.array([data['bbox'] for data in new_data])
+            object_bboxes = np.array(object_bboxes)
+
+            # Order bboxes and labels
+            order = np.argsort(bboxes[:, 0])
+            bboxes = bboxes[order]
+            new_lab = new_lab[order]
+
+            order2 = np.argsort(object_bboxes[:, 0])
+            object_bboxes = object_bboxes[order2]
+            object_labels = np.array(object_labels)[order2]
+
+            # Make all values of bboxes integers
+            bboxes = bboxes.astype(int)
+
+            tolerance = 1
+
+            object_labels = np.array(object_labels)
+
+
+            if len(object_bboxes) == len(bboxes):
+                # Calculate absolute differences
+                abs_diff = np.abs(object_bboxes - bboxes)
+                
+                for i in range(len(object_bboxes)):
+                    for j in range(len(object_bboxes[i])):
+                        if abs_diff[i][j] > tolerance:
+                            changes = True
+                            break
+
+                #check if labels are the same
+                if not np.array_equal(object_labels, new_lab):
+                    changes = True
+            else:   
+                changes = True                
+
             for i in range(len(bboxes)):
                 bboxes[i][2] = bboxes[i][2] + bboxes[i][0]
                 bboxes[i][3] = bboxes[i][3] + bboxes[i][1]
 
             object_scores = []
             object_keypoints = []
-            for i in range(len(new_labels)):
+            for i in range(len(new_data)):
                 object_scores.append(1.0)
                 object_keypoints.append([[0, 0, 0], [0, 0, 0]])
 
@@ -363,12 +401,15 @@ def modify_results(percentage_text_dist_thresh=0.5):
             new_keypoints = np.concatenate((object_keypoints, arrow_keypoints))
 
             
-            boxes, labels, scores, keypoints, bpmn_id, flow_links, best_points, pool_dict = develop_prediction(new_bbox, new_lab, new_scores, new_keypoints, class_dict, correction=False)
+            boxes, labels, scores, keypoints, bpmn_id, flow_links, best_points, pool_dict = develop_prediction(new_bbox, new_lab, new_scores, new_keypoints, class_dict, correction=True)
 
             st.session_state.prediction = generate_data(st.session_state.prediction['image'], boxes, labels, scores, keypoints, bpmn_id, flow_links, best_points, pool_dict)
             st.session_state.text_mapping = mapping_text(st.session_state.prediction, st.session_state.text_pred, print_sentences=False, percentage_thresh=percentage_text_dist_thresh)
 
-            st.rerun()
+            if changes:
+                st.rerun()
+
+        
 
 
 def display_bpmn_modeler(is_mobile, screen_width):

modules/toXML.py

@@ -176,7 +176,6 @@ def create_XML(full_pred, text_mapping, size_scale, scale):
 
     # Create BPMN process elements
     process = []
-    print(full_pred['pool_dict'])
     for idx in range (len(full_pred['pool_dict'].items())):
         process_id = f'process_{idx+1}'      
         process.append(ET.SubElement(definitions, 'bpmn:process', id=process_id, isExecutable='false'))
@@ -544,15 +543,20 @@ def calculate_pool_bounds(boxes, labels, keep_elements, size):
         if element in {None, 7, 13, 14, 15}: 
             continue
         
-        x, y = boxes[i][:2]
-        element_width, element_height = size[class_dict[labels[i]]]
+
+        if size == None:
+            element_width = boxes[i][2] - boxes[i][0]
+            element_height = boxes[i][3] - boxes[i][1]
+        else:
+            element_width, element_height = size[class_dict[labels[i]]]
         
+        x, y = boxes[i][:2]
         min_x = min(min_x, x)
         min_y = min(min_y, y)
         max_x = max(max_x, x + element_width)
         max_y = max(max_y, y + element_height)
     
-    return min_x-50, min_y-50, max_x+100, max_y+50
+    return min_x-50, min_y-50, max_x+50, max_y+50