Update doors_fasterrcnn.py

doors_fasterrcnn.py

@@ -167,39 +167,49 @@ def add_annotations_to_pdf(image, pdf_name, annotation_s, annotation_d):
     # pdf_document.close()
 
 def main_run(pdf_fullpath, weights_path, pdf_name):
-  img_pillow = convert2pillow(pdf_fullpath)
-  new_image = img_pillow.resize((2384,1684))
-  # Specify the number of classes (including the background)
-  num_classes = 6  # Ensure this matches the saved model's number of classes
-  # 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')
-  # model.load_state_dict(torch.load(weights_path, map_location=device))
-  model.load_state_dict(torch.load(weights_path, map_location=device), strict=False)
-  # Set the model to evaluation mode
-  model.eval()
-  # Move the model to the appropriate device
-  model.to(device)
-
-  #START INFERENCE
-  sbox, dbox = ev_model(new_image, model, device, 0.6)
-
-  single_info = mid_points_bbox(sbox)
-  double_info = mid_points_bbox(dbox)
-
-  single_annotations = create_annotations("single door", single_info)
-  double_annotations = create_annotations("double door", double_info)
-
-  pdf_document=add_annotations_to_pdf(new_image, pdf_name, single_annotations, double_annotations)
-  
-  page=pdf_document[0]
-  pix = page.get_pixmap()  # render page to an image
-  pl=Image.frombytes('RGB', [pix.width,pix.height],pix.samples)
-  img=np.array(pl)
-  annotatedimg = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-
-  return annotatedimg,pdf_document
+    img_pillow = convert2pillow(pdf_fullpath)
+    new_image = img_pillow.resize((2384, 1684))
+    
+    # Specify the number of classes (including the background)
+    num_classes = 6  # Ensure this matches the saved model's number of classes
+    
+    # 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)
+    
+    # START INFERENCE
+    sbox, dbox = ev_model(new_image, model, device, 0.6)
+    
+    single_info = mid_points_bbox(sbox)
+    double_info = mid_points_bbox(dbox)
+    
+    single_annotations = create_annotations("single door", single_info)
+    double_annotations = create_annotations("double door", double_info)
+    
+        # add_annotations_to_pdf(new_image, pdf_name, single_annotations, double_annotations)
+    pdf_document=add_annotations_to_pdf(new_image, pdf_name, single_annotations, double_annotations)
+    
+    page=pdf_document[0]
+    pix = page.get_pixmap()  # render page to an image
+    pl=Image.frombytes('RGB', [pix.width,pix.height],pix.samples)
+    img=np.array(pl)
+    annotatedimg = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+    
+    return annotatedimg,pdf_document
 
 # model_path =  '/content/drive/MyDrive/combined.pth'
 # #pdf_name = data