# -*- coding: utf-8 -*- """MartheDeployment_Doors_fasterRCNN.ipynb Automatically generated by Colab. Original file is located at https://colab.research.google.com/drive/1kgEtpfNt0jxSwPRhOzODIC6P_prg-c4L ## Libraries """ import cv2 import numpy as np import pandas as pd import statistics from statistics import mode from PIL import Image # pip install PyPDF2 # pip install PyMuPDF # pip install pip install PyMuPDF==1.19.0 import io # !pip install pypdfium2 import pypdfium2 as pdfium import fitz # PyMuPDF import os #drive.mount("/content/drive", force_remount=True) import torch from torchvision.models.detection.faster_rcnn import FastRCNNPredictor from PIL import Image, ImageDraw import torchvision.transforms.functional as F import matplotlib.pyplot as plt 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) page = pdf.get_page(0) pil_image = page.render().to_pil() return pil_image # Function to get the model def get_model(num_classes): # Load a pre-trained Faster R-CNN model with a ResNet-50-FPN backbone model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True) # Get the number of input features for the classifier in_features = model.roi_heads.box_predictor.cls_score.in_features # Replace the pre-trained head with a new one for our number of classes model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 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) 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 = [] singles = 0 door_type = [] 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) point_st = (int(xmin), int(ymax)) point_end = (int(xmax),int(ymax)) lines.append((point_st,point_end)) line_midpoint.append(calculate_midpoint(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) door_type.append(0) singles +=1 #horz_upper if door_inf[1] == 3: #for drawing point_st = (int(xmin),int(ymin)) point_end = (int(xmax),int(ymin)) #point_st = (int(xmin),int(ymin) -5) #point_end = (int(xmax),int(ymin) - 5) lines.append((point_st,point_end)) line_midpoint.append(calculate_midpoint(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) singles +=1 door_type.append(0) #vert_right if door_inf[1] == 4: #for drawing #point_st = (int(xmax) + 5,int(ymin)) #point_end = (int(xmax) + 5,int(ymax)) point_st = (int(xmax), int(ymin)) point_end = (int(xmax), int(ymax)) lines.append((point_st,point_end)) line_midpoint.append(calculate_midpoint(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) singles +=1 door_type.append(0) #vert_left if door_inf[1] == 5: #for drawing point_st = (int(xmin),int(ymin)) point_end = (int(xmin),int(ymax)) #point_st = (int(xmin) -5,int(ymin)) #point_end = (int(xmin) -5,int(ymax)) lines.append((point_st,point_end)) line_midpoint.append(calculate_midpoint(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) singles +=1 door_type.append(0) return width_pixels, lines, line_midpoint, singles, door_type def get_door_info_double(doors_info_double, width_pixels, lines, line_midpoint, door_type): doubles = 0 for door_inf in doors_info_double: xmin, ymin, xmax, ymax = door_inf[0] #double_bottom if door_inf[1] == 1: point_st = (int(xmin), int(ymax)) point_end = (int(xmax),int(ymax)) #point_st = (int(xmin), int(ymax) + 5) #point_end = (int(xmax),int(ymax) + 5) lines.append((point_st,point_end)) line_midpoint.append(calculate_midpoint(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) doubles +=1 door_type.append(1) #double_upper if door_inf[1] == 7: #for drawing point_st = (int(xmin),int(ymin)) point_end = (int(xmax),int(ymin)) #point_st = (int(xmin),int(ymin) -5) #point_end = (int(xmax),int(ymin) - 5) lines.append((point_st,point_end)) line_midpoint.append(calculate_midpoint(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) doubles +=1 door_type.append(1) #double_right if door_inf[1] == 8: #for drawing #point_st = (int(xmax) + 5,int(ymin)) #point_end = (int(xmax) + 5,int(ymax)) point_st = (int(xmax),int(ymin)) point_end = (int(xmax),int(ymax)) lines.append((point_st,point_end)) line_midpoint.append(calculate_midpoint(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) doubles +=1 door_type.append(1) #double_left if door_inf[1] == 4: #for drawing #point_st = (int(xmin) -5,int(ymin)) #point_end = (int(xmin) -5,int(ymax)) point_st = (int(xmin),int(ymin)) point_end = (int(xmin),int(ymax)) lines.append((point_st,point_end)) line_midpoint.append(calculate_midpoint(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) doubles +=1 door_type.append(1) return width_pixels, lines, line_midpoint, doubles, door_type 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}") 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, char_width, line_midpoint, door_type): image_width, image_height = image.size # Create a new PDF document 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) # # 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] annot = page.add_freetext_annot(rect, text, fontsize=10, fontname="helv", text_color=(1, 0, 0)) annot.update() #for i in range(len(doubleD_bbox)): #a7seb ana el midpoint beta3 el bbox 3ashan a7ot el width feha #x,y = calculate_midpoint((doubleD_bbox[i][0], doubleD_bbox[i][1]), (doubleD_bbox[i][2],doubleD_bbox[i][3])) #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_d_width[i] #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)): l_points = [fitz.Point(*lines[i][0])* derotationMatrix, fitz.Point(*lines[i][1])* derotationMatrix] #l_points = [fitz.Point(*sanda[i][0])* derotationMatrix, fitz.Point(*lines[i][0])* derotationMatrix, fitz.Point(*lines[i][1])* derotationMatrix, fitz.Point(*sanda[i][1])* derotationMatrix] annot = page.add_polyline_annot(l_points) 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(char_width[i])+' m',subject='Perimeter Measurement', title="ADR Team") annot.update() for i in range(len(door_type)): x, y = line_midpoint[i] p_midpoint = fitz.Point(x, y) * derotationMatrix if door_type[i] == 0: text = "Single Door" else: text = "Double Door" # Create an annotation (sticky note) annot = page.add_text_annot((p_midpoint.x, p_midpoint.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() page.set_rotation(rotationOld) return pdf_document def main_run(img_pillow,pdf_fullpath, weights_path, weights_path2, pdf_name,pdfpath,ratio): ####pdf_fullpath here is the data and not the path img_pillow = convert2pillow(pdf_fullpath) # For Single Doors num_classes = 10 # classes + background # 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) # For Double Doors num_classes2 = 12 # classes + background # Load the model with the specified number of classes model2 = get_model(num_classes2) # Load the saved model's state dictionary with map_location to handle CPU device2 = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') try: model2.load_state_dict(torch.load(weights_path2, map_location=device2), strict=False) except RuntimeError as e: print(f"Error loading model state_dict: {e}") return # Set the model to evaluation mode model2.eval() # Move the model to the appropriate device model2.to(device2) # START INFERENCE doors_info = ev_model(img_pillow, model, device, 0.9) doors_info_double = ev_model(img_pillow, model2, device2, 0.8) width_pixels, lines, line_midpoint, single_count, door_type = get_door_info(doors_info) width_pixels, lines, line_midpoint, double_count, door_type = get_door_info_double(doors_info_double, width_pixels, lines, line_midpoint, door_type) real_width = pxl2meter(width_pixels, ratio) char_width = width_as_char(real_width) pdf_document = add_annotations_to_pdf(img_pillow, pdf_fullpath, lines, char_width, line_midpoint, door_type) modified_pdf_data = pdf_document.tobytes() OutputPdfStage2=adjustannotations(modified_pdf_data) #Dataframe for Doors count doors_count = {'Type': ['Single Doors', 'Double Doors'], 'Quantity': [single_count, double_count]} 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) img=np.array(pl) annotatedimg = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) df_doors = df_doors.fillna(' ') gc,spreadsheet_service,spreadsheetId, spreadsheet_url , namepathArr=google_sheet_Legend.legendGoogleSheets(df_doors , pdf_name,pdfpath) list1=pd.DataFrame(columns=['content', 'id', 'subject','color']) # for page in doc: for page in doc2: # Iterate through annotations on the page for annot in page.annots(): # Get the color of the annotation annot_color = annot.colors if annot_color is not None: # annot_color is a dictionary with 'stroke' and 'fill' keys stroke_color = annot_color.get('stroke') # Border color print('strokeee',stroke_color) # fill_color = annot_color.get('fill') # Fill color # if fill_color: # v='fill' # print('fill') if stroke_color: v='stroke' # x,y,z=int(annot_color.get(v)[0]*255),int(annot_color.get(v)[1]*255),int(annot_color.get(v)[2]*255) list1.loc[len(list1)] =[annot.info['content'],annot.info['id'],annot.info['subject'],[255,0,0]] print('list1',list1) # modify this return # 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 # for i in range(len(fullpath)): # main_run(fullpath[i], model_path, pdf_name[i])