correct bug

app.py

@@ -2,246 +2,12 @@ import streamlit as st
 from torchvision.transforms import functional as F
 import gc
 import numpy as np
-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
-from streamlit_drawable_canvas import st_canvas
-from modules.streamlit_utils import *
-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, object_dict
-
-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 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(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=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))
-
-            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(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)
+from modules.streamlit_utils import *
 
-            #print('New prediction:', st.session_state.prediction['keypoints'])
-            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
 
 def main():
     is_mobile, screen_width = configure_page()