change is inside

modules/OCR.py

@@ -8,9 +8,8 @@ import numpy as np
 import networkx as nx
 from modules.utils import class_dict, proportion_inside
 import json
-from modules.utils import rescale_boxes as rescale 
+from modules.utils import rescale_boxes as rescale, is_vertical
 import streamlit as st
-from modules.utils import is_vertical
 
 VISION_KEY = os.getenv("VISION_KEY")
 VISION_ENDPOINT = os.getenv("VISION_ENDPOINT")
@@ -133,13 +132,6 @@ def min_distance_between_boxes(box1, box2):
                 min_dist = dist
     return min_dist
 
-
-def is_inside(box1, box2):
-    """Check if the center of box1 is inside box2."""
-    x_center = (box1[0] + box1[2]) / 2
-    y_center = (box1[1] + box1[3]) / 2
-    return box2[0] <= x_center <= box2[2] and box2[1] <= y_center <= box2[3]
-
 def are_close(box1, box2, threshold=50):
     """Determines if boxes are close based on their corners and center points."""
     corners1 = np.array([
@@ -307,12 +299,6 @@ def group_texts(task_boxes, text_boxes, texts, min_dist=50, iou_threshold=0.8, p
 
 def mapping_text(full_pred, text_pred, print_sentences=False,percentage_thresh=0.6,scale=1.0, iou_threshold=0.5):
 
-    ########### REFAIRE CETTE FONCTION ###########
-    #refaire la fonction pour qu'elle prenne en premier les elements qui sont dans les task et ensuite prendre un seuil de distance pour les autres elements
-    #ou sinon faire la distance entre les elements et non pas seulement les tasks
-
-
-     # Example usage
     boxes = rescale(scale, full_pred['boxes'])
 
     min_dist = 200

modules/streamlit_utils.py

@@ -81,6 +81,7 @@ def load_models():
             torch.save(model_arrow.state_dict(), output_arrow)
         elif 'model_arrow' not in st.session_state and Path(output_arrow).exists():
             model_arrow.load_state_dict(torch.load(output_arrow, map_location=device))
+            print()
             st.session_state.model_arrow = model_arrow
             print('Model arrow loaded from local file')
  
@@ -95,8 +96,9 @@ def load_models():
             torch.save(model_object.state_dict(), output_object)
         elif 'model_object' not in st.session_state and Path(output_object).exists():
             model_object.load_state_dict(torch.load(output_object, map_location=device))
+            print()
             st.session_state.model_object = model_object
-            print('Model object loaded from local file')
+            print('Model object loaded from local file\n')
 
 
         # Move models to device
@@ -184,14 +186,16 @@ def perform_inference(model_object, model_arrow, image, score_threshold, is_mobi
             width = screen_width//2
         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)
+
+    # 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)
+
+    #Mapping text to prediction
     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

modules/toXML.py

@@ -103,7 +103,7 @@ def expand_pool_bounding_boxes(modified_pred, pred, size_elements):
 
         position = find_position(pool_index, modified_pred['BPMN_id'])
 
-        if keep_elements == [] or position is None:
+        if keep_elements == [] and position is not None:
             min_x, min_y, max_x, max_y = modified_pred['boxes'][position]
         else:
             min_x, min_y, max_x, max_y = calculate_pool_bounds(modified_pred['boxes'], modified_pred['labels'], keep_elements, size_elements)
@@ -121,7 +121,7 @@ def adjust_pool_boundaries(modified_pred, pred):
     min_left, max_right = 0, 0
     for pool_index, element_indices in pred['pool_dict'].items():
         position = find_position(pool_index, modified_pred['BPMN_id'])
-        if position >= len(modified_pred['boxes']):
+        if position is None or position >= len(modified_pred['boxes']):
             continue
         x1, y1, x2, y2 = modified_pred['boxes'][position]
         left = x1
@@ -133,7 +133,7 @@ def adjust_pool_boundaries(modified_pred, pred):
 
     for pool_index, element_indices in pred['pool_dict'].items():
         position = find_position(pool_index, modified_pred['BPMN_id'])
-        if position >= len(modified_pred['boxes']):
+        if position is None or position >= len(modified_pred['boxes']):
             continue
         x1, y1, x2, y2 = modified_pred['boxes'][position]
         if x1 > min_left:
@@ -148,9 +148,9 @@ def align_boxes(pred, size, class_dict):
     pool_groups = calculate_centers_and_group_by_pool(pred, class_dict)
     align_elements_within_pool(modified_pred, pool_groups, class_dict, size)
     
-    if len(pred['pool_dict']) > 1:
-        expand_pool_bounding_boxes(modified_pred, pred, size)
-        adjust_pool_boundaries(modified_pred, pred)
+    #if len(pred['pool_dict']) > 1:
+        #expand_pool_bounding_boxes(modified_pred, pred, size)
+        #adjust_pool_boundaries(modified_pred, pred)
 
     return modified_pred['boxes']
 

modules/utils.py

@@ -57,6 +57,12 @@ class_dict = {
 }
 
 
+def is_inside(box1, box2):
+    """Check if the center of box1 is inside box2."""
+    x_center = (box1[0] + box1[2]) / 2
+    y_center = (box1[1] + box1[3]) / 2
+    return box2[0] <= x_center <= box2[2] and box2[1] <= y_center <= box2[3]
+
 def is_vertical(box):
     """Determine if the text in the bounding box is vertically aligned."""
     width = box[2] - box[0]