correction of bug with dataAssociation

app.py

@@ -249,7 +249,7 @@ def perform_inference(model_object, model_arrow, image, score_threshold):
     image_placeholder.image(uploaded_image, caption='Original Image', width=1000)
 
     # Prediction
-    _, st.session_state.prediction = full_prediction(model_object, model_arrow, img_tensor, score_threshold=score_threshold, iou_threshold=0.5)
+    _, st.session_state.prediction = full_prediction(model_object, model_arrow, img_tensor, score_threshold=score_threshold, iou_threshold=0.5, distance_treshold=30)
 
     # Perform OCR on the uploaded image
     ocr_results = text_prediction(uploaded_image)
@@ -311,21 +311,21 @@ def main():
 
     model_arrow = st.session_state.model_arrow
     model_object = st.session_state.model_object
+    
+    with st.expander("Use example images"):
+        img_selected = image_select("If you have no image and just want to test the demo, click on one of these images", ["./images/none.jpg", "./images/example1.jpg", "./images/example2.jpg", "./images/example3.jpg", "./images/example4.jpg"],
+                                captions=["None", "Example 1", "Example 2", "Example 3", "Example 4"], index=0, use_container_width=False, return_value="original")
+    
+    if img_selected== './images/none.jpg':
+        print('No example image selected')
+        #delete the prediction 
+        #if 'prediction' in st.session_state:
+            #del st.session_state['prediction']
+        img_selected = None
 
     #Create the layout for the app
     col1, col2 = st.columns(2)
     with col1:
-        with st.expander("Use example images"):
-            img_selected = image_select("If you have no image and just want to test the demo, click on one of these images", ["./images/none.jpg", "./images/example1.jpg", "./images/example2.jpg", "./images/example3.jpg"], 
-                                    captions=["None", "Example 1", "Example 2", "Example 3"], index=0, use_container_width=False, return_value="original")
-        
-        if img_selected== './images/none.jpg':
-            print('No example image selected')
-            #delete the prediction 
-            if 'prediction' in st.session_state:
-                del st.session_state['prediction']
-            img_selected = None
-
         # Create a file uploader for the user to upload an image
         if img_selected is not None:
             uploaded_file = img_selected
@@ -359,10 +359,10 @@ def main():
                     perform_inference(model_object, model_arrow, st.session_state.crop_image, score_threshold)
                     #st.session_state.prediction = modif_box_pos(st.session_state.prediction, object_dict)                    
                     st.balloons()
-    else:
+    #else:
         #delete the prediction 
-        if 'prediction' in st.session_state:
-            del st.session_state['prediction']
+        #if 'prediction' in st.session_state:
+            #del st.session_state['prediction']
 
 
     # If the prediction has been made and the user has uploaded an image, display the options for the user to annotate the image

modules/OCR.py

@@ -10,7 +10,7 @@ from modules.utils import class_dict, proportion_inside
 import json
 from modules.utils import rescale_boxes as rescale 
 import streamlit as st
-
+from modules.utils import is_vertical
 
 VISION_KEY = os.getenv("VISION_KEY")
 VISION_ENDPOINT = os.getenv("VISION_ENDPOINT")
@@ -188,11 +188,6 @@ def find_closest_box(text_box, all_boxes, labels, threshold, iou_threshold=0.5):
     return None
 
 
-def is_vertical(box):
-    """Determine if the text in the bounding box is vertically aligned."""
-    width = box[2] - box[0]
-    height = box[3] - box[1]
-    return (height > 2*width)
 
 def group_texts(task_boxes, text_boxes, texts, min_dist=50, iou_threshold=0.8, percentage_thresh=0.8):
     """Maps text boxes to task boxes and groups texts within each task based on proximity."""

modules/eval.py

@@ -3,6 +3,7 @@ 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.utils import is_vertical
 
 
 
@@ -75,7 +76,7 @@ def object_prediction(model, image, score_threshold=0.5, iou_threshold=0.5):
         for i in range(len(labels)):
             if labels[i] != list(object_dict.values()).index('task'):
                 continue
-            if boxes[i][2]-boxes[i][0] < boxes[i][3]-boxes[i][1]:
+            if is_vertical(boxes[i]):
                 vertical += 1
         horizontal = len(labels) - vertical
         for i in range(len(labels)):
@@ -83,12 +84,12 @@ def object_prediction(model, image, score_threshold=0.5, iou_threshold=0.5):
                 continue
 
             if vertical < horizontal:
-                if boxes[i][2]-boxes[i][0] < boxes[i][3]-boxes[i][1]:
+                if is_vertical(boxes[i]):
                     #find the element in the list and remove it
                     if i in selected_boxes:
                         selected_boxes.remove(i)
             elif vertical > horizontal:
-                if boxes[i][2]-boxes[i][0] > boxes[i][3]-boxes[i][1]:
+                if is_vertical(boxes[i]) == False:
                     #find the element in the list and remove it
                     if i in selected_boxes:
                         selected_boxes.remove(i)
@@ -261,19 +262,27 @@ 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 each link is in the same pool
+        #check if the label sequenceflow is good
         for i in range(len(flow_links)):
             if labels[i] == list(class_dict.values()).index('sequenceFlow'):
                 id1, id2 = flow_links[i]
-                if (id1 and id2) is not None:
+                if id1 is not None and id2 is not None:
+                    #check if each link is in the same pool
                     if id1 in elements and id2 in elements:
-                        continue
+                        #check if the link is between a dataObject or a dataStore
+                        if labels[id1] == 8 or labels[id2] == 8 or labels[id1] == 9 or labels[id2] == 9:
+                            print('change the link from sequenceFlow to dataAssociation')
+                            labels[i]=list(class_dict.values()).index('dataAssociation')
+                        else:
+                            continue
                     elif id1 not in elements and id2 not in elements:
                         continue
                     else:
                         print('change the link from sequenceFlow to messageFlow')
                         labels[i]=list(class_dict.values()).index('messageFlow')
 
+                    
+
     for i in range(len(labels)):
         #check if dataAssociation is connected to a dataObject
         if labels[i] == list(class_dict.values()).index('dataAssociation'):

modules/toXML.py

@@ -92,18 +92,35 @@ def check_status(link, keep_elements):
         return 'middle'
     
 def check_data_association(i, links, labels, keep_elements):
+    status, links_idx = [], []
     for j, (k,l) in enumerate(links):
-        if labels[j] == 14:
+        if labels[j] == list(class_dict.values()).index('dataAssociation'):
+            print('i,j,k,l',i,j,k,l)
             if k==i:
-                return 'output',j
+                status.append('output')
+                links_idx.append(j)
             elif l==i:
-                return 'input',j
-            
-    return 'no association', None
+                status.append('input')
+                links_idx.append(j)
 
-def create_data_Association(bpmn,data,size,element_id,source_id,target_id):
-    waypoints = calculate_waypoints(data, size, source_id, target_id)
+    return status, links_idx
+
+def create_data_Association(bpmn,data,size,element_id,current_idx,source_id,target_id):
+    waypoints = calculate_waypoints(data, size, current_idx, source_id, target_id)
     add_diagram_edge(bpmn, element_id, waypoints)
+
+def check_eventBasedGateway(i, links, labels):
+    status, links_idx = [], []
+    for j, (k,l) in enumerate(links):
+        if labels[j] == list(class_dict.values()).index('sequenceFlow'):
+            if k==i:
+                status.append('output')
+                links_idx.append(j)
+            elif l==i:
+                status.append('input')
+                links_idx.append(j)
+
+    return status, links_idx
         
 # Function to dynamically create and layout BPMN elements
 def create_bpmn_object(process, bpmnplane, text_mapping, definitions, size, data, keep_elements):
@@ -111,6 +128,8 @@ def create_bpmn_object(process, bpmnplane, text_mapping, definitions, size, data
     positions = data['boxes']
     links = data['links']
 
+    #print(links)
+
     for i in keep_elements:
         element_id = elements[i]
         if element_id is None:
@@ -127,25 +146,27 @@ def create_bpmn_object(process, bpmnplane, text_mapping, definitions, size, data
         # Task
         elif element_type == 'task':
             element = ET.SubElement(process, 'bpmn:task', id=element_id, name=text_mapping[element_id])
-            status, dataAssociation_idx = check_data_association(i, data['links'], data['labels'], keep_elements)
-
-            # Handle Data Input Association
-            if status == 'input':
-                dataObject_idx = links[dataAssociation_idx][0]
-                dataObject_name = elements[dataObject_idx]
-                dataObject_ref = f'DataObjectReference_{dataObject_name.split("_")[1]}'
-                sub_element = ET.SubElement(element, 'bpmn:dataInputAssociation', id=f'dataInputAssociation_{dataObject_ref.split("_")[1]}')
-                ET.SubElement(sub_element, 'bpmn:sourceRef').text = dataObject_ref
-                create_data_Association(bpmnplane, data, size, sub_element.attrib['id'], dataObject_name, element_id)
-
-            # Handle Data Output Association
-            elif status == 'output':
-                dataObject_idx = links[dataAssociation_idx][1]
-                dataObject_name = elements[dataObject_idx]
-                dataObject_ref = f'DataObjectReference_{dataObject_name.split("_")[1]}'
-                sub_element = ET.SubElement(element, 'bpmn:dataOutputAssociation', id=f'dataOutputAssociation_{dataObject_ref.split("_")[1]}')
-                ET.SubElement(sub_element, 'bpmn:targetRef').text = dataObject_ref
-                create_data_Association(bpmnplane, data, size, sub_element.attrib['id'], element_id, dataObject_name)
+            status, datasAssociation_idx = check_data_association(i, data['links'], data['labels'], keep_elements)
+
+            if len(status) != 0:
+                for state, dataAssociation_idx in zip(status, datasAssociation_idx):
+                    # Handle Data Input Association
+                    if state == 'input':
+                        dataObject_idx = links[dataAssociation_idx][0]
+                        dataObject_name = elements[dataObject_idx]
+                        dataObject_ref = f'DataObjectReference_{dataObject_name.split("_")[1]}'
+                        sub_element = ET.SubElement(element, 'bpmn:dataInputAssociation', id=f'dataInputAssociation_{dataObject_ref.split("_")[1]}')
+                        ET.SubElement(sub_element, 'bpmn:sourceRef').text = dataObject_ref
+                        create_data_Association(bpmnplane, data, size, sub_element.attrib['id'], dataAssociation_idx, dataObject_name, element_id)
+
+                    # Handle Data Output Association
+                    elif state == 'output':
+                        dataObject_idx = links[dataAssociation_idx][1]
+                        dataObject_name = elements[dataObject_idx]
+                        dataObject_ref = f'DataObjectReference_{dataObject_name.split("_")[1]}'
+                        sub_element = ET.SubElement(element, 'bpmn:dataOutputAssociation', id=f'dataOutputAssociation_{dataObject_ref.split("_")[1]}')
+                        ET.SubElement(sub_element, 'bpmn:targetRef').text = dataObject_ref
+                        create_data_Association(bpmnplane, data, size, sub_element.attrib['id'], dataAssociation_idx, element_id, dataObject_name)
 
             add_diagram_elements(bpmnplane, element_id, x, y, size['task'][0], size['task'][1])
 
@@ -158,6 +179,30 @@ def create_bpmn_object(process, bpmnplane, text_mapping, definitions, size, data
                 element = ET.SubElement(process, 'bpmn:intermediateCatchEvent', id=element_id, name=text_mapping[element_id])
             elif status == 'end':
                 element = ET.SubElement(process, 'bpmn:endEvent', id=element_id, name=text_mapping[element_id])
+
+            status, datasAssociation_idx = check_data_association(i, data['links'], data['labels'], keep_elements)
+            print('status',status)
+            print('datasAssociation_idx',datasAssociation_idx)
+            if len(status) != 0:
+                for state, dataAssociation_idx in zip(status, datasAssociation_idx):
+                    # Handle Data Input Association
+                    if state == 'input':
+                        dataObject_idx = links[dataAssociation_idx][0]
+                        dataObject_name = elements[dataObject_idx]
+                        dataObject_ref = f'DataObjectReference_{dataObject_name.split("_")[1]}'
+                        sub_element = ET.SubElement(element, 'bpmn:dataInputAssociation', id=f'dataInputAssociation_{dataObject_ref.split("_")[1]}')
+                        ET.SubElement(sub_element, 'bpmn:sourceRef').text = dataObject_ref
+                        create_data_Association(bpmnplane, data, size, sub_element.attrib['id'], dataAssociation_idx, dataObject_name, element_id)
+
+                    # Handle Data Output Association
+                    elif state == 'output':
+                        dataObject_idx = links[dataAssociation_idx][1]
+                        dataObject_name = elements[dataObject_idx]
+                        dataObject_ref = f'DataObjectReference_{dataObject_name.split("_")[1]}'
+                        sub_element = ET.SubElement(element, 'bpmn:dataOutputAssociation', id=f'dataOutputAssociation_{dataObject_ref.split("_")[1]}')
+                        ET.SubElement(sub_element, 'bpmn:targetRef').text = dataObject_ref
+                        create_data_Association(bpmnplane, data, size, sub_element.attrib['id'], dataAssociation_idx, element_id, dataObject_name)
+
             ET.SubElement(element, 'bpmn:messageEventDefinition', id=f'MessageEventDefinition_{i+1}')
             add_diagram_elements(bpmnplane, element_id, x, y, size['message'][0], size['message'][1])
 
@@ -172,13 +217,36 @@ def create_bpmn_object(process, bpmnplane, text_mapping, definitions, size, data
             element = ET.SubElement(process, f'bpmn:{gateway_type}', id=element_id)
             add_diagram_elements(bpmnplane, element_id, x, y, size[element_type][0], size[element_type][1])
 
+        elif element_type == 'eventBasedGateway':
+            element = ET.SubElement(process, 'bpmn:eventBasedGateway', id=element_id)
+            status, links_idx = check_eventBasedGateway(i, data['links'], data['labels'])
+
+            if len(status) != 0:
+                for state, link_idx in zip(status, links_idx):
+                    # Handle Data Input Association
+                    if state == 'input' :
+                        gateway_idx = links[link_idx][0]
+                        gateway_name = elements[gateway_idx]
+                        sub_element = ET.SubElement(element, 'bpmn:eventBasedGateway', id=f'eventBasedGateway{gateway_name.split("_")[1]}')
+                        create_data_Association(bpmnplane, data, size, sub_element.attrib['id'], i, gateway_name, element_id)
+
+                    # Handle Data Output Association
+                    elif state == 'output':
+                        gateway_idx = links[link_idx][1]
+                        gateway_name = elements[gateway_idx]
+                        sub_element = ET.SubElement(element, 'bpmn:eventBasedGateway', id=f'eventBasedGateway{gateway_name.split("_")[1]}')
+                        create_data_Association(bpmnplane, data, size, sub_element.attrib['id'], i, element_id, gateway_name)
+
+
+            add_diagram_elements(bpmnplane, element_id, x, y, size['eventBasedGateway'][0], size['eventBasedGateway'][1])
+
         # Data Object
-        elif element_type == 'dataObject':
+        elif element_type == 'dataObject' or element_type == 'dataStore':
             dataObject_idx = element_id.split('_')[1]
             dataObject_ref = f'DataObjectReference_{dataObject_idx}'
             element = ET.SubElement(process, 'bpmn:dataObjectReference', id=dataObject_ref, dataObjectRef=element_id, name=text_mapping[element_id])
-            ET.SubElement(process, 'bpmn:dataObject', id=element_id)
-            add_diagram_elements(bpmnplane, dataObject_ref, x, y, size['dataObject'][0], size['dataObject'][1])
+            ET.SubElement(process, f'bpmn:{element_type}', id=element_id)
+            add_diagram_elements(bpmnplane, dataObject_ref, x, y, size[element_type][0], size[element_type][1])
 
         # Timer Event
         elif element_type == 'timerEvent':
@@ -222,6 +290,8 @@ def calculate_pool_waypoints(idx, data, size, source_idx, target_idx, source_ele
 
     # Check if the connection involves a pool
     if source_element == 'pool':
+        if target_element == 'pool':
+            return [(source_mid_x, source_mid_y), (source_mid_x, source_mid_y)]
         pool_box = source_box
         element_box = (target_box[0], target_box[1], target_box[0]+size[target_element][0], target_box[1]+size[target_element][1])
         element_mid_x = (element_box[0] + element_box[2]) / 2
@@ -295,21 +365,24 @@ def create_flow_element(bpmn, text_mapping, idx, size, data, parent, message=Fal
     else:
         element_id = f'sequenceflow_{source_id}_{target_id}'
 
-    if source_id.split('_')[0] == 'pool' or target_id.split('_')[0] == 'pool':
-        waypoints = calculate_pool_waypoints(idx, data, size, source_idx, target_idx, source_id.split('_')[0], target_id.split('_')[0])
-        #waypoints = data['best_points'][idx]
-        if source_id.split('_')[0] == 'pool':
-            source_id = f"participant_{source_id.split('_')[1]}"
-        if target_id.split('_')[0] == 'pool':
-            target_id = f"participant_{target_id.split('_')[1]}"
-    else:
-        waypoints = calculate_waypoints(data, size, idx, source_id, target_id)
-        #waypoints = data['best_points'][idx]
-
-    #waypoints = data['best_points'][idx]
     if message:
-        element = ET.SubElement(parent, 'bpmn:messageFlow', id=element_id, sourceRef=source_id, targetRef=target_id, name=text_mapping[data['BPMN_id'][idx]])
+        if source_id.split('_')[0] == 'pool' or target_id.split('_')[0] == 'pool':
+            waypoints = calculate_pool_waypoints(idx, data, size, source_idx, target_idx, source_id.split('_')[0], target_id.split('_')[0])
+            if source_id.split('_')[0] == 'pool':
+                XML_source_id = f"participant_{source_id.split('_')[1]}"
+                XML_target_id = target_id
+            if target_id.split('_')[0] == 'pool':
+                XML_target_id = f"participant_{target_id.split('_')[1]}"
+                XML_source_id = source_id
+
+            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)
+            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)
         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

@@ -61,6 +61,13 @@ class_dict = {
     15: 'messageFlow',
 }
 
+
+def is_vertical(box):
+    """Determine if the text in the bounding box is vertically aligned."""
+    width = box[2] - box[0]
+    height = box[3] - box[1]
+    return (height > 2*width)
+
 def rescale_boxes(scale, boxes):
     for i in range(len(boxes)):
                 boxes[i] = [boxes[i][0]*scale,