Update 2.1_Counting_Columns.py

2.1_Counting_Columns.py

@@ -120,12 +120,14 @@ def getNearestText(point_list, p):
 
 def getColumnsTypes(nearbyy, x):
   found_tuple = []
+  selected_clms_points = []
   # Loop through the list of tuples
   for i in range(len(nearbyy)):
     for tpl in x:
       if (tpl[2] == nearbyy[i][0] and tpl[3] == nearbyy[i][1]) and tpl[4].startswith("C"):
         found_tuple.append(tpl[4])
-  return found_tuple
+        selected_clms_points.append(nearbyy[i])
+  return found_tuple, selected_clms_points
 
 def generate_legend(found_tuple):
   word_freq = {}
@@ -138,6 +140,34 @@ def generate_legend(found_tuple):
   df = pd.DataFrame(data.items(), columns=['Column Type', 'Count'])
   return df
 
+def add_annotations_to_pdf(image, pdf_name, columns_types, slctd_clms_pts):
+    image_width, image_height = image.size
+
+    # Create a new PDF document
+    pdf_document = fitz.open('pdf',pdf_name)
+    page=pdf_document[0]
+    rotationOld=page.rotation
+    derotationMatrix=page.derotation_matrix
+    print('rotationOld',rotationOld)
+    if page.rotation!=0:
+      rotationangle = page.rotation
+      page.set_rotation(0)
+    print('rotationnew',page.rotation)
+
+
+    #Annotation for drawin lines as in the markups
+     for i in range(len(columns_types)):
+        x, y = slctd_clms_pts[i]
+        text = columns_types[i]
+        # Create an annotation (sticky note)
+        annot = page.add_text_annot(fitz.Point(x, y), text)
+        annot.set_border(width=0.2, dashes=(1, 2))  # Optional border styling
+        annot.set_colors(stroke=(1, 0, 0), fill=None)  # Set the stroke color to red
+        annot.update()
+
+    page.set_rotation(rotationOld)
+    return pdf_document
+
 def mainfun(plan):
   texts_from_pdf = get_text_from_pdf(plan)
   img = convert2img(plan)
@@ -148,8 +178,9 @@ def mainfun(plan):
     # BROWN COLUMNS
     text_points = getTextsPoints(texts_from_pdf)
     nearby = getNearestText(text_points, column_points)
-    columns_types = getColumnsTypes(nearby, texts_from_pdf)
+    columns_types, slctd_clms_pts = getColumnsTypes(nearby, texts_from_pdf)
     legend = generate_legend(columns_types)
+    add_annotations_to_pdf(image, pdf_name, columns_types, slctd_clms_pts)
   else:
     # BLUE COLUMNS
     img_blue = changeGrayModify(img)
@@ -158,6 +189,7 @@ def mainfun(plan):
     column_points,mask_clmns, mask_walls = get_columns_info(outsu, img)
     text_points = getTextsPoints(texts_from_pdf)
     nearby = getNearestText(text_points, column_points)
-    columns_types = getColumnsTypes(nearby, texts_from_pdf)
+    columns_types, slctd_clms_pts = getColumnsTypes(nearby, texts_from_pdf)
     legend = generate_legend(columns_types)
+    add_annotations_to_pdf(image, pdf_name, columns_types, slctd_clms_pts)
   return legend