Update doors_fasterrcnn.py

doors_fasterrcnn.py

@@ -105,6 +105,9 @@ def get_door_info(doors_info):
   lines = []
   sanda = []
   line_midpoint = []
+  doubleD_bbox = []
+  doubleD_width = []
+  
   single_door = 0
   for door_inf in doors_info:
     xmin, ymin, xmax, ymax = door_inf[0]
@@ -169,7 +172,19 @@ def get_door_info(doors_info):
       width = distance((xmin,ymin), (xmin,ymax))
       width_pixels.append(width)
 
-  return width_pixels, lines, sanda, line_midpoint
+    #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
 
 def pxl2meter(width_pixels, ratio):
   real_width = []
@@ -238,7 +253,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):
+def add_annotations_to_pdf(image, pdf_name, lines, sanda, char_width, line_midpoint, doubleD_bbox):
     image_width, image_height = image.size
 
     # Create a new PDF document
@@ -276,6 +291,14 @@ def add_annotations_to_pdf(image, pdf_name, lines, sanda, char_width, line_midpo
         annot.set_border(width=2, dashes=None)  # Optional border styling
         annot.set_colors(stroke=(1, 0, 0))  # Set the line color to red
         annot.update()
+    for i in range(len(doubleD_bbox)):
+        # rect_coords should be a tuple: (x1, y1, x2, y2)
+        rect = fitz.Rect(fitz.Point(*doubleD_bbox[i])* derotationMatrix)
+        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])
         # # annot = page.add_line_annot(sanda[i][1], lines[i][1])
@@ -333,7 +356,7 @@ 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 = get_door_info(doors_info)
+    width_pixels, lines, sanda, line_midpoint, doubleD_bbox, doubleD_width = get_door_info(doors_info)
     real_width = pxl2meter(width_pixels, ratio)
     char_width = width_as_char(real_width)