change the model object, correct some bug and add warnings

app.py

@@ -13,7 +13,7 @@ import gdown
 
 from modules.htlm_webpage import display_bpmn_xml
 from modules.OCR import text_prediction, filter_text, mapping_text, rescale
-from modules.utils import class_dict, arrow_dict, object_dict, find_closest_object
+from modules.utils import class_dict, arrow_dict, object_dict
 from modules.toXML import calculate_pool_bounds, add_diagram_elements, create_bpmn_object, create_flow_element
 from modules.display import draw_stream
 from modules.eval import full_prediction
@@ -49,6 +49,8 @@ def modif_box_pos(pred, size):
             modified_pred['boxes'][i] = [center[0] - size[label][0] / 2, center[1] - size[label][1] / 2, center[0] + size[label][0] / 2, center[1] + size[label][1] / 2]
     return modified_pred['boxes']
 
+
+
 # Function to create a BPMN XML file from prediction results
 def create_XML(full_pred, text_mapping, scale):
     namespaces = {
@@ -58,22 +60,24 @@ def create_XML(full_pred, text_mapping, scale):
         'dc': 'http://www.omg.org/spec/DD/20100524/DC',
         'xsi': 'http://www.w3.org/2001/XMLSchema-instance'
     }
+
     
     size_elements = {
-        'event': (43.2, 43.2),
-        'task': (120, 96),
-        'message': (43.2, 43.2),
-        'messageEvent': (43.2, 43.2),
-        'exclusiveGateway': (60, 60),
-        'parallelGateway': (60, 60),
-        'dataObject': (48, 72),
-        'dataStore': (72, 72),
-        'subProcess': (144, 108),
-        'eventBasedGateway': (60, 60),
-        'timerEvent': (48, 48),
+        'event': (st.session_state.size_scale*43.2, st.session_state.size_scale*43.2),
+        'task': (st.session_state.size_scale*120, st.session_state.size_scale*96),
+        'message': (st.session_state.size_scale*43.2, st.session_state.size_scale*43.2),
+        'messageEvent': (st.session_state.size_scale*43.2, st.session_state.size_scale*43.2),
+        'exclusiveGateway': (st.session_state.size_scale*60, st.session_state.size_scale*60),
+        'parallelGateway': (st.session_state.size_scale*60, st.session_state.size_scale*60),
+        'dataObject': ( st.session_state.size_scale*48, st.session_state.size_scale*72),
+        'dataStore': (st.session_state.size_scale*72, st.session_state.size_scale*72),
+        'subProcess': (st.session_state.size_scale*144, st.session_state.size_scale*108),
+        'eventBasedGateway': (st.session_state.size_scale*60, st.session_state.size_scale*60),
+        'timerEvent': (st.session_state.size_scale*48, st.session_state.size_scale*48),
     }
 
 
+
     definitions = ET.Element('bpmn:definitions', {
         'xmlns:xsi': namespaces['xsi'],
         'xmlns:bpmn': namespaces['bpmn'],
@@ -153,7 +157,7 @@ def load_models():
         model_arrow = get_arrow_model(len(arrow_dict),2)
 
         url_arrow = 'https://drive.google.com/uc?id=1vv1X_r_lZ8gnzMAIKxcVEb_T_Qb-NkyA'
-        url_object = 'https://drive.google.com/uc?id=1lmeVBRfTrcrBsq3904m4TgppSA4nqsIi'
+        url_object = 'https://drive.google.com/uc?id=1b1bqogxqdPS-SnvaOfWJGV1I1qOrTKh5'
 
         # Define paths to save models
         output_arrow = 'model_arrow.pth'
@@ -190,7 +194,7 @@ def prepare_image(image, pad=True, new_size=(1333, 1333)):
 
     if pad:
         enhancer = ImageEnhance.Brightness(image)
-        image = enhancer.enhance(1.5)  # Adjust the brightness if necessary
+        image = enhancer.enhance(1.0)  # Adjust the brightness if necessary
         # Pad the resized image to make it exactly the desired size
         padding = [0, 0, new_size[0] - new_scaled_size[0], new_size[1] - new_scaled_size[1]]
         image = F.pad(image, padding, fill=200, padding_mode='edge')
@@ -324,7 +328,7 @@ def main():
     st.sidebar.subheader("Instructions:")
     st.sidebar.text("1. Upload you image")
     st.sidebar.text("2. Crop the image \n  (try to put the BPMN diagram \n   in the center of the image)")
-    st.sidebar.text("3. Set the score threshold \n   for prediction (default is 0.6)")
+    st.sidebar.text("3. Set the score threshold \n   for prediction (default is 0.5)")
     st.sidebar.text("4. Click on 'Launch Prediction'")
     st.sidebar.text("5. You can now see the annotation \n   and the BPMN XML result")
     st.sidebar.text("6. You can change the scale for \n   the XML file (default is 1.0)")
@@ -410,7 +414,7 @@ def main():
                 with col1:
                     score_threshold = st.slider("Set score threshold for prediction", min_value=0.0, max_value=1.0, value=0.5, step=0.05)
             else:
-                score_threshold = st.slider("Set score threshold for prediction", min_value=0.0, max_value=1.0, value=0.5, step=0.05)
+                score_threshold = st.slider("Set score threshold for prediction", min_value=0.0, max_value=1.0, value=0.6, step=0.05)
             
             if st.button("Launch Prediction"):
                 st.session_state.crop_image = cropped_image
@@ -425,7 +429,11 @@ def main():
         with st.spinner('Waiting for BPMN modeler...'):
             col1, col2 = st.columns(2)
             with col1:
-                st.session_state.scale = st.slider("Set scale for XML file", min_value=0.1, max_value=2.0, value=1.0, step=0.1)
+                st.session_state.scale = st.slider("Set distance scale for XML file", min_value=0.1, max_value=2.0, value=1.0, step=0.1)
+                if is_mobile is False:
+                    st.session_state.size_scale = st.slider("Set size object scale for XML file", min_value=0.5, max_value=2.0, value=1.0, step=0.1)
+                else:
+                    st.session_state.size_scale = 1.0
             st.session_state.bpmn_xml = create_XML(st.session_state.prediction.copy(), st.session_state.text_mapping, st.session_state.scale)
             display_bpmn_xml(st.session_state.bpmn_xml, is_mobile=is_mobile, screen_width=int(4/5*screen_width))
 

modules/eval.py

@@ -172,14 +172,17 @@ def mix_predictions(objects_pred, arrow_pred):
 
     return boxes, labels, scores, keypoints
 
-def regroup_elements_by_pool(boxes, labels, class_dict):
+
+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.
@@ -187,14 +190,29 @@ def regroup_elements_by_pool(boxes, labels, class_dict):
     # Initialize a dictionary to hold the elements in each pool
     pool_dict = {}
 
+    # Filter out pools with IoU greater than the threshold
+    to_delete = []
+    for i in range(len(boxes)):
+        for j in range(i + 1, len(boxes)):
+            if labels[i] == labels[j] and labels[i] == list(class_dict.values()).index('pool'):
+                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_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)))
+        pool_dict[len(labels) - 1] = list(range(len(boxes)))
     else:
         # Initialize each pool index with an empty list
         for pool_index in pool_indices:
@@ -232,7 +250,8 @@ def regroup_elements_by_pool(boxes, labels, class_dict):
     # Merge non-empty pools followed by empty pools
     pool_dict = {**non_empty_pools, **empty_pools}
 
-    return pool_dict, labels
+    return pool_dict, boxes, labels, scores, keypoints
+
 
 
 def create_links(keypoints, boxes, labels, class_dict):
@@ -260,8 +279,7 @@ def create_links(keypoints, boxes, labels, class_dict):
 
     return links, best_points
 
-def correction_labels(boxes, labels, class_dict, pool_dict, flow_links):
- 
+def correction_labels(boxes, labels, class_dict, pool_dict, flow_links): 
     for pool_index, elements in pool_dict.items():
         print(f"Pool {pool_index} contains elements: {elements}")
         #check if the label sequenceflow is good
@@ -307,10 +325,12 @@ def last_correction(boxes, labels, scores, keypoints, links, best_points, pool_d
     #delete pool that are have only messageFlow on it
     delete_pool = []
     for pool_index, elements in pool_dict.items():
-        if all([labels[i] == list(class_dict.values()).index('messageFlow') for i in elements]):
+        if all([labels[i] in [list(class_dict.values()).index('messageFlow'),
+                              list(class_dict.values()).index('sequenceFlow'),
+                              list(class_dict.values()).index('dataAssociation')] for i in elements]):
             if len(elements) > 0:
                 delete_pool.append(pool_dict[pool_index])
-                print(f"Pool {pool_index} contains only messageFlow elements, deleting it")
+                print(f"Pool {pool_index} contains only arrow elements, deleting it")
 
     #sort index
     delete_pool = sorted(delete_pool, reverse=True)
@@ -371,7 +391,7 @@ def full_prediction(model_object, model_arrow, image, score_threshold=0.5, iou_t
         boxes, labels, scores, keypoints = mix_predictions(objects_pred, arrow_pred)
     
         # Regroup elements by pool
-        pool_dict, labels = regroup_elements_by_pool(boxes,labels, class_dict)
+        pool_dict, boxes, labels, scores, keypoints = regroup_elements_by_pool(boxes, labels, scores, keypoints, class_dict)
         # 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

modules/toXML.py

@@ -1,5 +1,5 @@
 import xml.etree.ElementTree as ET
-from modules.utils import class_dict
+from modules.utils import class_dict, error, warning
 
 def rescale(scale, boxes):
     for i in range(len(boxes)):
@@ -317,6 +317,17 @@ def calculate_waypoints(data, size, current_idx, source_id, target_id):
 
     source_idx = data['BPMN_id'].index(source_id)
     target_idx = data['BPMN_id'].index(target_id)
+
+    if source_idx==target_idx:
+        warning()
+        #return [data['keypoints'][current_idx][0][:2], data['keypoints'][current_idx][1][:2]]
+        return None
+
+    if source_idx is None or target_idx is None:
+        warning()
+        return [(source_x, source_y), (target_x, target_y)]
+
+
     name_source = source_id.split('_')[0]
     name_target = target_id.split('_')[0]
 
@@ -324,6 +335,10 @@ def calculate_waypoints(data, size, current_idx, source_id, target_id):
     source_x, source_y = data['boxes'][source_idx][:2]
     target_x, target_y = data['boxes'][target_idx][:2]
 
+    if name_source == 'pool' or name_target == 'pool':
+        warning()
+        return [(source_x, source_y), (target_x, target_y)]
+
     if pos_source == 'left':
         source_x = source_x
         source_y += size[name_source][1]/2
@@ -352,8 +367,13 @@ def calculate_waypoints(data, size, current_idx, source_id, target_id):
 
     return [(source_x, source_y), (target_x, target_y)]
 
-def create_flow_element(bpmn, text_mapping, idx, size, data, parent, message=False):
+def create_flow_element(bpmn, text_mapping, idx, size, data, parent, message=False):  
     source_idx, target_idx = data['links'][idx]
+
+    if source_idx is None or target_idx is None:
+        warning()
+        return
+
     source_id, target_id = data['BPMN_id'][source_idx], data['BPMN_id'][target_idx]
     if message:
         element_id = f'messageflow_{source_id}_{target_id}'
@@ -373,9 +393,13 @@ def create_flow_element(bpmn, text_mapping, idx, size, data, parent, message=Fal
             element = ET.SubElement(parent, 'bpmn:messageFlow', id=element_id, sourceRef=XML_source_id, targetRef=XML_target_id, name=text_mapping[data['BPMN_id'][idx]])
         else:
             waypoints = calculate_waypoints(data, size, idx, source_id, target_id)
+            if waypoints is None:
+                return
             element = ET.SubElement(parent, 'bpmn:messageFlow', id=element_id, sourceRef=source_id, targetRef=target_id, name=text_mapping[data['BPMN_id'][idx]])
     else:
         waypoints = calculate_waypoints(data, size, idx, source_id, target_id)
+        if waypoints is None:
+            return
         element = ET.SubElement(parent, 'bpmn:sequenceFlow', id=element_id, sourceRef=source_id, targetRef=target_id, name=text_mapping[data['BPMN_id'][idx]])
     add_diagram_edge(bpmn, element_id, waypoints)
     

modules/utils.py

@@ -17,6 +17,7 @@ import time
 from torch.optim import AdamW
 import copy
 from torchvision import transforms
+import streamlit as st
 
 
 object_dict = {
@@ -912,8 +913,9 @@ def find_closest_object(keypoint, boxes, labels):
     Returns:
     - int or None: The index of the closest object to the keypoint, or None if no object is found.
     """
-    min_distance = float('inf')
     closest_object_idx = None
+    best_point = None  
+    min_distance = float('inf')
     # Iterate over each bounding box
     for i, box in enumerate(boxes):
         if labels[i] in [list(class_dict.values()).index('sequenceFlow'),
@@ -943,3 +945,9 @@ def find_closest_object(keypoint, boxes, labels):
 
     return closest_object_idx, best_point
 
+
+def error():
+    st.error('There is an error in the detection', icon="🚨")
+
+def warning():
+    st.warning('Some element are not detected, verify your parameters', icon="⚠️")