add verification of length and ready for demo

app.py

@@ -13,7 +13,7 @@ from glob import glob
 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
+from modules.utils import class_dict, object_dict
 
 def configure_page():
     st.set_page_config(layout="wide")
@@ -114,41 +114,114 @@ def launch_prediction(cropped_image, score_threshold, is_mobile, screen_width):
             score_threshold, is_mobile, screen_width, iou_threshold=0.3, distance_treshold=30, percentage_text_dist_thresh=0.5
         )
         st.balloons()    
+    
+def mix_new_pred(objects_pred, arrow_pred):
+    # Initialize the list of lists for keypoints
+    object_keypoints = []
+
+    # Number of boxes
+    num_boxes = len(objects_pred['boxes'])
+
+    # Iterate over the number of boxes
+    for _ in range(num_boxes):
+        # Each box has 2 keypoints, both initialized to [0, 0, 0]
+        keypoints = [[0, 0, 0], [0, 0, 0]]
+        object_keypoints.append(keypoints)
+
+    #concatenate the two predictions
+    boxes = np.concatenate((objects_pred['boxes'], arrow_pred['boxes']))
+    labels = np.concatenate((objects_pred['labels'], arrow_pred['labels']))
+
+    return boxes, labels, keypoints
 
 
 def modify_results(percentage_text_dist_thresh=0.5):
     with st.expander("Method and Style modification (beta version)"):
-        label_list = list(class_dict.values())
+        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
+        ignore_labels = [6, 7]
+        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])
+
+        print('Object bboxes:', object_bboxes)
+        print('Object labels:', object_labels)
+        print('Arrow bboxes:', arrow_bboxes)
+        print('Arrow labels:', arrow_labels)
+
+        original_obj_len = len(object_bboxes)
+
+
         uploaded_image = prepare_image(st.session_state.crop_image, new_size=(1333, 1333), pad=False)
         scale = 2000 / uploaded_image.size[0]
         new_labels = detection(
-            image=uploaded_image, bboxes=bboxes, labels=labels, 
+            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])
-  
+            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))
+
+            print('New labels:', new_lab)
 
             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
+            len_keypoints = len(keypoints)
+            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)
+
+            print('lenghts: ',len(bboxes), len(new_lab), len(scores), len(keypoints))
+            keypoints = np.array(keypoints)
+            scores = np.array(scores)
+
             #print('Old prediction:', st.session_state.prediction['keypoints'])
-            boxes, labels, scores, keypoints, flow_links, best_points, pool_dict = develop_prediction(bboxes, new_lab, scores, keypoints, class_dict, correction=False)
+            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)
 
             #print('New prediction:', st.session_state.prediction['keypoints'])
+            st.rerun()
 
 
 def display_bpmn_modeler(is_mobile, screen_width):
@@ -186,15 +259,16 @@ def main():
 
     if cropped_image is not None:
         get_score_threshold(is_mobile)
-        if st.button("Launch Prediction"):
+        if st.button("🚀 Launch Prediction"):
             launch_prediction(cropped_image, st.session_state.score_threshold, is_mobile, screen_width)
+            st.session_state.original_prediction = st.session_state.prediction.copy()
             st.rerun()
     
     if 'prediction' in st.session_state and uploaded_file:
-        if st.button("🔄 Refresh image"):
-            st.rerun()
+        #if st.button("🔄 Refresh image"):
+            #st.rerun()
 
-        with st.expander("Show result"):
+        with st.expander("Show result of prediction"):
             with st.spinner('Waiting for result display...'):
                 display_options(st.session_state.crop_image, st.session_state.score_threshold, is_mobile, int(5/6 * screen_width))
 

modules/display.py

@@ -94,6 +94,8 @@ def draw_stream(image,
         # Draw keypoints if available
         if draw_keypoints and 'keypoints' in prediction:
             for i in range(len(prediction['keypoints'])):
+                if i >= len(prediction['keypoints']):
+                    continue
                 kp = prediction['keypoints'][i]
                 for j in range(kp.shape[0]):
                     if prediction['labels'][i] != list(class_dict.values()).index('sequenceFlow') and prediction['labels'][i] != list(class_dict.values()).index('messageFlow') and prediction['labels'][i] != list(class_dict.values()).index('dataAssociation'):

modules/streamlit_utils.py

@@ -130,7 +130,7 @@ def display_options(image, score_threshold, is_mobile, screen_width):
 
     # 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,
+        np.array(image), prediction=st.session_state.original_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

modules/toXML.py

@@ -13,7 +13,7 @@ def align_boxes(pred, size):
     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']):
+            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]
@@ -138,7 +138,7 @@ def create_XML(full_pred, text_mapping, size_scale, scale):
             pool_width = max_x - min_x + 100  # Adding padding
             pool_height = max_y - min_y + 100  # Adding padding
             #check area
-            if pool_width < 400 or pool_height < 30:
+            if pool_width < 300 or pool_height < 30:
                 print("The pool is too small, please add more elements or increase the scale")
                 continue
         
@@ -157,6 +157,9 @@ def create_XML(full_pred, text_mapping, size_scale, scale):
     # Create sequence flow elements
     for idx, (pool_index, keep_elements) in enumerate(full_pred['pool_dict'].items()):
         for i in keep_elements:
+            if i >= len(full_pred['labels']):
+                print("Problem with the index")
+                continue
             if full_pred['labels'][i] == list(class_dict.values()).index('sequenceFlow'):
                 create_flow_element(bpmnplane, text_mapping, i, size_elements, full_pred, process[idx], message=False)
     
@@ -259,6 +262,8 @@ def add_diagram_edge(parent, element_id, waypoints):
         'id': element_id + '_di'
     })
     for x, y in waypoints:
+        if x is None or y is None:
+            return
         ET.SubElement(edge, 'di:waypoint', attrib={
             'x': str(x),
             'y': str(y)
@@ -312,6 +317,9 @@ def create_bpmn_object(process, bpmnplane, text_mapping, definitions, size, data
     links = data['links']
 
     for i in keep_elements:
+        if i >= len(elements):
+            print("Problem with the index")
+            continue
         element_id = elements[i]
 
         if element_id is None: