ready for demo

modules/display.py

@@ -52,10 +52,19 @@ def draw_stream(image,
     - resize (bool): Whether to resize annotations to fit the image size.
     """
 
+    #delete the global pool if it is the only one to show
+    """if len(prediction['pool_dict'])==1 and prediction['labels'][-1]==6:
+        pool_index = list(prediction['pool_dict'])[0]
+        if len(prediction['pool_dict'][pool_index])==(len(prediction['boxes'])-1):
+            prediction['boxes'] = prediction['boxes'][:-1]
+            prediction['labels'] = prediction['labels'][:-1]
+            prediction['scores'] = prediction['scores'][:-1]
+            prediction['keypoints'] = prediction['keypoints'][:-1]
+            prediction['links'] = prediction['links'][:-1]"""
+
     # Convert image to RGB (if not already in that format)
     if prediction is None:
         image = image.squeeze(0).permute(1, 2, 0).cpu().numpy()
-
     
     image_copy = image.copy()
     scale = max(image.shape[0], image.shape[1]) / 1000

modules/eval.py

@@ -2,8 +2,9 @@ import numpy as np
 import torch
 from modules.utils import class_dict, object_dict, arrow_dict, find_closest_object, find_other_keypoint, filter_overlap_boxes, iou
 from tqdm import tqdm
-from modules.toXML import create_BPMN_id
+from modules.toXML import get_size_elements, calculate_pool_bounds, create_BPMN_id
 from modules.utils import is_vertical
+import streamlit as st
 
 
 def non_maximum_suppression(boxes, scores, labels=None, iou_threshold=0.5):
@@ -207,11 +208,7 @@ def regroup_elements_by_pool(boxes, labels, scores, keypoints, class_dict, iou_t
     pool_boxes = [boxes[i] for i in pool_indices]
 
 
-    if not pool_indices:
-        # If no pools or lanes are detected, create a single pool with all elements
-        labels = np.append(labels, list(class_dict.values()).index('pool'))
-        pool_dict[len(labels) - 1] = list(range(len(boxes)))
-    else:
+    if pool_indices:
         # Initialize each pool index with an empty list
         for pool_index in pool_indices:
             pool_dict[pool_index] = []
@@ -250,10 +247,15 @@ def regroup_elements_by_pool(boxes, labels, scores, keypoints, class_dict, iou_t
                     count += 1
             if len(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)
+                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}")
-
+       
     # 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}
@@ -296,6 +298,8 @@ def correction_labels(boxes, labels, class_dict, pool_dict, flow_links):
     data_association_index = list(class_dict.values()).index('dataAssociation')
     data_object_index = list(class_dict.values()).index('dataObject')
     data_store_index = list(class_dict.values()).index('dataStore')
+    message_event_index = list(class_dict.values()).index('messageEvent')
+    senquence_flow_indexx = list(class_dict.values()).index('sequenceFlow')
 
     for pool_index, elements in pool_dict.items():
         print(f"Pool {pool_index} contains elements: {elements}")
@@ -326,9 +330,12 @@ def correction_labels(boxes, labels, class_dict, pool_dict, flow_links):
                 label2 = labels[id2]
                 if data_object_index in {label1, label2} or data_store_index in {label1, label2}:
                     continue
-                else:
+                elif message_event_index in {label1, label2}: 
                     print('Change the link from dataAssociation to messageFlow')
                     labels[i] = message_flow_index
+                else:
+                    print('Change the link from dataAssociation to sequenceFlow')
+                    labels[i] = senquence_flow_indexx
 
     return labels, flow_links
 
@@ -355,6 +362,10 @@ def last_correction(boxes, labels, scores, keypoints, links, best_points, pool_d
                 delete_pool.append(pool_index)
                 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")
+
 
     delete_elements = []
     # Check if there is an arrow that has the same links
@@ -374,7 +385,7 @@ def last_correction(boxes, labels, scores, keypoints, links, best_points, pool_d
     delete_elements = delete_elements + delete_pool
     #delete double value in delete_elements
     delete_elements = list(set(delete_elements))
-    
+
     boxes = np.delete(boxes, delete_elements, axis=0)
     labels = np.delete(labels, delete_elements)
     scores = np.delete(scores, delete_elements)
@@ -424,12 +435,16 @@ def generate_data(image, boxes, labels, scores, keypoints, flow_links, best_poin
             return data
 
 def develop_prediction(boxes, labels, scores, keypoints, class_dict, correction=True):
+    print("Lenghts 1", len(boxes), len(labels), len(scores), len(keypoints))
     pool_dict, boxes, labels, scores, keypoints = regroup_elements_by_pool(boxes, labels, scores, keypoints, class_dict)
+    print("Lenghts 2", len(boxes), len(labels), len(scores), len(keypoints))
     # Create links between elements
     flow_links, best_points = create_links(keypoints, boxes, labels, class_dict)
+    
     #Correct the labels of some sequenceflow that cross multiple pool
     if correction:
         labels, flow_links = correction_labels(boxes, labels, class_dict, pool_dict, flow_links)
+    
     #give a link to event to allow the creation of the BPMN id with start, indermediate and end event
     flow_links = give_link_to_element(flow_links, labels)
 
@@ -439,6 +454,7 @@ def develop_prediction(boxes, labels, scores, keypoints, class_dict, correction=
             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)  
+    print("Lenghts 3", len(boxes), len(labels), len(scores), len(keypoints))
 
     return boxes, labels, scores, keypoints, flow_links, best_points, pool_dict 
 
@@ -452,14 +468,13 @@ def full_prediction(model_object, model_arrow, image, score_threshold=0.5, iou_t
     with torch.no_grad():  # Disable gradient calculation for inference
         _, objects_pred = object_prediction(model_object, image, score_threshold=score_threshold, iou_threshold=0.1)
         _, arrow_pred = arrow_prediction(model_arrow, image, score_threshold=score_threshold, iou_threshold=iou_threshold, distance_treshold=distance_treshold)
-        
-        #print('Object prediction:', objects_pred)
-
 
+        st.session_state.arrow_pred = arrow_pred
+        
         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)
-
+        
         image = image.permute(1, 2, 0).cpu().numpy()
         image = (image * 255).astype(np.uint8)
     

modules/streamlit_utils.py

@@ -310,20 +310,20 @@ def modify_results(percentage_text_dist_thresh=0.5):
             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
         object_bboxes = []
-        object_labels = []
-        arrow_bboxes = []
-        arrow_labels = []
+        object_labels = []     
         for i in range(len(bboxes)):
             if labels[i] <= 12:
                 object_bboxes.append(bboxes[i])
                 object_labels.append(labels[i])
-            else:
-                arrow_bboxes.append(bboxes[i])
-                arrow_labels.append(labels[i])
-
-        original_obj_len = len(object_bboxes)
 
         uploaded_image = prepare_image(st.session_state.crop_image, new_size=(1333, 1333), pad=False)
 
@@ -339,34 +339,20 @@ def modify_results(percentage_text_dist_thresh=0.5):
             for i in range(len(bboxes)):
                 bboxes[i][2] = bboxes[i][2] + bboxes[i][0]
                 bboxes[i][3] = bboxes[i][3] + bboxes[i][1]
-            for i in range(len(arrow_bboxes)):
-                arrow_bboxes[i][2] = arrow_bboxes[i][2] + arrow_bboxes[i][0]
-                arrow_bboxes[i][3] = arrow_bboxes[i][3] + arrow_bboxes[i][1]
+
+            object_scores = []
+            object_keypoints = []
+            for i in range(len(new_labels)):
+                object_scores.append(1.0)
+                object_keypoints.append([[0, 0, 0], [0, 0, 0]])
 
             new_bbox = np.concatenate((bboxes, arrow_bboxes))
             new_lab = np.concatenate((new_lab, arrow_labels))
+            new_scores = np.concatenate((object_scores, arrow_score))
+            new_keypoints = np.concatenate((object_keypoints, arrow_keypoints))
 
-            scores = st.session_state.prediction['scores']
-            keypoints = st.session_state.prediction['keypoints']
-
-            #delete element in keypoints to make it match the new number of boxes
-            keypoints = keypoints.tolist()
-            scores = scores.tolist()
-
-            diff = original_obj_len-len(bboxes)
-            if diff > 0:
-                for i in range(diff):                  
-                            keypoints.pop(0)
-                            scores.pop(0)
-            elif diff < 0:
-                for i in range(-diff):
-                    keypoints.insert(0, [[0, 0, 0], [0, 0, 0]])
-                    scores.insert(0, 0.0)
-
-            keypoints = np.array(keypoints)
-            scores = np.array(scores)
-
-            boxes, labels, scores, keypoints, flow_links, best_points, pool_dict = develop_prediction(new_bbox, new_lab, scores, keypoints, class_dict, correction=False)
+            
+            boxes, labels, scores, keypoints, 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.text_mapping = mapping_text(st.session_state.prediction, st.session_state.text_pred, print_sentences=False, percentage_thresh=percentage_text_dist_thresh)

modules/toXML.py

@@ -79,32 +79,24 @@ def align_boxes(pred, size):
                     ]
 
     # Step 3: Expand the pool bounding boxes to fit the aligned elements
-    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:
-            min_x, min_y, max_x, max_y = calculate_pool_bounds(modified_pred['boxes'],modified_pred['labels'], keep_elements, size_elements)
-            marge = 50
-        else:
-            if pool_index >= len(modified_pred['boxes']):
-                print("Problem with the index")
-                continue
+    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      
+            
             min_x, min_y, max_x, max_y = modified_pred['boxes'][pool_index]
-            marge = 0
-
-        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
 
-        if pool_index >= len(modified_pred['boxes']):
-            new_box = np.array([min_x - marge, min_y - marge//2, min_x + pool_width + marge, min_y + pool_height + marge//2])
-            modified_pred['boxes'] = np.append(modified_pred['boxes'], [new_box], axis=0)
-        else:
             modified_pred['boxes'][pool_index] = [min_x -marge, min_y-marge//2, min_x+pool_width+marge, min_y+pool_height+marge//2]
 
-    min_left,max_right = 0, 0
-    if len(pred['pool_dict'])>1:
+        min_left,max_right = 0, 0
         for pool_index, element_indices in pred['pool_dict'].items():
                 if pool_index >= len(modified_pred['boxes']):
                     print(f"Problem with the index {pool_index} with a length of {len(modified_pred['boxes'])}")
@@ -152,6 +144,10 @@ def create_XML(full_pred, text_mapping, size_scale, scale):
     })
 
     size_elements = get_size_elements(size_scale)
+
+    #if there is no pool or lane, create a pool with all elements
+    if len(full_pred['pool_dict'])==0 or (len(full_pred['pool_dict'])==1 and list(full_pred['pool_dict'])[0]==len(full_pred['labels'])):
+        full_pred, text_mapping = create_big_pool(full_pred, text_mapping)
     
     #modify the boxes positions
     old_boxes = copy.deepcopy(full_pred)
@@ -217,6 +213,20 @@ def create_XML(full_pred, text_mapping, size_scale, scale):
 
     return pretty_xml_as_string
 
+def create_big_pool(full_pred, text_mapping):
+    # If no pools or lanes are detected, create a single pool with all elements
+    new_pool_index = len(full_pred['labels'])
+    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]
+    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')
+    text_mapping['pool_1'] = 'Process'
+    print(f"Created a new pool index {new_pool_index} with elements: {elements_pool}")
+    return full_pred, text_mapping
+
 def get_size_elements(size_scale):
     size_elements = {
         'event': (size_scale*43.2, size_scale*43.2),