app.py

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

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()    

from modules.eval import develop_prediction, generate_data
from modules.utils import class_dict
def modify_results(percentage_text_dist_thresh=0.5):
    with st.expander("Method and Style modification"):
        label_list = list(class_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']]
        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, 
            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]

            scores = st.session_state.prediction['scores']
            keypoints = st.session_state.prediction['keypoints']
            #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)

            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'])


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):
    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) if not is_mobile else 1.0
        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) if not is_mobile else 1.0

def main():
    is_mobile, screen_width = configure_page()
    display_banner(is_mobile)
    display_title(is_mobile)
    display_sidebar()
    initialize_session_state()

    cropped_image = None

    img_selected = load_example_image()
    uploaded_file = load_user_image(img_selected, is_mobile)
    if uploaded_file is not None:
        cropped_image = display_image(uploaded_file, screen_width, is_mobile)

    if cropped_image is not None:
        get_score_threshold(is_mobile)
        if st.button("Launch Prediction"):
            launch_prediction(cropped_image, st.session_state.score_threshold, is_mobile, screen_width)
            st.rerun()
    
    if 'prediction' in st.session_state and uploaded_file:
        if st.button("🔄 Refresh image"):
            st.rerun()

        with st.expander("Show result"):
            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))

        #if not is_mobile:
            #modify_results()

        with st.expander("Options for BPMN modeler"):
            modeler_options(is_mobile)
        
        display_bpmn_modeler(is_mobile, screen_width)

    gc.collect()

if __name__ == "__main__":
    print('Starting the app...')
    main()