Create OLD_doors_fasterrcnn.py

OLD_doors_fasterrcnn.py

@@ -0,0 +1,452 @@
+# -*- 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
+"""
+
+# from google.colab.patches import cv2_imshow
+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
+"""# updated for (fullpath, pdf_name)
+
+
+
+"""
+
+def convert2pillow(path):
+  pdf = pdfium.PdfDocument(path)
+  page = pdf.get_page(0)
+  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
+    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)
+
+  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)
+  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)
+
+    # 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]
+  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
+    
+    # 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)
+    
+    # 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)
+    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
+    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'])
+    for page in pdf_document:
+        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/separated_classes.pth'
+# #pdf_name = data
+# for i in range(len(fullpath)):
+#   main_run(fullpath[i], model_path, pdf_name[i])