pool align width and modification works better

modules/eval.py

@@ -239,8 +239,15 @@ def regroup_elements_by_pool(boxes, labels, scores, keypoints, class_dict, iou_t
                     elements_not_in_pool.append(i)
 
         if elements_not_in_pool:
-            new_pool_index = max(pool_dict.keys()) + 1
+            elements_not_in_pool_to_delete = []
+            new_pool_index = len(labels)
             labels = np.append(labels, list(class_dict.values()).index('pool'))
+            #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
+            elements_not_in_pool = [i for i in elements_not_in_pool if i not in elements_not_in_pool_to_delete]
             pool_dict[new_pool_index] = elements_not_in_pool
 
     # Separate empty pools
@@ -330,7 +337,8 @@ def last_correction(boxes, labels, scores, keypoints, links, best_points, pool_d
     for pool_index, elements in pool_dict.items():
         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]):
+                              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")
@@ -339,7 +347,6 @@ def last_correction(boxes, labels, scores, keypoints, links, best_points, pool_d
         if pool_index < len(boxes):
             pool = boxes[pool_index]
             area = (pool[2] - pool[0]) * (pool[3] - pool[1])
-            print("area: ",area)
             if len(pool_dict)>1 and area < limit_area:
                 delete_pool.append(pool_index)
                 print(f"Pool {pool_index} is too small, deleting it")

modules/streamlit_utils.py

@@ -4,8 +4,6 @@ import torch
 from torchvision.transforms import functional as F
 import gc
 import psutil
-import copy
-import xml.etree.ElementTree as ET
 import numpy as np
 from pathlib import Path
 import gdown
@@ -17,6 +15,17 @@ from modules.eval import full_prediction
 from modules.train import get_faster_rcnn_model, get_arrow_model
 from streamlit_image_comparison import image_comparison
 
+from streamlit_image_annotation import detection
+from modules.toXML import create_XML
+from modules.eval import develop_prediction, generate_data
+from modules.utils import class_dict, object_dict
+
+from modules.htlm_webpage import display_bpmn_xml
+from streamlit_cropper import st_cropper
+from streamlit_image_select import image_select
+from streamlit_js_eval import streamlit_js_eval
+
+
 
 
 def get_memory_usage():
@@ -186,4 +195,199 @@ def perform_inference(model_object, model_arrow, image, score_threshold, is_mobi
 
 @st.cache_data
 def get_image(uploaded_file):
-    return Image.open(uploaded_file).convert('RGB')
+    return Image.open(uploaded_file).convert('RGB')
+
+
+
+def configure_page():
+    st.set_page_config(layout="wide")
+    screen_width = streamlit_js_eval(js_expressions='screen.width', want_output=True, key='SCR')
+    is_mobile = screen_width is not None and screen_width < 800
+    return is_mobile, screen_width
+
+def display_banner(is_mobile):
+    if is_mobile:
+        st.image("./images/banner_mobile.png", use_column_width=True)
+    else:
+        st.image("./images/banner_desktop.png", use_column_width=True)
+
+def display_title(is_mobile):
+    title = "Welcome on the BPMN AI model recognition app"
+    if is_mobile:
+        title = "Welcome on the mobile version of BPMN AI model recognition app"
+    st.title(title)
+
+def display_sidebar():
+    sidebar()
+
+def initialize_session_state():
+    if 'pool_bboxes' not in st.session_state:
+        st.session_state.pool_bboxes = []
+    if 'model_object' not in st.session_state or 'model_arrow' not in st.session_state:
+        clear_memory()
+        load_models()
+
+def load_example_image():
+    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"
+        )
+        return img_selected
+
+def load_user_image(img_selected, is_mobile):
+    if img_selected == './images/none.jpg':
+        img_selected = None
+
+    if img_selected is not None:
+        uploaded_file = img_selected
+    else:
+        if is_mobile:
+            uploaded_file = st.file_uploader("Choose an image from my computer...", type=["jpg", "jpeg", "png"], accept_multiple_files=False)
+        else:
+            col1, col2 = st.columns(2)
+            with col1:
+                uploaded_file = st.file_uploader("Choose an image from my computer...", type=["jpg", "jpeg", "png"])
+
+    return uploaded_file
+
+def display_image(uploaded_file, screen_width, is_mobile):
+    
+    with st.spinner('Waiting for image display...'):
+        original_image = get_image(uploaded_file)
+        resized_image = original_image.resize((screen_width // 2, int(original_image.height * (screen_width // 2) / original_image.width)))
+
+        if not is_mobile:
+            cropped_image = crop_image(resized_image, original_image)
+        else:
+            st.image(resized_image, caption="Image", use_column_width=False, width=int(4/5 * screen_width))
+            cropped_image = original_image
+
+    return cropped_image
+
+def crop_image(resized_image, original_image):
+    marge = 10
+    cropped_box = st_cropper(
+        resized_image,
+        realtime_update=True,
+        box_color='#0000FF',
+        return_type='box',
+        should_resize_image=False,
+        default_coords=(marge, resized_image.width - marge, marge, resized_image.height - marge)
+    )
+    scale_x = original_image.width / resized_image.width
+    scale_y = original_image.height / resized_image.height
+    x0, y0, x1, y1 = int(cropped_box['left'] * scale_x), int(cropped_box['top'] * scale_y), int((cropped_box['left'] + cropped_box['width']) * scale_x), int((cropped_box['top'] + cropped_box['height']) * scale_y)
+    cropped_image = original_image.crop((x0, y0, x1, y1))
+    return cropped_image
+
+def get_score_threshold(is_mobile):
+    col1, col2 = st.columns(2)
+    with col1:
+        st.session_state.score_threshold = st.slider("Set score threshold for prediction", min_value=0.0, max_value=1.0, value=0.5 if not is_mobile else 0.6, step=0.05) 
+
+def launch_prediction(cropped_image, score_threshold, is_mobile, screen_width):
+    st.session_state.crop_image = cropped_image
+    with st.spinner('Processing...'):
+        perform_inference(
+            st.session_state.model_object, st.session_state.model_arrow, st.session_state.crop_image,
+            score_threshold, is_mobile, screen_width, iou_threshold=0.3, distance_treshold=30, percentage_text_dist_thresh=0.5
+        )
+        st.balloons()    
+    
+
+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']]
+        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']]
+
+        # Filter boxes and labels where label is less than 12
+        object_bboxes = []
+        object_labels = []
+        arrow_bboxes = []
+        arrow_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)
+
+        new_labels = 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])  
+            # Convert back to original format
+            bboxes = np.array([label['bbox'] for label in new_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]
+            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]
+
+            new_bbox = np.concatenate((bboxes, arrow_bboxes))
+            new_lab = np.concatenate((new_lab, arrow_labels))
+
+            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)
+
+            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)
+
+            st.rerun()
+
+
+def display_bpmn_modeler(is_mobile, screen_width):
+    with st.spinner('Waiting for BPMN modeler...'):
+        st.session_state.bpmn_xml = create_XML(
+            st.session_state.prediction.copy(), st.session_state.text_mapping, 
+            st.session_state.size_scale, st.session_state.scale
+        )
+        display_bpmn_xml(st.session_state.bpmn_xml, is_mobile=is_mobile, screen_width=int(4/5 * screen_width))
+
+def modeler_options(is_mobile):
+    if not is_mobile:
+        with st.expander("Options for BPMN modeler"):
+            col1, col2 = st.columns(2)
+            with col1:
+                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) 
+                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.scale = 1.0
+        st.session_state.size_scale = 1.0

modules/toXML.py

@@ -4,6 +4,7 @@ import streamlit as st
 from modules.utils import class_dict, rescale_boxes
 import copy
 from xml.dom import minidom
+import numpy as np
 
 def align_boxes(pred, size):
     modified_pred = copy.deepcopy(pred)  # Make a deep copy of the prediction
@@ -76,6 +77,44 @@ def align_boxes(pred, size):
                         new_center[0] + size[label][0] / 2, 
                         modified_pred['boxes'][idx][3]
                     ]
+
+    # 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)
+        else:
+            min_x, min_y, max_x, max_y = modified_pred['boxes'][pool_index]
+        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 - 50, min_y - 50, min_x + pool_width + 50, min_y + pool_height + 50])
+            modified_pred['boxes'] = np.append(modified_pred['boxes'], [new_box], axis=0)
+        else:
+            modified_pred['boxes'][pool_index] = [min_x -50, min_y-50, min_x+pool_width+50, min_y+pool_height+50]
+
+    min_left,max_right = 0, 0
+    if len(pred['pool_dict'])>1:
+        for pool_index, element_indices in pred['pool_dict'].items():
+                x1, y1, x2, y2 = modified_pred['boxes'][pool_index]
+                left = x1
+                right = x2
+                if left < min_left:
+                    min_left = left
+                if right > max_right:
+                    max_right = right
+
+        for pool_index, element_indices in pred['pool_dict'].items():
+                x1, y1, x2, y2 = modified_pred['boxes'][pool_index]
+                if x1 > min_left:
+                    x1 = min_left
+                if x2 < max_right:
+                    x2 = max_right
+                modified_pred['boxes'][pool_index] = [x1, y1, x2, y2]
+
     return modified_pred['boxes']
 
 # Function to create a BPMN XML file from prediction results
@@ -125,24 +164,12 @@ def create_XML(full_pred, text_mapping, size_scale, scale):
         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]]])
         
         # Calculate the bounding box for the pool
-        if len(keep_elements) == 0:
-            min_x, min_y, max_x, max_y = full_pred['boxes'][pool_index]
-            pool_width = max_x - min_x
-            pool_height = max_y - min_y
-            #check area
-            if pool_width < 400 or pool_height < 30:
-                print("The pool is too small, please add more elements or increase the scale")
-                continue
-        else:
-            min_x, min_y, max_x, max_y = calculate_pool_bounds(full_pred, keep_elements, size_elements)
-            pool_width = max_x - min_x + 100  # Adding padding
-            pool_height = max_y - min_y + 100  # Adding padding
-            #check area
-            if pool_width < 300 or pool_height < 30:
-                print("The pool is too small, please add more elements or increase the scale")
-                continue
+        #if len(keep_elements) == 0:
+        min_x, min_y, max_x, max_y = full_pred['boxes'][pool_index]
+        pool_width = max_x - min_x
+        pool_height = max_y - min_y
         
-        add_diagram_elements(bpmnplane, pool_id, min_x - 50, min_y - 50, pool_width, pool_height)
+        add_diagram_elements(bpmnplane, pool_id, min_x, min_y, pool_width, pool_height)
 
 
     # Create BPMN elements for each pool
@@ -446,21 +473,21 @@ def create_bpmn_object(process, bpmnplane, text_mapping, definitions, size, data
             add_diagram_elements(bpmnplane, element_id, x, y, size['timerEvent'][0], size['timerEvent'][1])
 
 
-def calculate_pool_bounds(data, keep_elements, size):
-    min_x = min_y = float('10000')
-    max_x = max_y = float('0')
-    
+def calculate_pool_bounds(boxes, labels, keep_elements, size):
+    min_x, min_y = float('inf'), float('inf')
+    max_x, max_y = float('-inf'), float('-inf')
+
     for i in keep_elements:
-        if i >= len(data['BPMN_id']):
+        if i >= len(labels):
             print("Problem with the index")
             continue
-        element = data['BPMN_id'][i]
-        if element is None or data['labels'][i] == 13 or data['labels'][i] == 14 or data['labels'][i] == 15 or data['labels'][i] == 7 or data['labels'][i] == 15: 
+        
+        element = labels[i]
+        if element in {None, 7, 13, 14, 15}: 
             continue
         
-        element_type = element.split('_')[0]
-        x, y = data['boxes'][i][:2]
-        element_width, element_height = size[element_type]
+        x, y = boxes[i][:2]
+        element_width, element_height = size[class_dict[labels[i]]]
         
         min_x = min(min_x, x)
         min_y = min(min_y, y)
@@ -469,6 +496,7 @@ def calculate_pool_bounds(data, keep_elements, size):
     
     return min_x, min_y, max_x, max_y
 
+
     
 def calculate_pool_waypoints(idx, data, size, source_idx, target_idx, source_element, target_element):
     # Get the bounding boxes of the source and target elements
@@ -489,9 +517,9 @@ def calculate_pool_waypoints(idx, data, size, source_idx, target_idx, source_ele
         element_mid_y = (element_box[1] + element_box[3]) / 2
         # Connect the pool's bottom or top side to the target element's top or bottom center
         if pool_box[3] < element_box[1]:  # Pool is above the target element
-            waypoints = [(element_mid_x, pool_box[3]-50), (element_mid_x, element_box[1])]
+            waypoints = [(element_mid_x, pool_box[3]), (element_mid_x, element_box[1])]
         else:  # Pool is below the target element
-            waypoints = [(element_mid_x, element_box[3]), (element_mid_x, pool_box[1]-50)]
+            waypoints = [(element_mid_x, element_box[3]), (element_mid_x, pool_box[1])]
     else:
         pool_box = target_box
         element_box = (source_box[0], source_box[1], source_box[0]+size[source_element][0], source_box[1]+size[source_element][1])
@@ -500,9 +528,9 @@ def calculate_pool_waypoints(idx, data, size, source_idx, target_idx, source_ele
 
         # Connect the element's bottom or top center to the pool's top or bottom side
         if pool_box[3] < element_box[1]:  # Pool is above the target element
-            waypoints = [(element_mid_x, element_box[1]), (element_mid_x, pool_box[3]-50)]
+            waypoints = [(element_mid_x, element_box[1]), (element_mid_x, pool_box[3])]
         else:  # Pool is below the target element
-            waypoints = [(element_mid_x, element_box[3]), (element_mid_x, pool_box[1]-50)]
+            waypoints = [(element_mid_x, element_box[3]), (element_mid_x, pool_box[1])]
 
     return waypoints