Update Counting_Columns_2_1.py

Counting_Columns_2_1.py

@@ -9,6 +9,7 @@ import google_sheet_Legend
 import pypdfium2 as pdfium
 import fitz  # PyMuPDF
 import os
+import random
 
 def get_text_from_pdf(input_pdf_path):
     pdf_document = fitz.open('pdf',input_pdf_path)
@@ -167,7 +168,21 @@ def generate_legend(found_tuple):
   data = word_freq
   df = pd.DataFrame(data.items(), columns=['Column Type', 'Count'])
   return df
-
+    
+def color_groups(txtpts_ky_vlu):
+  unique_labels = list(set(txtpts_ky_vlu.values()))
+  def generate_rgb():
+        return (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))  # RGB tuple
+  key_colors = {key: generate_rgb() for key in unique_labels}  # Assign a unique RGB color to each key
+  return key_colors
+
+def get_drawing_info(txt_clmn,txtpts_ky_vlu,key_colors):
+  #Search for each word in the txt_clmn to get the word associated to it
+  huge_list_clmn_clr_loc = []
+  for text_location, column_location in txt_clmn:
+    word = txtpts_ky_vlu[text_location]
+    huge_list_clmn_clr_loc.append((text_location, column_location, word, key_colors[word]))
+  return huge_list_clmn_clr_loc #text_location, column_location, word, color
 '''def add_annotations_to_pdf(image, pdf_name, slctd_clm, columns_types_v):
     image_width = image.shape[1]
     image_height = image.shape[0]
@@ -191,7 +206,7 @@ def generate_legend(found_tuple):
     page.set_rotation(rotationOld)
     return pdf_document'''
     
-def add_annotations_to_pdf(image, pdf_name, slctd_clm, columns_types_v):
+def add_annotations_to_pdf(image, pdf_name, huge_list_clmn_clr_loc):
     image_width = image.shape[1]
     image_height = image.shape[0]
     # Create a new PDF document
@@ -202,21 +217,24 @@ def add_annotations_to_pdf(image, pdf_name, slctd_clm, columns_types_v):
     if page.rotation!=0:
         rotationangle = page.rotation
         page.set_rotation(0)
-    for i in range(len(slctd_clm)):
-        x, y = slctd_clm[i]
+    #for i in range(len(slctd_clm)):
+    for text_loc, column_loc, word, color in huge_list_clmn_clr_loc:
+        x, y = column_loc
+        clr = (clr[0] / 255, clr[1] / 255, clr[2] / 255)
+        #x, y = slctd_clm[i]
         p_midpoint = fitz.Point(x, y) * derotationMatrix
         annot = page.add_circle_annot(
                 fitz.Rect(p_midpoint.x - 10, p_midpoint.y - 10, p_midpoint.x + 10,p_midpoint.y + 10)  # Small circle
             )
         # ✅ Assign required Bluebeam metadata
-        annot.set_colors(stroke=(1, 0, 0), fill=(1, 1, 1))  # Set stroke color and fill white
+        annot.set_colors(stroke=clr, fill=(1, 1, 1))  # Set stroke color and fill white
         annot.set_border(width=2)  # Border thickness
         annot.set_opacity(1)  # Fully visible
         text = columns_types_v[i]
         # ✅ Set annotation properties for Bluebeam Count detection
-        annot.set_info("name", text)  # Unique name for each count
+        annot.set_info("name", word)  # Unique name for each count
         annot.set_info("subject", "Count")  # ✅ Bluebeam uses "Count" for Count markups
-        annot.set_info("title", text)  # Optional
+        annot.set_info("title", word)  # Optional
         annot.update()  # Apply changes
     page.set_rotation(rotationOld)
     return pdf_document
@@ -237,12 +255,14 @@ def mainfun(pdf_name,pdfpath,planname):
   imgResult = segment_brown(img)
   outsu = threshold(imgResult)
   column_points,mask_clmns, mask_walls = get_columns_info(outsu, img)
-
+  key_colors = color_groups(txtpts_ky_vlu)
+    
   if len(column_points) > 10:
       # BROWN COLUMNS
       nearby, slctd_clm = getNearestText(text_points, column_points)
       columns_types_v = getColumnsTypesKeyValue(nearby, txtpts_ky_vlu)
       legend = generate_legend(columns_types_v)
+      huge_list_clmn_clr_loc = get_drawing_info(txt_clmn,txtpts_ky_vlu,key_colors)
 
   else:
       # BLUE COLUMNS
@@ -253,7 +273,9 @@ def mainfun(pdf_name,pdfpath,planname):
       nearby, slctd_clm = getNearestText(text_points, column_points)
       columns_types_v = getColumnsTypesKeyValue(nearby, txtpts_ky_vlu)
       legend = generate_legend(columns_types_v)
-  pdf_document = add_annotations_to_pdf(img, pdf_name, slctd_clm, columns_types_v)
+      huge_list_clmn_clr_loc = get_drawing_info(txt_clmn,txtpts_ky_vlu,key_colors)
+      
+  pdf_document = add_annotations_to_pdf(img, pdf_name, huge_list_clmn_clr_loc)
   page=pdf_document[0]
   pix = page.get_pixmap()  # render page to an image
   pl=Image.frombytes('RGB', [pix.width,pix.height],pix.samples)