Update doors_fasterrcnn.py

doors_fasterrcnn.py

@@ -100,7 +100,7 @@ def calculate_midpoint(p1, p2):
     xm = int((x1 + x2) / 2)
     ym = int((y1 + y2) / 2)
     return (xm, ym)
-def get_door_info(doors_info):
+def get_door_info(doors_info, doors_info_double):
   width_pixels = []
   lines = []
   sanda = []
@@ -214,6 +214,81 @@ def get_door_info(doors_info):
 
   return width_pixels, lines, sanda, line_midpoint, singles, doubles, door_type
 
+def get_door_info_double(doors_info_double, width_pixels, lines, sanda, line_midpoint, doubles, door_type):
+  for door_inf in doors_info_double:
+    xmin, ymin, xmax, ymax = door_inf[0]
+      
+    #double_bottom
+    if door_inf[1] == 1:
+      point_st = (int(xmin), int(ymax) + 5)
+      point_end = (int(xmax),int(ymax) + 5)
+      lines.append((point_st,point_end))
+
+      sanda_st = (int(xmin), int(ymax))
+      sand_end = (int(xmax),int(ymax))
+      sanda.append((sanda_st, sand_end))
+
+      line_midpoint.append(calculate_midpoint(sanda_st,sand_end))
+      #for calculation
+      width = distance((xmin,ymax), (xmax,ymax))
+      width_pixels.append(width)
+      doubles +=1
+      door_type.append(1)
+      
+
+    #double_upper
+    if door_inf[1] == 7:
+      #for drawing
+      point_st = (int(xmin),int(ymin) -5)
+      point_end = (int(xmax),int(ymin) - 5)
+      lines.append((point_st,point_end))
+
+      sanda_st = (int(xmin),int(ymin))
+      sand_end = (int(xmax),int(ymin))
+      sanda.append((sanda_st, sand_end))
+      line_midpoint.append(calculate_midpoint(sanda_st,sand_end))
+      #for calculation
+      width = distance((xmin,ymin), (xmax,ymin))
+      width_pixels.append(width)
+      doubles +=1
+      door_type.append(1)
+
+    #double_right
+    if door_inf[1] == 8:
+      #for drawing
+      point_st = (int(xmax) + 5,int(ymin))
+      point_end = (int(xmax) + 5,int(ymax))
+      lines.append((point_st,point_end))
+
+      sanda_st = (int(xmax), int(ymin))
+      sand_end = (int(xmax), int(ymax))
+      sanda.append((sanda_st, sand_end))
+      line_midpoint.append(calculate_midpoint(sanda_st,sand_end))
+      #for calculation
+      width = distance((xmax,ymin), (xmax,ymax))
+      width_pixels.append(width)
+      doubles +=1
+      door_type.append(1)
+          
+    #double_left
+    if door_inf[1] == 4:
+      #for drawing
+      point_st = (int(xmin) -5,int(ymin))
+      point_end = (int(xmin) -5,int(ymax))
+      lines.append((point_st,point_end))
+
+      sanda_st = (int(xmin),int(ymin))
+      sand_end = (int(xmin),int(ymax))
+      sanda.append((sanda_st, sand_end))
+      line_midpoint.append(calculate_midpoint(sanda_st,sand_end))
+      #for calculation
+      width = distance((xmin,ymin), (xmin,ymax))
+      width_pixels.append(width)
+      doubles +=1
+      door_type.append(1)
+
+  return width_pixels, lines, sanda, line_midpoint, singles, doubles, door_type
+
 def pxl2meter(width_pixels, ratio):
   real_width = []
   for width in width_pixels:
@@ -350,31 +425,47 @@ def add_annotations_to_pdf(image, pdf_name, lines, sanda, char_width, line_midpo
 
 def main_run(img_pillow,pdf_fullpath, weights_path, pdf_name,pdfpath,ratio): ####pdf_fullpath here is the data and not the path 
     img_pillow = convert2pillow(pdf_fullpath)
-    
+
+    # For Single Doors
     num_classes = 12  # classes + background
-    
     # Load the model with the specified number of classes
     model = get_model(num_classes)
-    
     # Load the saved model's state dictionary with map_location to handle CPU
     device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
-    
     try:
         model.load_state_dict(torch.load(weights_path, map_location=device), strict=False)
     except RuntimeError as e:
         print(f"Error loading model state_dict: {e}")
         return
-    
     # Set the model to evaluation mode
     model.eval()
-    
     # Move the model to the appropriate device
     model.to(device)
     
+
+    # For Double Doors
+    num_classes2 = 11  # classes + background
+    # Load the model with the specified number of classes
+    model2 = get_model(num_classes2)
+    # Load the saved model's state dictionary with map_location to handle CPU
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+    try:
+        model.load_state_dict(torch.load(weights_path2, map_location=device), strict=False)
+    except RuntimeError as e:
+        print(f"Error loading model state_dict: {e}")
+        return
+    # Set the model to evaluation mode
+    model2.eval()
+    # Move the model to the appropriate device
+    model2.to(device)
+    
     # START INFERENCE
     doors_info = ev_model(img_pillow, model, device, 0.8)
+    doors_info_double = ev_model(img_pillow, model2, device, 0.8)
+
+    
     #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)
+    width_pixels, lines, sanda, line_midpoint, single_count, double_count, door_type = get_door_info(doors_info, doors_info_double)
     real_width = pxl2meter(width_pixels, ratio)
     #real_double = pxl2meter(doubleD_width, ratio)
     char_width = width_as_char(real_width)