from PIL import Image import numpy as np import matplotlib.pyplot as plt from skimage import feature from xml.dom.minidom import getDOMImplementation from utils import * def run_AI(): print('AI analysis') with open("image_paths.txt", "r") as file: imageList = [line.strip() for line in file] wave_corect="auto" IMAGE_DIR = './data/data1' # 1. stitching if len(imageList) > 1: print("Stitching images...") stitcher = Stitcher(wave_correct_kind=wave_corect) panorama = stitcher.stitch(imageList) else: panorama = cv2.imread(imageList[0]) SAVED_DIR = './results/{}'.format(IMAGE_DIR.split("/")[-1]) if not os.path.exists(SAVED_DIR): os.makedirs(SAVED_DIR) cv2.imwrite('./results/{}/panorama.jpg'.format(IMAGE_DIR.split("/")[-1]), panorama) # 2. detection saved_panorama = panorama config_file = './mmyolo/configs/1_crack/final_model.py' checkpoint_file = './pretrained_weights/best_pretrained.pth' # build the model from a config file and a checkpoint file model = init_detector(config_file, checkpoint_file, device='cuda:0') crop_height = 3000 crop_width = 4000 total_bboxes = [] total_labels = [] pano_height = panorama.shape[0] pano_width = panorama.shape[1] for i in tqdm(range(0,panorama.shape[0],crop_height)): for j in range(0,panorama.shape[1],crop_width): crop_img = panorama[i:i+crop_height, j:j+crop_width] # plt.imshow(crop_img) # plt.show() # inference results = inference_detector(model, crop_img) raw_public_bboxes, raw_public_labels, raw_public_scores = mmdet3x_convert_to_bboxes_mmdet(results,0.1) merged_boxes, merged_labels = merge_boxes_with_labels(raw_public_bboxes, raw_public_labels) for i, box in enumerate(merged_boxes): # image = cv2.rectangle(crop_img, (box[0], box[1]), (box[2], box[3]), (255,0,0), 2) # image = cv2.putText(image, str(merged_labels[i]).split("_")[-1], (box[0], box[1]), cv2.FONT_HERSHEY_SIMPLEX, 3, (255, 0, 0),4) # add confidence score org_coordinate = [box[0]+j, box[1]+i, box[2]+j, box[3]+i] total_bboxes.append(org_coordinate) total_labels.append(merged_labels[i]) # plt.imshow(crop_img) # plt.show() for i,box in enumerate(total_bboxes): # saved_panaroma = cv2.putText(saved_panaroma, str(total_labels[i]).split("_")[-1], (box[0], box[1]), cv2.FONT_HERSHEY_SIMPLEX, 3, (0, 0, 255),5) cropped_image = saved_panorama[box[1]:box[3], box[0]:box[2]] # edge detection edges = cv2.Canny(cropped_image,100,200) # calculate the number of pixels in the edge # if the number of pixels is less than 1000, then remove the bbox saved_panaroma = cv2.rectangle(saved_panorama, (box[0], box[1]), (box[2], box[3]), (0,0,255), 10) cv2.imwrite('./results/{}/panorama_detected.jpg'.format(IMAGE_DIR.split("/")[-1]), panorama) # 3. Extract Table from matplotlib.patches import Rectangle # raw_panorama = cv2.imread('./results/{}/panorama.jpg'.format(IMAGE_DIR.split("/")[-1]),0) img = cv2.imread('./results/{}/panorama.jpg'.format(IMAGE_DIR.split("/")[-1])) gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) _, binary = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY) binary_inv = cv2.bitwise_not(binary) cv2.imwrite('test.png', binary_inv) img_path = 'test.png' # Load the binary image binary_img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) # Invert the binary image since we're interested in the black shape on white background binary_img_inv = cv2.bitwise_not(binary_img) contours, _ = cv2.findContours(binary_img_inv, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) sorted_contours = sorted(contours, key=cv2.contourArea, reverse=True) largest_contour = sorted_contours[0] contour_img = np.zeros_like(binary_img) contour_img = cv2.drawContours(contour_img, [largest_contour], -1, (255),10 ) vis_panorama = np.zeros((contour_img.shape[0], contour_img.shape[1])) plt.figure(figsize=(15, 10)) print(total_bboxes) print(total_labels) for i,box in enumerate(total_bboxes): # add pacth here x1, y1, x2, y2 = box # Add a rectangle patch with "*" hatching pattern if total_labels[i] == 'class_1': plt.gca().add_patch(Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, hatch='///', edgecolor='red')) elif total_labels[i] == 'class_2': plt.gca().add_patch(Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, hatch='\\', edgecolor='blue')) elif total_labels[i] == 'class_3': plt.gca().add_patch(Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, hatch='-', edgecolor='green')) elif total_labels[i] == 'class_4': plt.gca().add_patch(Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, hatch='|', edgecolor='yellow')) box_area = (x2-x1)*(y2-y1) print(box_area) # print total_bboxes and total_labels as a table import pandas as pd df = pd.DataFrame(total_bboxes, columns=['x1', 'y1', 'x2', 'y2']) df['label'] = total_labels df.to_csv('./results/{}/detection_coor.csv'.format(IMAGE_DIR.split("/")[-1]), index=False) # 4. Convert to CAD contour_img[contour_img<255] = 0 plt.imshow(contour_img) plt.axis("off") plt.tight_layout() plt.savefig('./results/{}/contour_panorama_detected.jpg'.format(IMAGE_DIR.split("/")[-1]), format='jpg', dpi=300) # Adjust dpi for desired quality raw_image = cv2.imread('./results/{}/contour_panorama_detected.jpg'.format(IMAGE_DIR.split("/")[-1])) edges = cv2.Canny(raw_image,200,250) plt.imshow(edges) plt.axis('off') cv2.imwrite('./results/{}/edge.jpg'.format(IMAGE_DIR.split("/")[-1]), edges) # Function to create SVG path from edges def create_svg_path(edges_array): dom = getDOMImplementation().createDocument(None, "svg", None) svg = dom.documentElement svg.setAttribute('xmlns', "http://www.w3.org/2000/svg") svg.setAttribute('xmlns:xlink', "http://www.w3.org/1999/xlink") svg.setAttribute('width', str(edges_array.shape[1])) svg.setAttribute('height', str(edges_array.shape[0])) # Go through the array and create line segments for edges for i in range(edges_array.shape[0]): for j in range(edges_array.shape[1]): if edges_array[i, j]: # Create a small line segment (a dot) new_element = dom.createElement('line') new_element.setAttribute('x1', str(j)) new_element.setAttribute('y1', str(i)) new_element.setAttribute('x2', str(j+1)) new_element.setAttribute('y2', str(i)) new_element.setAttribute('style', "stroke:rgb(0,0,0);stroke-width:1") svg.appendChild(new_element) return dom.toxml() # Load the image edge_image_path = './results/{}/edge.jpg'.format(IMAGE_DIR.split("/")[-1]) edge_image = Image.open(edge_image_path).convert("L") # Convert image to numpy array image_array = np.array(edge_image) # Detect edges edges = feature.canny(image_array) # Generate SVG svg_data = create_svg_path(edges) # Save the SVG svg_output_path = "./results/{}/CAD_FORMAT.svg".format(IMAGE_DIR.split("/")[-1]) with open(svg_output_path, "w") as f: f.write(svg_data) import svgpathtools import ezdxf # Read the SVG file using svgpathtools or any other library that can parse SVG files. paths, attributes = svgpathtools.svg2paths('./results/{}/CAD_FORMAT.svg'.format(IMAGE_DIR.split("/")[-1])) # Create a new DXF document using ezdxf. doc = ezdxf.new(dxfversion='R2010') msp = doc.modelspace() # Iterate over the parsed SVG paths and convert them to DXF entities. for path in paths: # Break the path down into individual segments. for segment in path: # Depending on the segment type (Line, CubicBezier, etc.), create the appropriate DXF entity. if isinstance(segment, svgpathtools.Line): msp.add_line((segment.start.real, segment.start.imag), (segment.end.real, segment.end.imag)) # Add additional cases for other segment types like CubicBezier, QuadraticBezier, Arc, etc. # Save the DXF document to a file. doc.saveas('./results/{}/dxf_file.dxf'.format(IMAGE_DIR.split("/")[-1])) print('done') run_AI()