Update doors_fasterrcnn.py

doors_fasterrcnn.py

@@ -9,7 +9,6 @@ Original file is located at
 ## Libraries
 """
 
-# from google.colab.patches import cv2_imshow
 import cv2
 import numpy as np
 import pandas as pd
@@ -42,11 +41,9 @@ from PIL import Image, ImageDraw
 import torchvision.transforms.functional as F
 import matplotlib.pyplot as plt
 import google_sheet_Legend
-"""# updated for (fullpath, pdf_name)
-
-
-
-"""
+import torch
+import torchvision
+from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
 
 def convert2pillow(path):
   pdf = pdfium.PdfDocument(path)
@@ -54,10 +51,6 @@ def convert2pillow(path):
   pil_image = page.render().to_pil()
   return pil_image
 
-import torch
-import torchvision
-from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
-
 # Function to get the model
 def get_model(num_classes):
     # Load a pre-trained Faster R-CNN model with a ResNet-50-FPN backbone
@@ -71,25 +64,6 @@ def get_model(num_classes):
 
     return model
 
-'''def ev_model(img, model, device, threshold):
-  image_tensor = F.to_tensor(img).unsqueeze(0)
-  image_tensor = image_tensor.to(device)
-  model.eval()
-
-  with torch.no_grad():
-    predictions = model(image_tensor)
-
-  single_boxes = []
-  double_boxes = []
-  for element in range(len(predictions[0]['boxes'])):
-    score = predictions[0]['scores'][element].item()
-    if score > threshold:
-      if predictions[0]['labels'][element].item() == 1:
-        single_boxes.append(predictions[0]['boxes'][element].tolist())
-      else:
-        double_boxes.append(predictions[0]['boxes'][element].tolist())
-
-  return single_boxes, double_boxes'''
 
 def ev_model(img, model, device, threshold):
   image_tensor = F.to_tensor(img).unsqueeze(0)
@@ -241,142 +215,10 @@ def add_annotations_to_pdf(image, pdf_name, lines, sanda, char_width, line_midpo
     
 
 
-'''def calculate_width(bbox):
-  #if looking right or left, width < height
-  bbox_width = bbox[2] - bbox[0]
-  bbox_height = bbox[3] - bbox[1]
-  if bbox_width > bbox_height:
-    door_width = bbox_width
-  else:
-    door_width = bbox_height
-  return door_width
-'''
-'''def width_annotations(bbox, ratio):
-  lines = []
-  width = []
-
-  for box in bbox:
-    door_width = calculate_width(box)
-    door_width = round(door_width*ratio, 2)
-    x1,y1,x2,y2 = box
-    b_width = x2 - x1
-    b_height = y2 - y1
-
-    if b_width > b_height:
-      lines.append(((x1, y1), (x2, y1)))
-      x = (x1+x2)/2
-      y = (y1+y1)/2
-      width.append(((x,y),door_width))
-    else:
-      lines.append(((x1, y1), (x1, y2)))
-      x = (x1+x1)/2
-      y = (y1+y2)/2
-      width.append(((x,y), door_width))
-  return lines, width'''
-
-'''def create_width_annotations(width_info):
-  annotations = []
-  for i in range(len(width_info)):
-    annotations.append(((width_info[i][0][0]),(width_info[i][0][1]),f"{width_info[i][1]} m"))
-  return annotations'''
-
-'''def calculate_midpoint(top_left, bottom_right):
-    x1, y1 = top_left
-    x2, y2 = bottom_right
-    # Calculate the midpoint
-    xm = int((x1 + x2) / 2)
-    ym = int((y1 + y2) / 2)
-    return (xm, ym)'''
-
-'''def mid_points_bbox(bbox):
-  midpoint = []
-  for i in range(len(bbox)):
-    x1 = int(bbox[i][0])
-    y1 = int(bbox[i][1])
-    x2 = int(bbox[i][2])
-    y2 = int(bbox[i][3])
-    top_left_corner = (x1, y1)
-    bottom_right_corner = (x2, y2)
-    midpoint.append(calculate_midpoint(top_left_corner, bottom_right_corner))
-  return midpoint'''
-
-'''def create_annotations(door_kind, midpoints):
-  door = door_kind
-  annotations = []
-  for i in range(len(midpoints)):
-    annotations.append((midpoints[i][0],midpoints[i][1], door))
-  return annotations
-
-def add_annotations_to_pdf(image, pdf_name, annotation_s, annotation_d,width_ann_single, width_ann_double,line_single,line_double):
-    image_width, image_height = image.size
-
-    # Create a new PDF document
-    pdf_document = fitz.open()
-
-    # Add a new page to the document with the same dimensions as the image
-    page = pdf_document.new_page(width=image_width, height=image_height)
-
-    # Insert the image into the PDF page
-    image_stream = io.BytesIO()
-    image.save(image_stream, format="PNG")
-    page.insert_image(page.rect, stream=image_stream.getvalue())
-
-    # Add annotations
-    for annotation in annotation_s:
-        x, y, text = annotation
-        # Create an annotation (sticky note)
-        annot = page.add_text_annot(fitz.Point(x, 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 annotation in annotation_d:
-        x, y, text = annotation
-        # Create an annotation (sticky note)
-        annot = page.add_text_annot(fitz.Point(x, y), text)
-        annot.set_border(width=0.2, dashes=(1, 2))  # Optional border styling
-        annot.set_colors(stroke=(0, 1, 0), fill=None)  # Set the stroke color to red
-        annot.update()
-
-    #Annotations for width measurement (marra single we marra double)
-    for annotation in width_ann_single:
-        x, y, text = annotation
-        rect = fitz.Rect(x, y, x + 200, y + 50)  # Adjust the width and height as needed
-        annot = page.add_freetext_annot(rect, text, fontsize=10, fontname="helv", text_color=(1,0,0))
-        annot.update()
-    for annotation in width_ann_double:
-        x, y, text = annotation
-        rect = fitz.Rect(x, y, x + 200, y + 50)  # Adjust the width and height as needed
-        annot = page.add_freetext_annot(rect, text, fontsize=10, fontname="helv", text_color=(1,0,0))
-        annot.update()
-
-    #Annotation kind of the line drawings (marra single we marra double)
-    for line in line_single:
-        start_point, end_point = line
-        annot = page.add_line_annot(start_point, end_point)
-        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 line in line_double:
-        start_point, end_point = line
-        annot = page.add_line_annot(start_point, end_point)
-        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()
-
-
-    output_pdf_path = pdf_name+"_annotated.pdf"
-    # Save the PDF with annotations
-
-    return pdf_document
-    # pdf_document.save(output_pdf_path)
-    # pdf_document.close()
-'''
 def main_run(pdf_fullpath, weights_path, pdf_name,pdfpath,ratio): ####pdf_fullpath here is the data and not the path 
     img_pillow = convert2pillow(pdf_fullpath)
-    #new_image = img_pillow.resize((2384, 1684))
     
-    # Specify the number of classes (including the background)
-    num_classes = 10  # Ensure this matches the saved model's number of classes
+    num_classes = 10  # classes + background
     
     # Load the model with the specified number of classes
     model = get_model(num_classes)
@@ -397,29 +239,6 @@ def main_run(pdf_fullpath, weights_path, pdf_name,pdfpath,ratio): ####pdf_fullpa
     model.to(device)
     
     # START INFERENCE
-    #sbox, dbox = ev_model(img_pillow, model, device, 0.6)
-    
-    #Dataframe for Doors count
-    #doors_count = {'Type': ['Single Doors', 'Double Doors'], 'Quantity': [len(sbox), len(dbox)]}
-    #df_doors = pd.DataFrame(doors_count)
-
-    #single_midpoint = mid_points_bbox(sbox)
-    #double_midpoint = mid_points_bbox(dbox)
-
-    #Kind Annotations
-    #single_annotations = create_annotations("single door", single_midpoint)
-    #double_annotations = create_annotations("double door", double_midpoint)
-    #Lines Annotations
-    #line_single, width_signle = width_annotations(sbox, 1)
-    #line_double, width_double = width_annotations(dbox, 1)
-    #Width Annotations
-    #width_single_ann = create_width_annotations(width_signle)
-    #width_double_ann = create_width_annotations(width_double)
-
-    # add_annotations_to_pdf(new_image, pdf_name, single_annotations, double_annotations)
-
-
-    #NEW
     doors_info = ev_model(img_pillow, model, device, 0.9)
     width_pixels, lines, sanda, line_midpoint = get_door_info(doors_info)
     real_width = pxl2meter(width_pixels, ratio)
@@ -427,7 +246,6 @@ def main_run(pdf_fullpath, weights_path, pdf_name,pdfpath,ratio): ####pdf_fullpa
 
     pdf_document = add_annotations_to_pdf(img_pillow, plan, lines, sanda, char_width, line_midpoint)
     
-    #pdf_document=add_annotations_to_pdf(img_pillow, pdf_name, single_annotations, double_annotations,width_single_ann,width_double_ann,line_single,line_double)
 
     page=pdf_document[0]
     pix = page.get_pixmap()  # render page to an image
@@ -444,7 +262,7 @@ def main_run(pdf_fullpath, weights_path, pdf_name,pdfpath,ratio): ####pdf_fullpa
 
 
     # modify this return
-    return annotatedimg, pdf_document , spreadsheet_url, list1 , df_doors
+    return annotatedimg, pdf_document , spreadsheet_url, list1 
 
 # model_path =  '/content/drive/MyDrive/combined.pth'
 # #pdf_name = data