python - 如何使用python从蓝图中分割房间(How to segment rooms from a blueprint using python)

Question

So this code is able to roughly segment this blueprint (Image 1), into rooms which it identifies with a randomized color.

(因此，此代码能够将该蓝图（图1）粗略地划分为多个房间，并以随机的颜色进行标识。)

However as shown, it's not quite as good at reading the rooms as it should be, sometimes even ignoring the room and just erasing it as seen in the output (image 2).

(但是，如图所示，在读取房间方面并没有达到应有的水平，有时甚至忽略了房间并只是擦除了输出中的图像（图2）。)

I would like to ask, how can i improve the code such that it's able to give me a better output.

(我想问一下，我该如何改善代码，使其能够为我提供更好的输出。)

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?

  ask by Shaun translate from so

Answer 1

等待大神答复