Update Counting_Columns_2_1.py

Counting_Columns_2_1.py

@@ -96,10 +96,37 @@ def get_columns_info(outsu4, img):
 
 def getTextsPoints(x):
   point_list = []
+  pt_clm = {}
   for h in x:
-    point_list.append((h[2],h[3]))
-  return point_list
+    point_list.append(calculate_midpoint(h[1],h[0],h[3],h[2]))
+    pt_clm[calculate_midpoint(h[1],h[0],h[3],h[2])] = h[4]
+  return point_list, pt_clm
+
+def fix_90_ky_val(pt_clm):
+  new_derotated = {}
+  for ky in pt_clm:
+    pts = fitz.Point(ky[0], ky[1]) * derotationMatrix
+    new_ky = ((int(pts.y),int(pts.x)))
+    new_derotated[new_ky] = pt_clm[ky]
+  return new_derotated
+
+def calculate_midpoint(x1,y1,x2,y2):
+  xm = int((x1 + x2) / 2)
+  ym = int((y1 + y2) / 2)
+  return (xm, ym)
+
+def getColumnsTypesKeyValue(nearbyy, pt_clm):
+  words = []
+  for i in range(len(nearbyy)):
+    words.append(pt_clm[nearbyy[i]])
+  return words
 
+def fix_rotation_90(pc_coordinates, derotationMatrix):
+  coor = []
+  for coordinate in pc_coordinates:
+    pts = fitz.Point(coordinate[0], coordinate[1]) * derotationMatrix
+    coor.append((int(pts.y),int(pts.x)))
+  return coor
 
 def distance(point1, point2):
     x1, y1 = point1
@@ -139,6 +166,40 @@ def generate_legend(found_tuple):
   return df
 
 def mainfun(plan):
+  pdf_document = fitz.open(plan)
+  page = pdf_document[0]
+  rotation = page.rotation
+  derotationMatrix=page.derotation_matrix
+  texts_from_pdf = get_text_from_pdf(plan)
+  text_points, txtpts_ky_vlu = getTextsPoints(texts_from_pdf)
+  if rotation != 0:
+    if rotation ==90:
+      text_points = fix_rotation_90(text_points, derotationMatrix)
+      txtpts_ky_vlu = fix_90_ky_val(txtpts_ky_vlu)
+
+  img = convert2img(plan)
+  imgResult = segment_brown(img)
+  outsu = threshold(imgResult)
+  column_points,mask_clmns, mask_walls = get_columns_info(outsu, img)
+
+  if len(column_points) > 10:
+      # BROWN COLUMNS
+      nearby = getNearestText(text_points, column_points)
+      columns_types_v = getColumnsTypesKeyValue(nearby, txtpts_ky_vlu)
+      legend = generate_legend(columns_types_v)
+
+  else:
+      # BLUE COLUMNS
+      img_blue = changeGrayModify(img)
+      imgResult = segment_blue(img_blue)
+      outsu = threshold(imgResult)
+      column_points,mask_clmns, mask_walls = get_columns_info(outsu, img)
+      nearby = getNearestText(text_points, column_points)
+      columns_types_v = getColumnsTypesKeyValue(nearby, txtpts_ky_vlu)
+      legend = generate_legend(columns_types_v)
+  return legend
+
+'''def mainfun(plan):
   texts_from_pdf = get_text_from_pdf(plan)
   img = convert2img(plan)
   imgResult = segment_brown(img)
@@ -166,4 +227,4 @@ def mainfun(plan):
         nearby = fix_rotation_90(pc_coordinates)
     columns_types = getColumnsTypes(nearby, texts_from_pdf)
     legend = generate_legend(columns_types)
-  return legend
+  return legend'''