Upload doors_fasterrcnn.py

doors_fasterrcnn.py

@@ -0,0 +1,206 @@
+# -*- 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
+
+"""# 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 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 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):
+  midpoints = []
+  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)
+    door_width = calculate_width(bbox[i])
+    midpoint = calculate_midpoint(top_left_corner, bottom_right_corner)
+    midpoints.append((midpoint, door_width))
+  return midpoints
+
+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]} pixels width"))
+  return annotations
+
+def add_annotations_to_pdf(image, pdf_name, annotation_s, annotation_d):
+    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()
+
+    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):
+  img_pillow = convert2pillow(pdf_fullpath)
+  new_image = img_pillow.resize((2384,1684))
+  # Specify the number of classes (including the background)
+  num_classes = 6  # 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')
+  model.load_state_dict(torch.load(weights_path, map_location=device))
+  # Set the model to evaluation mode
+  model.eval()
+  # Move the model to the appropriate device
+  model.to(device)
+
+  #START INFERENCE
+  sbox, dbox = ev_model(new_image, model, device, 0.6)
+
+  single_info = mid_points_bbox(sbox)
+  double_info = mid_points_bbox(dbox)
+
+  single_annotations = create_annotations("single door", single_info)
+  double_annotations = create_annotations("double door", double_info)
+
+  pdf_document=add_annotations_to_pdf(new_image, pdf_name, single_annotations, double_annotations)
+  
+  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)
+
+  return annotatedimg,pdf_document
+
+# 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])