add some file to reduce code in the app.py

app.py

@@ -1,155 +1,25 @@
 import streamlit as st
-from PIL import Image, ImageEnhance
-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 xml.dom import minidom
-from pathlib import Path
-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
-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
-from modules.train import get_faster_rcnn_model, get_arrow_model
-from streamlit_image_comparison import image_comparison
+from modules.utils import class_dict, rescale_boxes
+from modules.toXML import calculate_pool_bounds, add_diagram_elements, create_bpmn_object, create_flow_element, get_size_elements, definitions
 from streamlit_cropper import st_cropper
-from streamlit_drawable_canvas import st_canvas
 from streamlit_image_select import image_select
 from streamlit_js_eval import streamlit_js_eval
-
-def get_memory_usage():
-    process = psutil.Process()
-    mem_info = process.memory_info()
-    return mem_info.rss / (1024 ** 2)  # Return memory usage in MB
-
-def clear_memory():
-    st.session_state.clear()
-    gc.collect()
-
-# Function to read XML content from a file
-def read_xml_file(filepath):
-    """ Read XML content from a file """
-    with open(filepath, 'r', encoding='utf-8') as file:
-        return file.read()
-
-def modif_box_pos(pred, size):
-    modified_pred = copy.deepcopy(pred)  # Make a deep copy of the prediction
-
-    # Step 1: Calculate the center of each bounding box and group them by pool
-    pool_groups = {}
-    for pool_index, element_indices in pred['pool_dict'].items():
-        pool_groups[pool_index] = []
-        for i in element_indices:
-            if class_dict[modified_pred['labels'][i]] != 'dataObject' or class_dict[modified_pred['labels'][i]] != 'dataStore':
-                x1, y1, x2, y2 = modified_pred['boxes'][i]
-                center = [(x1 + x2) / 2, (y1 + y2) / 2]
-                pool_groups[pool_index].append((center, i))
-
-    # Function to group centers within a specified range
-    def group_centers(centers, axis, range_=50):
-        groups = []
-        while centers:
-            center, idx = centers.pop(0)
-            group = [(center, idx)]
-            for other_center, other_idx in centers[:]:
-                if abs(center[axis] - other_center[axis]) <= range_:
-                    group.append((other_center, other_idx))
-                    centers.remove((other_center, other_idx))
-            groups.append(group)
-        return groups
-
-    # Step 2: Align the elements within each pool
-    for pool_index, centers in pool_groups.items():
-        # Group bounding boxes by checking if their centers are within ±50 pixels on the y-axis
-        y_groups = group_centers(centers.copy(), axis=1)
-
-        # Align the y-coordinates of the centers of grouped bounding boxes
-        for group in y_groups:
-            avg_y = sum([c[0][1] for c in group]) / len(group)  # Calculate the average y-coordinate
-            for (center, idx) in group:
-                label = class_dict[modified_pred['labels'][idx]]
-                if label in size:
-                    new_center = (center[0], avg_y)  # Align the y-coordinate
-                    modified_pred['boxes'][idx] = [
-                        new_center[0] - size[label][0] / 2, 
-                        new_center[1] - size[label][1] / 2, 
-                        new_center[0] + size[label][0] / 2, 
-                        new_center[1] + size[label][1] / 2
-                    ]
-
-        # Recalculate centers after vertical alignment
-        centers = []
-        for group in y_groups:
-            for center, idx in group:
-                x1, y1, x2, y2 = modified_pred['boxes'][idx]
-                center = [(x1 + x2) / 2, (y1 + y2) / 2]
-                centers.append((center, idx))
-
-        # Group bounding boxes by checking if their centers are within ±50 pixels on the x-axis
-        x_groups = group_centers(centers.copy(), axis=0)
-
-        # Align the x-coordinates of the centers of grouped bounding boxes
-        for group in x_groups:
-            avg_x = sum([c[0][0] for c in group]) / len(group)  # Calculate the average x-coordinate
-            for (center, idx) in group:
-                label = class_dict[modified_pred['labels'][idx]]
-                if label in size:
-                    new_center = (avg_x, center[1])  # Align the x-coordinate
-                    modified_pred['boxes'][idx] = [
-                        new_center[0] - size[label][0] / 2, 
-                        modified_pred['boxes'][idx][1], 
-                        new_center[0] + size[label][0] / 2, 
-                        modified_pred['boxes'][idx][3]
-                    ]
-
-    return modified_pred['boxes']
-
+from modules.streamlit_utils import get_memory_usage, clear_memory, get_image, load_models, perform_inference, display_options, align_boxes, sidebar
 
 
 # Function to create a BPMN XML file from prediction results
-def create_XML(full_pred, text_mapping, scale):
-    namespaces = {
-        'bpmn': 'http://www.omg.org/spec/BPMN/20100524/MODEL',
-        'bpmndi': 'http://www.omg.org/spec/BPMN/20100524/DI',
-        'di': 'http://www.omg.org/spec/DD/20100524/DI',
-        'dc': 'http://www.omg.org/spec/DD/20100524/DC',
-        'xsi': 'http://www.w3.org/2001/XMLSchema-instance'
-    }
+def create_XML(full_pred, text_mapping, size_scale, scale):
 
+    size_elements = get_size_elements(size_scale)
     
-    size_elements = {
-        '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'],
-        'xmlns:bpmndi': namespaces['bpmndi'],
-        'xmlns:di': namespaces['di'],
-        'xmlns:dc': namespaces['dc'],
-        'targetNamespace': "http://example.bpmn.com",
-        'id': "simpleExample"
-    })
-
     #modify the boxes positions
     old_boxes = copy.deepcopy(full_pred)
 
@@ -165,8 +35,8 @@ def create_XML(full_pred, text_mapping, scale):
     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(scale, old_boxes['boxes'])
-    full_pred['boxes'] = modif_box_pos(full_pred, size_elements)
+    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()):
@@ -185,6 +55,7 @@ def create_XML(full_pred, text_mapping, scale):
         
         add_diagram_elements(bpmnplane, pool_id, min_x - 50, min_y - 50, pool_width, pool_height)
 
+
     # Create BPMN elements for each pool
     for idx, (pool_index, keep_elements) in enumerate(full_pred['pool_dict'].items()):
         create_bpmn_object(process[idx], bpmnplane, text_mapping, definitions, size_elements, full_pred, keep_elements)
@@ -206,163 +77,15 @@ def create_XML(full_pred, text_mapping, scale):
     reparsed = minidom.parseString(rough_string)
     pretty_xml_as_string = reparsed.toprettyxml(indent="  ")
 
-    full_pred['boxes'] = rescale(1/scale, full_pred['boxes'])
+    full_pred['boxes'] = rescale_boxes(1/scale, full_pred['boxes'])
     full_pred['boxes'] = old_boxes
 
     return pretty_xml_as_string
 
 
-# Function to load the models only once and use session state to keep track of it
-def load_models():
-    with st.spinner('Loading model...'):     
-        model_object = get_faster_rcnn_model(len(object_dict))
-        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=1b1bqogxqdPS-SnvaOfWJGV1I1qOrTKh5'
-
-        # Define paths to save models
-        output_arrow = 'model_arrow.pth'
-        output_object = 'model_object.pth'
-
-        # Download models using gdown
-        if not Path(output_arrow).exists():
-            # Download models using gdown
-            gdown.download(url_arrow, output_arrow, quiet=False)
-        else:
-            print('Model arrow downloaded from local')
-        if not Path(output_object).exists():
-            gdown.download(url_object, output_object, quiet=False)
-        else:
-            print('Model object downloaded from local')
-
-        # Load models
-        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        model_arrow.load_state_dict(torch.load(output_arrow, map_location=device))
-        model_object.load_state_dict(torch.load(output_object, map_location=device))
-
-        st.session_state.model_loaded = True
-        st.session_state.model_arrow = model_arrow
-        st.session_state.model_object = model_object
-
-        return model_object, model_arrow
-
-# Function to prepare the image for processing
-def prepare_image(image, pad=True, new_size=(1333, 1333)):
-    original_size = image.size
-    # Calculate scale to fit the new size while maintaining aspect ratio
-    scale = min(new_size[0] / original_size[0], new_size[1] / original_size[1])
-    new_scaled_size = (int(original_size[0] * scale), int(original_size[1] * scale))
-    # Resize image to new scaled size
-    image = F.resize(image, (new_scaled_size[1], new_scaled_size[0]))
-
-    if pad:
-        enhancer = ImageEnhance.Brightness(image)
-        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')
-
-    return new_scaled_size, image
-
-# Function to display various options for image annotation
-def display_options(image, score_threshold, is_mobile, screen_width):
-    col1, col2, col3, col4, col5 = st.columns(5)
-    with col1:
-        write_class = st.toggle("Write Class", value=True)
-        draw_keypoints = st.toggle("Draw Keypoints", value=True)
-        draw_boxes = st.toggle("Draw Boxes", value=True)
-    with col2:
-        draw_text = st.toggle("Draw Text", value=False)
-        write_text = st.toggle("Write Text", value=False)
-        draw_links = st.toggle("Draw Links", value=False)
-    with col3:
-        write_score = st.toggle("Write Score", value=True)
-        write_idx = st.toggle("Write Index", value=False)
-    with col4:
-        # Define options for the dropdown menu
-        dropdown_options = [list(class_dict.values())[i] for i in range(len(class_dict))]
-        dropdown_options[0] = 'all'
-        selected_option = st.selectbox("Show class", dropdown_options)
-
-    # Draw the annotated image with selected options
-    annotated_image = draw_stream(
-        np.array(image), prediction=st.session_state.prediction, text_predictions=st.session_state.text_pred,
-        draw_keypoints=draw_keypoints, draw_boxes=draw_boxes, draw_links=draw_links, draw_twins=False, draw_grouped_text=draw_text,
-        write_class=write_class, write_text=write_text, keypoints_correction=True, write_idx=write_idx, only_show=selected_option,
-        score_threshold=score_threshold, write_score=write_score, resize=True, return_image=True, axis=True
-    )
-
-    if is_mobile is True:
-        width = screen_width
-    else:
-        width = screen_width//2
-
-    # Display the original and annotated images side by side
-    image_comparison(
-        img1=annotated_image,
-        img2=image,
-        label1="Annotated Image",
-        label2="Original Image",
-        starting_position=99,
-        width=width,
-    )
-
-# Function to perform inference on the uploaded image using the loaded models
-def perform_inference(model_object, model_arrow, image, score_threshold, is_mobile, screen_width, iou_threshold=0.5, distance_treshold=30, percentage_text_dist_thresh=0.5):
-    _, uploaded_image = prepare_image(image, pad=False)
-              
-    img_tensor = F.to_tensor(prepare_image(image.convert('RGB'))[1])
-
-    # Display original image
-    if 'image_placeholder' not in st.session_state:
-        image_placeholder = st.empty()  # Create an empty placeholder
-    if is_mobile is False:
-        width = screen_width
-        if is_mobile is False:
-            width = screen_width//3
-        image_placeholder.image(uploaded_image, caption='Original Image', width=width)
-
-    # Prediction
-    _, st.session_state.prediction = full_prediction(model_object, model_arrow, img_tensor, score_threshold=score_threshold, iou_threshold=iou_threshold, distance_treshold=distance_treshold)
-
-    # Perform OCR on the uploaded image
-    ocr_results = text_prediction(uploaded_image)
-
-    # Filter and map OCR results to prediction results
-    st.session_state.text_pred = filter_text(ocr_results, threshold=0.6)
-    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)
-                
-    # Remove the original image display
-    image_placeholder.empty()
-
-    # Force garbage collection
-    gc.collect()
-
-@st.cache_data
-def get_image(uploaded_file):
-    return Image.open(uploaded_file).convert('RGB')
-
-
-
 def main():
     st.set_page_config(layout="wide")
 
-    # Apply CSS to change the sidebar width
-    st.markdown(
-        """
-        <style>
-        [data-testid="stSidebar"] {
-            width: 350px;
-        }
-        [data-testid="stSidebar"][aria-expanded="true"] {
-            width: 350px;
-        }
-        </style>
-        """,
-        unsafe_allow_html=True,
-    )
-
     screen_width = streamlit_js_eval(js_expressions='screen.width', want_output = True, key = 'SCR')
     print("Screen width:", screen_width)
 
@@ -387,23 +110,7 @@ def main():
         
 
 
-    # Sidebar content
-    st.sidebar.header("This BPMN AI model recognition is proposed by: \n ELCA in collaboration with EPFL.")
-    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.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)")
-    st.sidebar.text("7. You can modify and download \n   the result in right format")
-
-    st.sidebar.subheader("If there is an error, try to:")
-    st.sidebar.text("1. Change the score threshold")
-    st.sidebar.text("2. Re-crop the image by placing\n   the BPMN diagram in the center\n   of the image")
-    st.sidebar.text("3. Re-Launch the prediction")
-
-    st.sidebar.subheader("You can close this sidebar")
+    sidebar()  # Display the sidebar
 
     
     # Display current memory usage
@@ -497,7 +204,7 @@ def main():
                     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)
+            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))
 
         gc.collect()

modules/streamlit_utils.py

@@ -0,0 +1,262 @@
+import streamlit as st
+from PIL import Image, ImageEnhance
+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
+
+
+from modules.htlm_webpage import display_bpmn_xml
+from modules.OCR import text_prediction, filter_text, mapping_text
+from modules.utils import class_dict, arrow_dict, object_dict, rescale_boxes
+from modules.display import draw_stream
+from modules.eval import full_prediction
+from modules.train import get_faster_rcnn_model, get_arrow_model
+from streamlit_image_comparison import image_comparison
+
+
+
+def get_memory_usage():
+    process = psutil.Process()
+    mem_info = process.memory_info()
+    return mem_info.rss / (1024 ** 2)  # Return memory usage in MB
+
+def clear_memory():
+    st.session_state.clear()
+    gc.collect()
+
+def sidebar():# Sidebar content
+    st.sidebar.header("This BPMN AI model recognition is proposed by: \n ELCA in collaboration with EPFL.")
+    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.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)")
+    st.sidebar.text("7. You can modify and download \n   the result in right format")
+
+    st.sidebar.subheader("If there is an error, try to:")
+    st.sidebar.text("1. Change the score threshold")
+    st.sidebar.text("2. Re-crop the image by placing\n   the BPMN diagram in the center\n   of the image")
+    st.sidebar.text("3. Re-Launch the prediction")
+
+    st.sidebar.subheader("You can close this sidebar")
+
+
+# Function to read XML content from a file
+def read_xml_file(filepath):
+    """ Read XML content from a file """
+    with open(filepath, 'r', encoding='utf-8') as file:
+        return file.read()
+
+def align_boxes(pred, size):
+    modified_pred = copy.deepcopy(pred)  # Make a deep copy of the prediction
+
+    # Step 1: Calculate the center of each bounding box and group them by pool
+    pool_groups = {}
+    for pool_index, element_indices in pred['pool_dict'].items():
+        pool_groups[pool_index] = []
+        for i in element_indices:
+            if i > len(modified_pred['labels']):
+                continue 
+            if class_dict[modified_pred['labels'][i]] != 'dataObject' or class_dict[modified_pred['labels'][i]] != 'dataStore':
+                x1, y1, x2, y2 = modified_pred['boxes'][i]
+                center = [(x1 + x2) / 2, (y1 + y2) / 2]
+                pool_groups[pool_index].append((center, i))
+                
+    # Function to group centers within a specified range
+    def group_centers(centers, axis, range_=50):
+        groups = []
+        while centers:
+            center, idx = centers.pop(0)
+            group = [(center, idx)]
+            for other_center, other_idx in centers[:]:
+                if abs(center[axis] - other_center[axis]) <= range_:
+                    group.append((other_center, other_idx))
+                    centers.remove((other_center, other_idx))
+            groups.append(group)
+        return groups
+
+    # Step 2: Align the elements within each pool
+    for pool_index, centers in pool_groups.items():
+        # Group bounding boxes by checking if their centers are within ±50 pixels on the y-axis
+        y_groups = group_centers(centers.copy(), axis=1)
+
+        # Align the y-coordinates of the centers of grouped bounding boxes
+        for group in y_groups:
+            avg_y = sum([c[0][1] for c in group]) / len(group)  # Calculate the average y-coordinate
+            for (center, idx) in group:
+                label = class_dict[modified_pred['labels'][idx]]
+                if label in size:
+                    new_center = (center[0], avg_y)  # Align the y-coordinate
+                    modified_pred['boxes'][idx] = [
+                        new_center[0] - size[label][0] / 2, 
+                        new_center[1] - size[label][1] / 2, 
+                        new_center[0] + size[label][0] / 2, 
+                        new_center[1] + size[label][1] / 2
+                    ]
+
+        # Recalculate centers after vertical alignment
+        centers = []
+        for group in y_groups:
+            for center, idx in group:
+                x1, y1, x2, y2 = modified_pred['boxes'][idx]
+                center = [(x1 + x2) / 2, (y1 + y2) / 2]
+                centers.append((center, idx))
+
+        # Group bounding boxes by checking if their centers are within ±50 pixels on the x-axis
+        x_groups = group_centers(centers.copy(), axis=0)
+
+        # Align the x-coordinates of the centers of grouped bounding boxes
+        for group in x_groups:
+            avg_x = sum([c[0][0] for c in group]) / len(group)  # Calculate the average x-coordinate
+            for (center, idx) in group:
+                label = class_dict[modified_pred['labels'][idx]]
+                if label in size:
+                    new_center = (avg_x, center[1])  # Align the x-coordinate
+                    modified_pred['boxes'][idx] = [
+                        new_center[0] - size[label][0] / 2, 
+                        modified_pred['boxes'][idx][1], 
+                        new_center[0] + size[label][0] / 2, 
+                        modified_pred['boxes'][idx][3]
+                    ]
+    return modified_pred['boxes']
+
+
+
+# Function to load the models only once and use session state to keep track of it
+def load_models():
+    with st.spinner('Loading model...'):     
+        model_object = get_faster_rcnn_model(len(object_dict))
+        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=1b1bqogxqdPS-SnvaOfWJGV1I1qOrTKh5'
+
+        # Define paths to save models
+        output_arrow = 'model_arrow.pth'
+        output_object = 'model_object.pth'
+
+        # Download models using gdown
+        if not Path(output_arrow).exists():
+            # Download models using gdown
+            gdown.download(url_arrow, output_arrow, quiet=False)
+        else:
+            print('Model arrow downloaded from local')
+        if not Path(output_object).exists():
+            gdown.download(url_object, output_object, quiet=False)
+        else:
+            print('Model object downloaded from local')
+
+        # Load models
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        model_arrow.load_state_dict(torch.load(output_arrow, map_location=device))
+        model_object.load_state_dict(torch.load(output_object, map_location=device))
+
+        st.session_state.model_loaded = True
+        st.session_state.model_arrow = model_arrow
+        st.session_state.model_object = model_object
+
+        return model_object, model_arrow
+
+# Function to prepare the image for processing
+def prepare_image(image, pad=True, new_size=(1333, 1333)):
+    original_size = image.size
+    # Calculate scale to fit the new size while maintaining aspect ratio
+    scale = min(new_size[0] / original_size[0], new_size[1] / original_size[1])
+    new_scaled_size = (int(original_size[0] * scale), int(original_size[1] * scale))
+    # Resize image to new scaled size
+    image = F.resize(image, (new_scaled_size[1], new_scaled_size[0]))
+
+    if pad:
+        enhancer = ImageEnhance.Brightness(image)
+        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')
+
+    return new_scaled_size, image
+
+# Function to display various options for image annotation
+def display_options(image, score_threshold, is_mobile, screen_width):
+    col1, col2, col3, col4, col5 = st.columns(5)
+    with col1:
+        write_class = st.toggle("Write Class", value=True)
+        draw_keypoints = st.toggle("Draw Keypoints", value=True)
+        draw_boxes = st.toggle("Draw Boxes", value=True)
+    with col2:
+        draw_text = st.toggle("Draw Text", value=False)
+        write_text = st.toggle("Write Text", value=False)
+        draw_links = st.toggle("Draw Links", value=False)
+    with col3:
+        write_score = st.toggle("Write Score", value=True)
+        write_idx = st.toggle("Write Index", value=False)
+    with col4:
+        # Define options for the dropdown menu
+        dropdown_options = [list(class_dict.values())[i] for i in range(len(class_dict))]
+        dropdown_options[0] = 'all'
+        selected_option = st.selectbox("Show class", dropdown_options)
+
+    # Draw the annotated image with selected options
+    annotated_image = draw_stream(
+        np.array(image), prediction=st.session_state.prediction, text_predictions=st.session_state.text_pred,
+        draw_keypoints=draw_keypoints, draw_boxes=draw_boxes, draw_links=draw_links, draw_twins=False, draw_grouped_text=draw_text,
+        write_class=write_class, write_text=write_text, keypoints_correction=True, write_idx=write_idx, only_show=selected_option,
+        score_threshold=score_threshold, write_score=write_score, resize=True, return_image=True, axis=True
+    )
+
+    if is_mobile is True:
+        width = screen_width
+    else:
+        width = screen_width//2
+
+    # Display the original and annotated images side by side
+    image_comparison(
+        img1=annotated_image,
+        img2=image,
+        label1="Annotated Image",
+        label2="Original Image",
+        starting_position=99,
+        width=width,
+    )
+
+# Function to perform inference on the uploaded image using the loaded models
+def perform_inference(model_object, model_arrow, image, score_threshold, is_mobile, screen_width, iou_threshold=0.5, distance_treshold=30, percentage_text_dist_thresh=0.5):
+    _, uploaded_image = prepare_image(image, pad=False)
+              
+    img_tensor = F.to_tensor(prepare_image(image.convert('RGB'))[1])
+
+    # Display original image
+    if 'image_placeholder' not in st.session_state:
+        image_placeholder = st.empty()  # Create an empty placeholder
+    if is_mobile is False:
+        width = screen_width
+        if is_mobile is False:
+            width = screen_width//3
+        image_placeholder.image(uploaded_image, caption='Original Image', width=width)
+
+    # Prediction
+    _, st.session_state.prediction = full_prediction(model_object, model_arrow, img_tensor, score_threshold=score_threshold, iou_threshold=iou_threshold, distance_treshold=distance_treshold)
+
+    # Perform OCR on the uploaded image
+    ocr_results = text_prediction(uploaded_image)
+
+    # Filter and map OCR results to prediction results
+    st.session_state.text_pred = filter_text(ocr_results, threshold=0.6)
+    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)
+                
+    # Remove the original image display
+    image_placeholder.empty()
+
+    # Force garbage collection
+    gc.collect()
+
+@st.cache_data
+def get_image(uploaded_file):
+    return Image.open(uploaded_file).convert('RGB')

modules/toXML.py

@@ -1,5 +1,41 @@
 import xml.etree.ElementTree as ET
 from modules.utils import class_dict, error, warning
+import streamlit as st
+
+namespaces = {
+    'bpmn': 'http://www.omg.org/spec/BPMN/20100524/MODEL',
+    'bpmndi': 'http://www.omg.org/spec/BPMN/20100524/DI',
+    'di': 'http://www.omg.org/spec/DD/20100524/DI',
+    'dc': 'http://www.omg.org/spec/DD/20100524/DC',
+    'xsi': 'http://www.w3.org/2001/XMLSchema-instance'
+}
+
+
+definitions = ET.Element('bpmn:definitions', {
+    'xmlns:xsi': namespaces['xsi'],
+    'xmlns:bpmn': namespaces['bpmn'],
+    'xmlns:bpmndi': namespaces['bpmndi'],
+    'xmlns:di': namespaces['di'],
+    'xmlns:dc': namespaces['dc'],
+    'targetNamespace': "http://example.bpmn.com",
+    'id': "simpleExample"
+})
+
+def get_size_elements(size_scale):
+    size_elements = {
+        'event': (size_scale*43.2, size_scale*43.2),
+        'task': (size_scale*120, size_scale*96),
+        'message': (size_scale*43.2, size_scale*43.2),
+        'messageEvent': (size_scale*43.2, size_scale*43.2),
+        'exclusiveGateway': (size_scale*60, size_scale*60),
+        'parallelGateway': (size_scale*60, size_scale*60),
+        'dataObject': (size_scale*48, size_scale*72),
+        'dataStore': (size_scale*72, size_scale*72),
+        'subProcess': (size_scale*144, size_scale*108),
+        'eventBasedGateway': (size_scale*60, size_scale*60),
+        'timerEvent': (size_scale*48, size_scale*48),
+    }
+    return size_elements
 
 def rescale(scale, boxes):
     for i in range(len(boxes)):
@@ -184,13 +220,10 @@ def create_bpmn_object(process, bpmnplane, text_mapping, definitions, size, data
                 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, datasAssociation_idx, element_id)
             if len(status) != 0:
-                print('ici')
                 for state, dataAssociation_idx in zip(status, datasAssociation_idx):
                     # Handle Data Input Association
                     if state == 'input':
-                        print('input')
                         dataObject_idx = links[dataAssociation_idx][0]
                         dataObject_name = elements[dataObject_idx]
                         dataObject_ref = f'DataObjectReference_{dataObject_name.split("_")[1]}'
@@ -200,7 +233,6 @@ def create_bpmn_object(process, bpmnplane, text_mapping, definitions, size, data
 
                     # Handle Data Output Association
                     elif state == 'output':
-                        print('output')
                         dataObject_idx = links[dataAssociation_idx][1]
                         dataObject_name = elements[dataObject_idx]
                         dataObject_ref = f'DataObjectReference_{dataObject_name.split("_")[1]}'