alignment of the element

app.py

@@ -39,14 +39,76 @@ def read_xml_file(filepath):
     with open(filepath, 'r', encoding='utf-8') as file:
         return file.read()
 
-# Function to modify bounding box positions based on the given sizes
 def modif_box_pos(pred, size):
     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[modified_pred['labels'][i]]
-        if label in size:
-            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]
+
+    # Step 1: Calculate the center of each bounding box and group them by pool
+    pool_groups = {}
+    for pool_index, element_indices in pred['pool_dict'].items():
+        pool_groups[pool_index] = []
+        for i in element_indices:
+            if class_dict[modified_pred['labels'][i]] != 'dataObject' or class_dict[modified_pred['labels'][i]] != 'dataStore':
+                x1, y1, x2, y2 = modified_pred['boxes'][i]
+                center = [(x1 + x2) / 2, (y1 + y2) / 2]
+                pool_groups[pool_index].append((center, i))
+
+    # Function to group centers within a specified range
+    def group_centers(centers, axis, range_=50):
+        groups = []
+        while centers:
+            center, idx = centers.pop(0)
+            group = [(center, idx)]
+            for other_center, other_idx in centers[:]:
+                if abs(center[axis] - other_center[axis]) <= range_:
+                    group.append((other_center, other_idx))
+                    centers.remove((other_center, other_idx))
+            groups.append(group)
+        return groups
+
+    # Step 2: Align the elements within each pool
+    for pool_index, centers in pool_groups.items():
+        # Group bounding boxes by checking if their centers are within ±50 pixels on the y-axis
+        y_groups = group_centers(centers.copy(), axis=1)
+
+        # Align the y-coordinates of the centers of grouped bounding boxes
+        for group in y_groups:
+            avg_y = sum([c[0][1] for c in group]) / len(group)  # Calculate the average y-coordinate
+            for (center, idx) in group:
+                label = class_dict[modified_pred['labels'][idx]]
+                if label in size:
+                    new_center = (center[0], avg_y)  # Align the y-coordinate
+                    modified_pred['boxes'][idx] = [
+                        new_center[0] - size[label][0] / 2, 
+                        new_center[1] - size[label][1] / 2, 
+                        new_center[0] + size[label][0] / 2, 
+                        new_center[1] + size[label][1] / 2
+                    ]
+
+        # Recalculate centers after vertical alignment
+        centers = []
+        for group in y_groups:
+            for center, idx in group:
+                x1, y1, x2, y2 = modified_pred['boxes'][idx]
+                center = [(x1 + x2) / 2, (y1 + y2) / 2]
+                centers.append((center, idx))
+
+        # Group bounding boxes by checking if their centers are within ±50 pixels on the x-axis
+        x_groups = group_centers(centers.copy(), axis=0)
+
+        # Align the x-coordinates of the centers of grouped bounding boxes
+        for group in x_groups:
+            avg_x = sum([c[0][0] for c in group]) / len(group)  # Calculate the average x-coordinate
+            for (center, idx) in group:
+                label = class_dict[modified_pred['labels'][idx]]
+                if label in size:
+                    new_center = (avg_x, center[1])  # Align the x-coordinate
+                    modified_pred['boxes'][idx] = [
+                        new_center[0] - size[label][0] / 2, 
+                        modified_pred['boxes'][idx][1], 
+                        new_center[0] + size[label][0] / 2, 
+                        modified_pred['boxes'][idx][3]
+                    ]
+
     return modified_pred['boxes']
 
 
@@ -245,7 +307,7 @@ def display_options(image, score_threshold, is_mobile, screen_width):
     )
 
 # Function to perform inference on the uploaded image using the loaded models
-def perform_inference(model_object, model_arrow, image, score_threshold, is_mobile, screen_width):
+def perform_inference(model_object, model_arrow, image, score_threshold, is_mobile, screen_width, iou_threshold=0.5, distance_treshold=30, percentage_text_dist_thresh=0.5):
     _, uploaded_image = prepare_image(image, pad=False)
               
     img_tensor = F.to_tensor(prepare_image(image.convert('RGB'))[1])
@@ -260,14 +322,14 @@ def perform_inference(model_object, model_arrow, image, score_threshold, is_mobi
         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=0.5, distance_treshold=30)
+    _, 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)
-    st.session_state.text_mapping = mapping_text(st.session_state.prediction, st.session_state.text_pred, print_sentences=False, percentage_thresh=0.5)
+    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
     image_placeholder.empty()
@@ -419,7 +481,7 @@ def main():
             if st.button("Launch Prediction"):
                 st.session_state.crop_image = cropped_image
                 with st.spinner('Processing...'):
-                    perform_inference(model_object, model_arrow, st.session_state.crop_image, score_threshold, is_mobile, screen_width)
+                    perform_inference(model_object, 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()
 
     if 'prediction' in st.session_state and uploaded_file is not None: