Update doors_fasterrcnn.py

doors_fasterrcnn.py

@@ -105,14 +105,16 @@ def get_door_info(doors_info):
   lines = []
   sanda = []
   line_midpoint = []
-  doubleD_bbox = []
-  doubleD_width = []
-  
-  single_door = 0
+  #doubleD_bbox = []
+  #doubleD_width = []
+  singles = 0 
+  doubles = 0
+  door_type = []
   for door_inf in doors_info:
     xmin, ymin, xmax, ymax = door_inf[0]
-    #horz_bottom
-    if door_inf[1] == 2:
+    #horz_bottom #double_bottom
+    if door_inf[1] == 2 or door_inf[1] == 1:
+    #if door_inf[1] == 2:
       #for drawing
       point_st = (int(xmin), int(ymax) + 5)
       point_end = (int(xmax),int(ymax) + 5)
@@ -126,10 +128,16 @@ def get_door_info(doors_info):
       #for calculation
       width = distance((xmin,ymax), (xmax,ymax))
       width_pixels.append(width)
-      #single_door +=1
-
-    #horz_upper
-    if door_inf[1] == 3:
+      if door_inf[1] == 2:
+          singles +=1
+          door_type.append(0)
+      if door_inf[1] == 1:
+          doubles +=1
+          door_type.append(1)
+      
+
+    #horz_upper #double_upper
+    if door_inf[1] == 3 or door_inf[1] == 7:
       #for drawing
       point_st = (int(xmin),int(ymin) -5)
       point_end = (int(xmax),int(ymin) - 5)
@@ -142,9 +150,15 @@ def get_door_info(doors_info):
       #for calculation
       width = distance((xmin,ymin), (xmax,ymin))
       width_pixels.append(width)
-
-    #vert_right
-    if door_inf[1] == 4:
+      if door_inf[1] == 3:
+          singles +=1
+          door_type.append(0)
+      if door_inf[1] == 7:
+          doubles +=1
+          door_type.append(1)
+
+    #vert_right #double_right
+    if door_inf[1] == 5 or door_inf[1] == 8:
       #for drawing
       point_st = (int(xmax) + 5,int(ymin))
       point_end = (int(xmax) + 5,int(ymax))
@@ -157,8 +171,15 @@ def get_door_info(doors_info):
       #for calculation
       width = distance((xmax,ymin), (xmax,ymax))
       width_pixels.append(width)
-    #vert_left
-    if door_inf[1] == 5:
+      if door_inf[1] == 5:
+          singles +=1
+          door_type.append(0)
+      if door_inf[1] == 8:
+          doubles +=1
+          door_type.append(1)
+          
+    #vert_left #double_left
+    if door_inf[1] == 6 or door_inf[1] == 4:
       #for drawing
       point_st = (int(xmin) -5,int(ymin))
       point_end = (int(xmin) -5,int(ymax))
@@ -171,20 +192,27 @@ def get_door_info(doors_info):
       #for calculation
       width = distance((xmin,ymin), (xmin,ymax))
       width_pixels.append(width)
+      if door_inf[1] ==6:
+          singles +=1
+          door_type.append(0)
+      if door_inf[1] == 4:
+          doubles +=1
+          door_type.append(1)
+          
 
     #double door  
-    if door_inf[1] == 1:
-      doubleD_bbox.append([int(xmin), int(ymin), int(xmax), int(ymax)])
-      bbox_width = xmax - xmin
-      bbox_height = ymax - ymin
-      if bbox_width > bbox_height:
-        door_width = bbox_width
-        doubleD_width.append(door_width)
-      else:
-        door_width = bbox_height
-        doubleD_width.append(door_width)
-
-  return width_pixels, lines, sanda, line_midpoint, doubleD_bbox, doubleD_width
+    #if door_inf[1] == 1:
+     # doubleD_bbox.append([int(xmin), int(ymin), int(xmax), int(ymax)])
+      #bbox_width = xmax - xmin
+      #bbox_height = ymax - ymin
+      #if bbox_width > bbox_height:
+       # door_width = bbox_width
+        #doubleD_width.append(door_width)
+      #else:
+       # door_width = bbox_height
+        #doubleD_width.append(door_width)
+
+  return width_pixels, lines, sanda, line_midpoint, singles, doubles, door_type
 
 def pxl2meter(width_pixels, ratio):
   real_width = []
@@ -253,7 +281,7 @@ def adjustannotations(OutputPdfStage1):
   print("Annotations updated and saved ")
   return output_pdf_io.read()
 
-def add_annotations_to_pdf(image, pdf_name, lines, sanda, char_width, line_midpoint, doubleD_bbox, char_d_width):
+def add_annotations_to_pdf(image, pdf_name, lines, sanda, char_width, line_midpoint, door_type):
     image_width, image_height = image.size
 
     # Create a new PDF document
@@ -284,14 +312,14 @@ def add_annotations_to_pdf(image, pdf_name, lines, sanda, char_width, line_midpo
       annot = page.add_freetext_annot(rect, text, fontsize=10, fontname="helv", text_color=(1, 0, 0))
       annot.update()
 
-    for i in range(len(doubleD_bbox)):
+    #for i in range(len(doubleD_bbox)):
         #a7seb ana el midpoint beta3 el bbox 3ashan a7ot el width feha 
-        x,y = calculate_midpoint((doubleD_bbox[i][0], doubleD_bbox[i][1]), (doubleD_bbox[i][2],doubleD_bbox[i][3]))
-        p_midpoint = fitz.Point(x, y) * derotationMatrix
-        rect = fitz.Rect(p_midpoint.x, p_midpoint.y, p_midpoint.x + 200, p_midpoint.y + 50)
-        text = char_d_width[i]
-        annot = page.add_freetext_annot(rect, text, fontsize=10, fontname="helv", text_color=(1, 0, 0))
-        annot.update()
+        #x,y = calculate_midpoint((doubleD_bbox[i][0], doubleD_bbox[i][1]), (doubleD_bbox[i][2],doubleD_bbox[i][3]))
+        #p_midpoint = fitz.Point(x, y) * derotationMatrix
+        #rect = fitz.Rect(p_midpoint.x, p_midpoint.y, p_midpoint.x + 200, p_midpoint.y + 50)
+        #text = char_d_width[i]
+        #annot = page.add_freetext_annot(rect, text, fontsize=10, fontname="helv", text_color=(1, 0, 0))
+        #annot.update()
 
         
     #Annotation for drawin lines as in the markups
@@ -300,44 +328,59 @@ def add_annotations_to_pdf(image, pdf_name, lines, sanda, char_width, line_midpo
         annot = page.add_polyline_annot(l_points) 
         annot.set_border(width=2, dashes=None)  # Optional border styling
         annot.set_colors(stroke=(1, 0, 0))  # Set the line color to red
+        annot.set_info(content=str(char_width[i])+' m',subject='Perimeter Measurement', title="ADR Team")
         annot.update()
-    for i in range(len(doubleD_bbox)):
+
+    for i in range(len(door_type)):
+        x, y = line_midpoint[i]
+        p_midpoint = fitz.Point(x, y) * derotationMatrix
+        if door_type[i] == 0:
+            text = "Single Door"
+        else: 
+            text = "Double Door"
+        # Create an annotation (sticky note)
+        annot = page.add_text_annot((p_midpoint.x p_midpoint.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()
+        
+    #for i in range(len(doubleD_bbox)):
         # rect_coords should be a tuple: (x1, y1, x2, y2)
-        tpl1 = (doubleD_bbox[i][0], doubleD_bbox[i][1])
-        tpl2 = (doubleD_bbox[i][2], doubleD_bbox[i][3])
-        tpl_derotated = fitz.Point(*tpl1) * derotationMatrix
-        tpl2_derotated = fitz.Point(*tpl2) * derotationMatrix
-        print("Derotated tuples:")
-        print(tpl_derotated)
-        print("Derotated tuples 2:")
-        print(tpl2_derotated)
+        #tpl1 = (doubleD_bbox[i][0], doubleD_bbox[i][1])
+        #tpl2 = (doubleD_bbox[i][2], doubleD_bbox[i][3])
+        #tpl_derotated = fitz.Point(*tpl1) * derotationMatrix
+        #tpl2_derotated = fitz.Point(*tpl2) * derotationMatrix
+        #print("Derotated tuples:")
+        #print(tpl_derotated)
+        #print("Derotated tuples 2:")
+        #print(tpl2_derotated)
         #coordinates_derotated = [tpl_derotated[0],tpl_derotated[1],tpl2_derotated[0],tpl2_derotated[1]]
-        x1, y1 = tpl_derotated
-        x2, y2 = tpl2_derotated
-        x1 = int(x1)
-        y1 = int(y1)
-        x2 = int(x2)
-        y2 = int(y2)
-        x1, x2 = min(x1, x2), max(x1, x2)
-        y1, y2 = min(y1, y2), max(y1, y2)
+        #x1, y1 = tpl_derotated
+        #x2, y2 = tpl2_derotated
+        #x1 = int(x1)
+        #y1 = int(y1)
+        #x2 = int(x2)
+        #y2 = int(y2)
+        #x1, x2 = min(x1, x2), max(x1, x2)
+        #y1, y2 = min(y1, y2), max(y1, y2)
 
         
-        print("x1:")
-        print(x1)
-        print("y1:")
-        print(y1)
-        print("x2:")
-        print(x2)
-        print("y2:")
-        print(y2)
+        #print("x1:")
+        #print(x1)
+        #print("y1:")
+        #print(y1)
+        #print("x2:")
+        #print(x2)
+        #print("y2:")
+        #print(y2)
         
-        rect = fitz.Rect(x1, y1, x2, y2)
-        print("RECT:")
-        print(rect)
-        rect_annot = page.add_rect_annot(rect)
-        rect_annot.set_border(width=2)  # Adjust border width
-        rect_annot.set_colors(stroke=(1, 0, 0))  # Red rectangle
-        rect_annot.update()
+        #rect = fitz.Rect(x1, y1, x2, y2)
+        #print("RECT:")
+        #print(rect)
+        #rect_annot = page.add_rect_annot(rect)
+        #rect_annot.set_border(width=2)  # Adjust border width
+        #rect_annot.set_colors(stroke=(1, 0, 0))  # Red rectangle
+        #rect_annot.update()
 
         # annot = page.add_line_annot(lines[i][0], lines[i][1])
         # # annot = page.add_line_annot(sanda[i][0], lines[i][0])
@@ -396,17 +439,18 @@ def main_run(img_pillow,pdf_fullpath, weights_path, pdf_name,pdfpath,ratio): ###
     
     # START INFERENCE
     doors_info = ev_model(img_pillow, model, device, 0.9)
-    width_pixels, lines, sanda, line_midpoint, doubleD_bbox, doubleD_width = get_door_info(doors_info)
+    #width_pixels, lines, sanda, line_midpoint, doubleD_bbox, doubleD_width = get_door_info(doors_info)
+    width_pixels, lines, sanda, line_midpoint, single_count, double_count, door_type = get_door_info(doors_info)
     real_width = pxl2meter(width_pixels, ratio)
-    real_double = pxl2meter(doubleD_width, ratio)
+    #real_double = pxl2meter(doubleD_width, ratio)
     char_width = width_as_char(real_width)
-    char_d_width = width_as_char(real_double)
+    #char_d_width = width_as_char(real_double)
 
-    pdf_document = add_annotations_to_pdf(img_pillow, pdf_fullpath, lines, sanda, char_width, line_midpoint, doubleD_bbox, char_d_width)
+    pdf_document = add_annotations_to_pdf(img_pillow, pdf_fullpath, lines, sanda, char_width, line_midpoint, door_type)
     modified_pdf_data = pdf_document.tobytes()
     OutputPdfStage2=adjustannotations(modified_pdf_data)
     #Dataframe for Doors count
-    doors_count = {'Type': ['Single Doors', 'Double Doors'], 'Quantity': [len(real_width), len(real_double)]}
+    doors_count = {'Type': ['Single Doors', 'Double Doors'], 'Quantity': [single_count, double_count]}
     #doors_count = {'Type': ['Single Doors'], 'Quantity': [len(real_width)]}
     df_doors = pd.DataFrame(doors_count)