correct bug and better pool format

.gitignore

@@ -23,3 +23,7 @@ result_bpmn.bpmn
 BPMN_creation.ipynb
 
 *.png
+
+*.ipynb
+
+*.pmw

modules/display.py

@@ -75,6 +75,9 @@ def draw_stream(image,
     new_scaled_size = (int(original_size[0] * scale_), int(original_size[1] * scale_))
 
     for i in range(len(prediction['boxes'])):
+        if only_show is not None and only_show != 'all':
+                if prediction['labels'][i] != list(class_dict.values()).index(only_show):
+                    continue
         box = prediction['boxes'][i]
         x1, y1, x2, y2 = box
         if resize:
@@ -83,9 +86,6 @@ def draw_stream(image,
         if score < score_threshold:
             continue
         if draw_boxes:
-            if only_show is not None and only_show != 'all':
-                if prediction['labels'][i] != list(class_dict.values()).index(only_show):
-                    continue
             #dont show the lanes
             if prediction['labels'][i] == list(class_dict.values()).index('lane'):
                     continue

modules/eval.py

@@ -240,21 +240,26 @@ def regroup_elements_by_pool(boxes, labels, scores, keypoints, class_dict, iou_t
                 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
-            elements_not_in_pool = [i for i in elements_not_in_pool if i not in elements_not_in_pool_to_delete]
+            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
-            if len(elements_not_in_pool) > 1 and 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:
                 new_pool_index = len(labels)
                 size_elements = get_size_elements(1)
-                box = calculate_pool_bounds(boxes,labels, elements_not_in_pool, size_elements)
+                box = calculate_pool_bounds(boxes,labels, new_elements_not_in_pool, size_elements)
                 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] = elements_not_in_pool
-                print(f"Created a new pool index {new_pool_index} with elements: {elements_not_in_pool}")
+                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}")
        
     # Separate empty pools
     non_empty_pools = {k: v for k, v in pool_dict.items() if v}
@@ -341,30 +346,35 @@ def correction_labels(boxes, labels, class_dict, pool_dict, flow_links):
 
 
 
-def last_correction(boxes, labels, scores, keypoints, links, best_points, pool_dict, limit_area=10000):
+def last_correction(boxes, labels, scores, keypoints, bpmn_id, links, best_points, pool_dict, limit_area=10000):
 
     #delete pool that are have only messageFlow on it
     delete_pool = []
     for pool_index, elements in pool_dict.items():
+        #find the position of the pool_index in the bpmn_id
+        if pool_index in bpmn_id:
+            position = bpmn_id.index(pool_index)
+        else:
+            continue 
         if all([labels[i] in [list(class_dict.values()).index('messageFlow'),
                               list(class_dict.values()).index('sequenceFlow'),
                               list(class_dict.values()).index('dataAssociation'),
                               list(class_dict.values()).index('lane')] for i in elements]):
             if len(elements) > 0:
-                delete_pool.append(pool_index)
-                print(f"Pool {pool_index} contains only arrow elements, deleting it")
+                delete_pool.append(position)
+                print(f"Pool {pool_index} contains only arrow elements, deleting it")   
 
         #calcul the area of the pool$
-        if pool_index < len(boxes):
-            pool = boxes[pool_index]
+        if position < len(boxes):
+            pool = boxes[position]
             area = (pool[2] - pool[0]) * (pool[3] - pool[1])
             if len(pool_dict)>1 and area < limit_area:
-                delete_pool.append(pool_index)
+                delete_pool.append(position)
                 print(f"Pool {pool_index} is too small, deleting it")     
 
-            if is_vertical(boxes[pool_index]):
-                delete_pool.append(pool_index)
-                print(f"Pool {pool_index} is vertical, deleting it")
+            if is_vertical(boxes[position]):
+                delete_pool.append(position)
+                print(f"Pool {position} is vertical, deleting it")
 
 
     delete_elements = []
@@ -390,6 +400,7 @@ def last_correction(boxes, labels, scores, keypoints, links, best_points, pool_d
     labels = np.delete(labels, delete_elements)
     scores = np.delete(scores, delete_elements)
     keypoints = np.delete(keypoints, delete_elements, axis=0)
+    bpmn_id = [point for i, point in enumerate(bpmn_id) if i not in delete_elements]
     links = np.delete(links, delete_elements, axis=0)
     best_points = [point for i, point in enumerate(best_points) if i not in delete_elements]
 
@@ -397,7 +408,7 @@ def last_correction(boxes, labels, scores, keypoints, links, best_points, pool_d
     for pool_index, elements in pool_dict.items():
         pool_dict[pool_index] = [i for i in elements if i not in delete_elements]
 
-    return boxes, labels, scores, keypoints, links, best_points, pool_dict
+    return boxes, labels, scores, keypoints, bpmn_id, links, best_points, pool_dict
 
 def give_link_to_element(links, labels):
     #give a link to event to allow the creation of the BPMN id with start, indermediate and end event
@@ -410,11 +421,11 @@ def give_link_to_element(links, labels):
         return links
 
 
-def generate_data(image, boxes, labels, scores, keypoints, flow_links, best_points, pool_dict, class_dict):
+def generate_data(image, boxes, labels, scores, keypoints, bpmn_id, flow_links, best_points, pool_dict):
             idx = []
             for i in range(len(labels)):
                 idx.append(i) 
-            bpmn_id = [class_dict[labels[i]] for i in range(len(labels))]   
+              
 
             data = {
                 'image': image,
@@ -429,8 +440,6 @@ def generate_data(image, boxes, labels, scores, keypoints, flow_links, best_poin
                 'BPMN_id': bpmn_id,
             }
 
-            # give a unique BPMN id to each element
-            data = create_BPMN_id(data)
 
             return data
 
@@ -438,6 +447,10 @@ def develop_prediction(boxes, labels, scores, keypoints, class_dict, correction=
 
     pool_dict, boxes, labels, scores, keypoints = regroup_elements_by_pool(boxes, labels, scores, keypoints, class_dict)
 
+    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)
     
@@ -453,9 +466,9 @@ def develop_prediction(boxes, labels, scores, keypoints, class_dict, correction=
         if labels[i] == list(class_dict.values()).index('dataStore'):
             labels[i] = list(class_dict.values()).index('dataObject')
     
-    boxes,labels,scores,keypoints,flow_links,best_points,pool_dict = last_correction(boxes,labels,scores,keypoints,flow_links,best_points, pool_dict)  
+    boxes,labels,scores,keypoints,bpmn_id, flow_links,best_points,pool_dict = last_correction(boxes,labels,scores,keypoints,bpmn_id,flow_links,best_points, pool_dict)  
 
-    return boxes, labels, scores, keypoints, flow_links, best_points, pool_dict 
+    return boxes, labels, scores, keypoints, bpmn_id, flow_links, best_points, pool_dict 
 
     
 
@@ -472,12 +485,12 @@ def full_prediction(model_object, model_arrow, image, score_threshold=0.5, iou_t
         
         boxes, labels, scores, keypoints = mix_predictions(objects_pred, arrow_pred)
         
-        boxes, labels, scores, keypoints, flow_links, best_points, pool_dict = develop_prediction(boxes, labels, scores, keypoints, class_dict)
+        boxes, labels, scores, keypoints, bpmn_id, flow_links, best_points, pool_dict = develop_prediction(boxes, labels, scores, keypoints, class_dict)
         
         image = image.permute(1, 2, 0).cpu().numpy()
         image = (image * 255).astype(np.uint8)
     
-        data = generate_data(image, boxes, labels, scores, keypoints, flow_links, best_points, pool_dict, class_dict)
+        data = generate_data(image, boxes, labels, scores, keypoints, bpmn_id, flow_links, best_points, pool_dict)
 
         return image, data
 

modules/streamlit_utils.py

@@ -363,9 +363,9 @@ def modify_results(percentage_text_dist_thresh=0.5):
             new_keypoints = np.concatenate((object_keypoints, arrow_keypoints))
 
             
-            boxes, labels, scores, keypoints, 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=False)
 
-            st.session_state.prediction = generate_data(st.session_state.prediction['image'], boxes, labels, scores, keypoints, flow_links, best_points, pool_dict, class_dict)
+            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()

modules/toXML.py

@@ -6,6 +6,16 @@ import copy
 from xml.dom import minidom
 import numpy as np
 
+def find_position(pool_index, BPMN_id):
+    #find the position of the pool_index in the bpmn_id
+    if pool_index in BPMN_id:
+        position = BPMN_id.index(pool_index)
+    else:
+        position = None
+        error(f"Problem with the pool index {pool_index} in the BPMN_id")
+
+    return position
+
 def align_boxes(pred, size):
     modified_pred = copy.deepcopy(pred)  # Make a deep copy of the prediction
 
@@ -85,23 +95,31 @@ def align_boxes(pred, size):
             if len(keep_elements) != 0:
                 marge = 50
             else:
-                marge = 0      
+                marge = 0     
+
+            position = find_position(pool_index, modified_pred['BPMN_id'])
             
-            min_x, min_y, max_x, max_y = modified_pred['boxes'][pool_index]
+            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
 
-            modified_pred['boxes'][pool_index] = [min_x -marge, min_y-marge//2, min_x+pool_width+marge, min_y+pool_height+marge//2]
+            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():
-                if pool_index >= len(modified_pred['boxes']):
+                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'][pool_index]
+                x1, y1, x2, y2 = modified_pred['boxes'][position]
                 left = x1
                 right = x2
                 if left < min_left:
@@ -110,15 +128,16 @@ def align_boxes(pred, size):
                     max_right = right
 
         for pool_index, element_indices in pred['pool_dict'].items():
-                if pool_index >= len(modified_pred['boxes']):
+                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'][pool_index]
+                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'][pool_index] = [x1, y1, x2, y2]
+                modified_pred['boxes'][position] = [x1, y1, x2, y2]
 
     return modified_pred['boxes']
 
@@ -157,27 +176,30 @@ def create_XML(full_pred, text_mapping, size_scale, scale):
 
     # Create BPMN process elements
     process = []
-    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', name=text_mapping[full_pred['BPMN_id'][list(full_pred['pool_dict'].keys())[idx]]]))
+    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'))
 
     bpmndi = ET.SubElement(definitions, 'bpmndi:BPMNDiagram', id='BPMNDiagram_1')
     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)
+         
 
     # Add diagram elements for each pool
     for idx, (pool_index, keep_elements) in enumerate(full_pred['pool_dict'].items()):
         pool_id = f'participant_{idx+1}'
-        pool = ET.SubElement(collaboration, 'bpmn:participant', id=pool_id, processRef=f'process_{idx+1}', name=text_mapping[full_pred['BPMN_id'][list(full_pred['pool_dict'].keys())[idx]]])
+        pool = ET.SubElement(collaboration, 'bpmn:participant', id=pool_id, processRef=f'process_{idx+1}', name=text_mapping[pool_index])
         
+        position = find_position(pool_index, full_pred['BPMN_id'])
         # Calculate the bounding box for the pool
         #if len(keep_elements) == 0:
-        if pool_index >= len(full_pred['boxes']):
+        if position >= len(full_pred['boxes']):
             print("Problem with the index")
             continue
-        min_x, min_y, max_x, max_y = full_pred['boxes'][pool_index]
+        min_x, min_y, max_x, max_y = full_pred['boxes'][position]
         pool_width = max_x - min_x
         pool_height = max_y - min_y
         
@@ -219,7 +241,7 @@ def create_big_pool(full_pred, text_mapping):
     size_elements = get_size_elements(st.session_state.size_scale)
     elements_pool = list(range(len(full_pred['boxes'])))
     min_x, min_y, max_x, max_y = calculate_pool_bounds(full_pred['boxes'],full_pred['labels'], elements_pool, size_elements)
-    box = [min_x-50, min_y-50, max_x+100, max_y+100]
+    box = [min_x, min_y, max_x, max_y]
     full_pred['boxes'] = np.append(full_pred['boxes'], [box], axis=0)
     full_pred['pool_dict'][new_pool_index] = elements_pool
     full_pred['BPMN_id'].append('pool_1')
@@ -251,7 +273,10 @@ def rescale(scale, boxes):
                             boxes[i][3]*scale]
     return boxes
 
-def create_BPMN_id(data):
+def create_BPMN_id(labels,pool_dict):
+
+    BPMN_id = [class_dict[labels[i]] for i in range(len(labels))] 
+
     enums = {
         'event': 1,
         'task': 1,
@@ -268,7 +293,7 @@ def create_BPMN_id(data):
         'eventBasedGateway': 1
     }
 
-    BPMN_name = [class_dict[label] for label in data['labels']]
+    BPMN_name = [class_dict[label] for label in labels]
 
     for idx, Bpmn_id in enumerate(BPMN_name):
         key = {
@@ -288,10 +313,17 @@ def create_BPMN_id(data):
         }.get(Bpmn_id, None)
 
         if key:
-            data['BPMN_id'][idx] = f'{key}_{enums[key]}'
+            BPMN_id[idx] = f'{key}_{enums[key]}'
             enums[key] += 1
+        
+    # Update the pool_dict keys with their corresponding BPMN_id values
+    updated_pool_dict = {}
+    for key, value in pool_dict.items():
+        if key < len(BPMN_id):
+            new_key = BPMN_id[key]
+            updated_pool_dict[new_key] = value
 
-    return data
+    return BPMN_id, updated_pool_dict
 
 
 
@@ -520,7 +552,7 @@ def calculate_pool_bounds(boxes, labels, keep_elements, size):
         max_x = max(max_x, x + element_width)
         max_y = max(max_y, y + element_height)
     
-    return min_x, min_y, max_x, max_y
+    return min_x-50, min_y-50, max_x+100, max_y+50
 
 
     

toWizard.py

@@ -0,0 +1,45 @@
+import xml.etree.ElementTree as ET
+from utils import class_dict
+
+def rescale(scale, boxes):
+    for i in range(len(boxes)):
+        boxes[i] = [boxes[i][0] * scale, boxes[i][1] * scale, boxes[i][2] * scale, boxes[i][3] * scale]
+    return boxes
+
+def create_BPMN_id(data):
+    enum_end, enum_start, enum_task, enum_sequence, enum_dataflow, enum_messflow, enum_messageEvent, enum_exclusiveGateway, enum_parallelGateway, enum_pool = 1, 1, 1, 1, 1, 1, 1, 1, 1, 1
+    BPMN_name = [class_dict[data['labels'][i]] for i in range(len(data['labels']))]
+    for idx, Bpmn_id in enumerate(BPMN_name):
+        if Bpmn_id == 'event':
+            if data['links'][idx][0] is not None and data['links'][idx][1] is None:
+                data['BPMN_id'][idx] = f'end_event_{enum_end}'
+                enum_end += 1
+            elif data['links'][idx][0] is None and data['links'][idx][1] is not None:
+                data['BPMN_id'][idx] = f'start_event_{enum_start}'
+                enum_start += 1
+        elif Bpmn_id == 'task' or Bpmn_id == 'dataObject':
+            data['BPMN_id'][idx] = f'task_{enum_task}'
+            enum_task += 1
+        elif Bpmn_id == 'sequenceFlow':
+            data['BPMN_id'][idx] = f'sequenceFlow_{enum_sequence}'
+            enum_sequence += 1
+        elif Bpmn_id == 'messageFlow':
+            data['BPMN_id'][idx] = f'messageFlow_{enum_messflow}'
+            enum_messflow += 1
+        elif Bpmn_id == 'messageEvent':
+            data['BPMN_id'][idx] = f'message_event_{enum_messageEvent}'
+            enum_messageEvent += 1
+        elif Bpmn_id == 'exclusiveGateway':
+            data['BPMN_id'][idx] = f'exclusiveGateway_{enum_exclusiveGateway}'
+            enum_exclusiveGateway += 1
+        elif Bpmn_id == 'parallelGateway':
+            data['BPMN_id'][idx] = f'parallelGateway_{enum_parallelGateway}'
+            enum_parallelGateway += 1
+        elif Bpmn_id == 'dataAssociation':
+            data['BPMN_id'][idx] = f'dataAssociation_{enum_sequence}'
+            enum_dataflow += 1
+        elif Bpmn_id == 'pool':
+            data['BPMN_id'][idx] = f'pool_{enum_pool}'
+            enum_pool += 1
+
+    return data