python - 如何使用python从蓝图中分割房间

import numpy as np def find_rooms(img, noise_reduction=10, corners_threshold=0.0000001, room_close=2, gap_in_wall_threshold=0.000001): # :param img: grey scale image of rooms, already eroded and doors removed etc. # :param noise_reduction: Amount of noise removed. # :param corners_threshold: Corners to retained, higher value = more of house removed. # :param room_close: Maximum line length to add to close off open doors. # :param gap_in_wall_threshold: Minimum number of pixels to identify component as room instead of hole in the wall. # :return: rooms: list of numpy arrays containing boolean masks for each detected room # colored_house: Give room a color. assert 0 <= corners_threshold <= 1 # Remove noise left from door removal img[img < 128] = 0 img[img > 128] = 255 contours, _ = cv2.findContours(~img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) mask = np.zeros_like(img) for contour in contours: area = cv2.contourArea(contour) if area > noise_reduction: cv2.fillPoly(mask, [contour], 255) img = ~mask # Detect corners (you can play with the parameters here) #harris corner detection dst = cv2.cornerHarris(img, 4,3,0.000001) dst = cv2.dilate(dst,None) corners = dst > corners_threshold * dst.max() # Draw lines to close the rooms off by adding a line between corners on the same x or y coordinate # This gets some false positives. # Can try disallowing drawing through other existing lines, need to test. for y,row in enumerate(corners): x_same_y = np.argwhere(row) for x1, x2 in zip(x_same_y[:-1], x_same_y[1:]): if x2[0] - x1[0] < room_close: color = 0 cv2.line(img, (x1, y), (x2, y), color, 1) for x,col in enumerate(corners.T): y_same_x = np.argwhere(col) for y1, y2 in zip(y_same_x[:-1], y_same_x[1:]): if y2[0] - y1[0] < room_close: color = 0 cv2.line(img, (x, y1), (x, y2), color, 1) # Mark the outside of the house as black contours, _ = cv2.findContours(~img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) contour_sizes = [(cv2.contourArea(contour), contour) for contour in contours] biggest_contour = max(contour_sizes, key=lambda x: x[0])[1] mask = np.zeros_like(mask) cv2.fillPoly(mask, [biggest_contour], 255) img[mask == 0] = 0 # Find the connected components in the house ret, labels = cv2.connectedComponents(img) img = cv2.cvtColor(img,cv2.COLOR_GRAY2RGB) unique = np.unique(labels) rooms = [] for label in unique: component = labels == label if img[component].sum() == 0 or np.count_nonzero(component) < gap_in_wall_threshold: color = 0 else: rooms.append(component) color = np.random.randint(0, 255, size=3) img[component] = color return rooms, img #Read gray image img = cv2.imread('mask6.png', 0) rooms, colored_house = find_rooms(img.copy()) cv2.imshow('result', colored_house) cv2.waitKey() cv2.destroyAllWindows()``` I have tried a few methods such as binarization, and to some extent, erode and dilate to not a great improvement. Are there any other methods I am missing out, or ways that i can improve my code?