Update doors_fasterrcnn.py

doors_fasterrcnn.py

@@ -71,7 +71,7 @@ def get_model(num_classes):
 
     return model
 
-def ev_model(img, model, device, threshold):
+'''def ev_model(img, model, device, threshold):
   image_tensor = F.to_tensor(img).unsqueeze(0)
   image_tensor = image_tensor.to(device)
   model.eval()
@@ -89,9 +89,159 @@ def ev_model(img, model, device, threshold):
       else:
         double_boxes.append(predictions[0]['boxes'][element].tolist())
 
-  return single_boxes, double_boxes
+  return single_boxes, double_boxes'''
+
+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)
+
+  doors_info = []
+  for element in range(len(predictions[0]['boxes'])):
+    score = predictions[0]['scores'][element].item()
+    if score > threshold:
+        box = predictions[0]['boxes'][element].tolist()
+        label = predictions[0]['labels'][element].item()
+        doors_info.append((box,label))
+  return doors_info
+
+
+def distance(point1, point2):
+    x1, y1 = point1
+    x2, y2 = point2
+    return np.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2)
+
+def calculate_midpoint(p1, p2):
+    x1, y1 = p1
+    x2, y2 = p2
+    # Calculate the midpoint
+    xm = int((x1 + x2) / 2)
+    ym = int((y1 + y2) / 2)
+    return (xm, ym)
+def get_door_info(doors_info):
+  width_pixels = []
+  lines = []
+  sanda = []
+  line_midpoint = []
+  for door_inf in doors_info:
+    xmin, ymin, xmax, ymax = door_inf[0]
+    #horz_bottom
+    if door_inf[1] == 2:
+      #for drawing
+      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)
+
+    #horz_upper
+    if door_inf[1] == 3:
+      #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)
+
+    #vert_right
+    if door_inf[1] == 4:
+      #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)
+    #vert_left
+    if door_inf[1] == 5:
+      #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)
+
+  return width_pixels, lines, sanda, line_midpoint
+
+def pxl2meter(width_pixels, ratio):
+  real_width = []
+  for width in width_pixels:
+    real_width.append(round(width*ratio, 2))
+  return real_width
+
+def width_as_char(real_width):
+  char_width = []
+  for width in real_width:
+    char_width.append(f"{width}m")
+  return char_width
+
+def add_annotations_to_pdf(image, pdf_name, lines, sanda, char_width, line_midpoint):
+    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)
 
-def calculate_width(bbox):
+    # 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())
+
+
+    #Annotation for drawin lines as in the markups
+    for i in range(len(line_midpoint)):
+      x, y = line_midpoint[i]
+      text = char_width[i]
+      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 for drawin lines as in the markups
+    for i in range(len(lines)):
+        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])
+        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
+    
+
+
+'''def calculate_width(bbox):
   #if looking right or left, width < height
   bbox_width = bbox[2] - bbox[0]
   bbox_height = bbox[3] - bbox[1]
@@ -100,8 +250,8 @@ def calculate_width(bbox):
   else:
     door_width = bbox_height
   return door_width
-
-def width_annotations(bbox, ratio):
+'''
+'''def width_annotations(bbox, ratio):
   lines = []
   width = []
 
@@ -122,23 +272,23 @@ def width_annotations(bbox, ratio):
       x = (x1+x1)/2
       y = (y1+y2)/2
       width.append(((x,y), door_width))
-  return lines, width
+  return lines, width'''
 
-def create_width_annotations(width_info):
+'''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
+  return annotations'''
 
-def calculate_midpoint(top_left, bottom_right):
+'''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)
+    return (xm, ym)'''
 
-def mid_points_bbox(bbox):
+'''def mid_points_bbox(bbox):
   midpoint = []
   for i in range(len(bbox)):
     x1 = int(bbox[i][0])
@@ -148,9 +298,9 @@ def mid_points_bbox(bbox):
     top_left_corner = (x1, y1)
     bottom_right_corner = (x2, y2)
     midpoint.append(calculate_midpoint(top_left_corner, bottom_right_corner))
-  return midpoint
+  return midpoint'''
 
-def create_annotations(door_kind, midpoints):
+'''def create_annotations(door_kind, midpoints):
   door = door_kind
   annotations = []
   for i in range(len(midpoints)):
@@ -220,13 +370,13 @@ def add_annotations_to_pdf(image, pdf_name, annotation_s, annotation_d,width_ann
     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 = 3  # Ensure this matches the saved model's number of classes
+    num_classes = 10  # Ensure this matches the saved model's number of classes
     
     # Load the model with the specified number of classes
     model = get_model(num_classes)
@@ -247,27 +397,37 @@ 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)
+    #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)
+    #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)
+    #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)
+    #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)
+    #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)
+    #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)
-    pdf_document=add_annotations_to_pdf(img_pillow, pdf_name, single_annotations, double_annotations,width_single_ann,width_double_ann,line_single,line_double)
+
+
+    #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)
+    char_width = width_as_char(real_width)
+
+    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
@@ -282,6 +442,8 @@ def main_run(pdf_fullpath, weights_path, pdf_name,pdfpath,ratio): ####pdf_fullpa
         for annot in page.annots():
             list1.loc[len(list1)] =annot.info
 
+
+    # modify this return
     return annotatedimg, pdf_document , spreadsheet_url, list1 , df_doors
 
 # model_path =  '/content/drive/MyDrive/combined.pth'