Update doors_fasterrcnn.py

doors_fasterrcnn.py

@@ -101,6 +101,35 @@ def calculate_width(bbox):
     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, 4)
+    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
@@ -109,8 +138,8 @@ def calculate_midpoint(top_left, bottom_right):
     ym = int((y1 + y2) / 2)
     return (xm, ym)
 
-def mid_points_bbox(bbox, ratio):
-  midpoints = []
+def mid_points_bbox(bbox):
+  midpoint = []
   for i in range(len(bbox)):
     x1 = int(bbox[i][0])
     y1 = int(bbox[i][1])
@@ -118,20 +147,17 @@ def mid_points_bbox(bbox, ratio):
     y2 = int(bbox[i][3])
     top_left_corner = (x1, y1)
     bottom_right_corner = (x2, y2)
-    door_width = calculate_width(bbox[i])
-    door_width = door_width * ratio
-    midpoint = calculate_midpoint(top_left_corner, bottom_right_corner)
-    midpoints.append((midpoint, door_width))
-  return midpoints
+    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][0],midpoints[i][0][1], door+f" with {midpoints[i][1]} meters width"))
+    annotations.append((midpoints[i][0],midpoints[i][1], door))
   return annotations
 
-def add_annotations_to_pdf(image, pdf_name, annotation_s, annotation_d):
+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
@@ -161,8 +187,36 @@ def add_annotations_to_pdf(image, pdf_name, annotation_s, annotation_d):
         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=12, fontname="helv", text_color=(0,1,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=12, fontname="helv", text_color=(0,1,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()