example image proposed

app.py

@@ -7,6 +7,8 @@ from PIL import Image, ImageEnhance
 from htlm_webpage import display_bpmn_xml
 import gc
 import psutil
+import copy
+
 
 from OCR import text_prediction, filter_text, mapping_text, rescale
 from train import prepare_model
@@ -22,6 +24,7 @@ from streamlit_image_comparison import image_comparison
 from xml.dom import minidom
 from streamlit_cropper import st_cropper
 from streamlit_drawable_canvas import st_canvas
+from streamlit_image_select import image_select
 from utils import find_closest_object
 from train import get_faster_rcnn_model, get_arrow_model
 import gdown
@@ -43,12 +46,13 @@ def read_xml_file(filepath):
 
 # Function to modify bounding box positions based on the given sizes
 def modif_box_pos(pred, size):
-    for i, (x1, y1, x2, y2) in enumerate(pred['boxes']):
+    modified_pred = copy.deepcopy(pred)  # Make a deep copy of the prediction
+    for i, (x1, y1, x2, y2) in enumerate(modified_pred['boxes']):
         center = [(x1 + x2) / 2, (y1 + y2) / 2]
-        label = class_dict[pred['labels'][i]]
+        label = class_dict[modified_pred['labels'][i]]
         if label in size:
-            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 pred
+            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):
@@ -69,7 +73,6 @@ def create_XML(full_pred, text_mapping, scale):
         'exclusiveGateway': (60, 60),
         'event': (43.2, 43.2),
         'parallelGateway': (60, 60),
-        'sequenceFlow': (180, 12),
         'dataObject': (48, 72),
         'dataStore': (72, 72),
         'subProcess': (144, 108),
@@ -89,7 +92,8 @@ def create_XML(full_pred, text_mapping, scale):
     })
 
     #modify the boxes positions
-    full_pred = modif_box_pos(full_pred, size_elements)
+    old_boxes = copy.deepcopy(full_pred)
+    full_pred['boxes'] = modif_box_pos(full_pred, size_elements)
 
     # Create BPMN collaboration element
     collaboration = ET.SubElement(definitions, 'bpmn:collaboration', id='collaboration_1')
@@ -144,6 +148,7 @@ def create_XML(full_pred, text_mapping, scale):
     pretty_xml_as_string = reparsed.toprettyxml(indent="  ")
 
     full_pred['boxes'] = rescale(1/scale, full_pred['boxes'])
+    full_pred['boxes'] = old_boxes
 
     return pretty_xml_as_string
 
@@ -314,8 +319,22 @@ def main():
     #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
-        uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+        if img_selected is not None:
+            uploaded_file = img_selected
+        else:
+            uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 
     # Display the uploaded image if the user has uploaded an image
     if uploaded_file is not None:
@@ -342,7 +361,7 @@ def main():
                 st.session_state.crop_image = cropped_image
                 with st.spinner('Processing...'):
                     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.session_state.prediction = modif_box_pos(st.session_state.prediction, object_dict)                    
                     st.balloons()
     else:
         #delete the prediction 

eval.py

@@ -239,10 +239,7 @@ def create_links(keypoints, boxes, labels, class_dict):
         if labels[i]==list(class_dict.values()).index('sequenceFlow') or labels[i]==list(class_dict.values()).index('messageFlow'):
             closest1, point_start = find_closest_object(keypoints[i][0], boxes, labels)
             closest2, point_end = find_closest_object(keypoints[i][1], boxes, labels)
-
-            print('closest1:', closest1, 'closest2:', closest2)
-            print('point_start:', point_start, 'point_end:', point_end)
-
+            
             if closest1 is not None and closest2 is not None:
                 best_points.append([point_start, point_end])
                 links.append([closest1, closest2])