Update doors_fasterrcnn.py

doors_fasterrcnn.py

@@ -44,6 +44,11 @@ import google_sheet_Legend
 import torch
 import torchvision
 from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
+from PyPDF2 import PdfReader, PdfWriter
+from PyPDF2.generic import TextStringObject, NameObject, ArrayObject, FloatObject
+from PyPDF2.generic import NameObject, TextStringObject, DictionaryObject, FloatObject, ArrayObject
+from PyPDF2 import PdfReader
+from io import BytesIO
 
 def convert2pillow(path):
   pdf = pdfium.PdfDocument(path)
@@ -178,46 +183,127 @@ def width_as_char(real_width):
     char_width.append(f"{width}m")
   return char_width
 
+def adjustannotations(OutputPdfStage1):
+  input_pdf_path = OutputPdfStage1
+  # output_pdf_path = "AnnotationAdjusted.pdf"
+
+  # Load the input PDF
+  pdf_bytes_io = BytesIO(OutputPdfStage1)
+
+  reader = PdfReader(pdf_bytes_io)
+  writer = PdfWriter()
+
+  # Append all pages to the writer
+  writer.append_pages_from_reader(reader)
+
+  # Add metadata (optional)
+  metadata = reader.metadata
+  writer.add_metadata(metadata)
+
+  # Iterate over pages
+  for page_index, page in enumerate(writer.pages):
+      if "/Annots" in page:
+          annotations = page["/Annots"]
+          for annot_index, annot in enumerate(annotations):
+              obj = annot.get_object()
+
+              # print("obj", obj)
+              # print(obj.get("/IT"))
+
+              if obj.get("/Subtype") == "/Line":
+                            print("AWL ANNOT IF")
+                            # Check the /IT value to differentiate annotations
+                            if "/Contents" in obj and "m" in obj["/Contents"]:
+                                print("Tany IF")
+                                obj.update({
+                                    NameObject("/Measure"): DictionaryObject({
+                                        NameObject("/Type"): NameObject("/Measure"),
+                                        NameObject("/L"): DictionaryObject({
+                                            NameObject("/G"): FloatObject(1),
+                                            NameObject("/U"): TextStringObject("sq m"),  # Unit of measurement for area
+                                        }),
+
+                                    }),
+                                    NameObject("/IT"): NameObject("/LineDimension"),  # Use more distinctive name
+                                    NameObject("/Subj"): TextStringObject("Length Measurement"),   # Intent explicitly for Area
+                                })
+              print("After Update:", obj)
+  
+  # # Save the modified PDF
+
+  output_pdf_io = BytesIO()
+  writer.write(output_pdf_io)
+  output_pdf_io.seek(0)
+
+  print("Annotations updated and saved ")
+  return output_pdf_io.read()
+
 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()
-
+    pdf_document = fitz.open('pdf',pdf_name)
+    page=pdf_document[0]
+    rotationOld=page.rotation
+    derotationMatrix=page.derotation_matrix
+    print('rotationOld',rotationOld)
+    if page.rotation!=0:
+      rotationangle = page.rotation
+      page.set_rotation(0)
+    print('rotationnew',page.rotation)
     # 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)
+    # 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())
+    # # 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]
+      p_midpoint = fitz.Point(x, y) * derotationMatrix
+      rect = fitz.Rect(p_midpoint.x, p_midpoint.y, p_midpoint.x + 200, p_midpoint.y + 50)
       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 = 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()
+        # 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()
+       
+      annot = page.add_line_annot(fitz.Point(*lines[i][0]) * derotationMatrix, fitz.Point(*lines[i][1]) * derotationMatrix)
+      annot.set_border(width=2, dashes=None)  # Optional border styling
+      annot.set_colors(stroke=(1, 0, 0))  # Set the line color to red
+      annot.set_info(content=str(1.0)+' m',subject='Perimeter Measurement', title="ADR Team")
+      annot.update()
 
+      annot = page.add_line_annot(fitz.Point(*sanda[i][0]) * derotationMatrix, fitz.Point(*lines[i][0]) * derotationMatrix)
+      annot.set_border(width=2, dashes=None)  # Optional border styling
+      annot.set_colors(stroke=(1, 0, 0))  # Set the line color to red
+      annot.set_info(content=str(1.0)+' m',subject='Perimeter Measurement', title="ADR Team")
+      annot.update()
 
-    output_pdf_path = pdf_name+"_annotated.pdf"
+      annot = page.add_line_annot(fitz.Point(*sanda[i][1]) * derotationMatrix, fitz.Point(*lines[i][1]) * derotationMatrix)
+      annot.set_border(width=2, dashes=None)  # Optional border styling
+      annot.set_colors(stroke=(1, 0, 0))  # Set the line color to red
+      annot.set_info(content=str(1.0)+' m',subject='Perimeter Measurement', title="ADR Team")
+      annot.update()
+    # output_pdf_path = pdf_name+"_annotated.pdf"
     # Save the PDF with annotations
 
+    page.set_rotation(rotationOld)
     return pdf_document
     
 
 
-def main_run(pdf_fullpath, weights_path, pdf_name,pdfpath,ratio): ####pdf_fullpath here is the data and not the path 
+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)
     
     num_classes = 10  # classes + background
@@ -246,13 +332,15 @@ def main_run(pdf_fullpath, weights_path, pdf_name,pdfpath,ratio): ####pdf_fullpa
     real_width = pxl2meter(width_pixels, ratio)
     char_width = width_as_char(real_width)
 
-    pdf_document = add_annotations_to_pdf(img_pillow, pdf_name, lines, sanda, char_width, line_midpoint)
-
+    pdf_document = add_annotations_to_pdf(img_pillow, pdf_fullpath, lines, sanda, char_width, line_midpoint)
+    modified_pdf_data = pdf_document.tobytes()
+    OutputPdfStage2=adjustannotations(modified_pdf_data)
     #Dataframe for Doors count
     #doors_count = {'Type': ['Single Doors', 'Double Doors'], 'Quantity': [len(sbox), len(dbox)]}
     doors_count = {'Type': ['Single Doors'], 'Quantity': [len(real_width)]}
     df_doors = pd.DataFrame(doors_count)
-
+    
+    doc2 =fitz.open('pdf',OutputPdfStage2)
     page=pdf_document[0]
     pix = page.get_pixmap()  # render page to an image
     pl=Image.frombytes('RGB', [pix.width,pix.height],pix.samples)
@@ -268,7 +356,9 @@ 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 
+    # OutputPdfStage2=OutputPdfStage2.read()
+    # doc2 =fitz.open('pdf',OutputPdfStage2)
+    return annotatedimg, doc2 , spreadsheet_url, list1, df_doors 
 
 # model_path =  '/content/drive/MyDrive/combined.pth'
 # #pdf_name = data