Вот кое-что, что я придумал.Это не идеально (я сделал несколько комментариев, что вы, возможно, захотите попробовать), и будет лучше, если вы улучшите качество входного изображения.
import cv2
import numpy as np
def find_rooms(img, noise_removal_threshold=25, corners_threshold=0.1,
room_closing_max_length=100, gap_in_wall_threshold=500):
"""
:param img: grey scale image of rooms, already eroded and doors removed etc.
:param noise_removal_threshold: Minimal area of blobs to be kept.
:param corners_threshold: Threshold to allow corners. Higher removes more of the house.
:param room_closing_max_length: 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: A colored version of the input image, where each room has a random 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_removal_threshold:
cv2.fillPoly(mask, [contour], 255)
img = ~mask
# Detect corners (you can play with the parameters here)
dst = cv2.cornerHarris(img ,2,3,0.04)
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.
# You could try to disallow drawing through other existing lines for example.
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_closing_max_length:
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_closing_max_length:
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("/home/veith/Pictures/room.png", 0)
rooms, colored_house = find_rooms(img.copy())
cv2.imshow('result', colored_house)
cv2.waitKey()
cv2.destroyAllWindows()
Это покажет такое изображение, где каждая комнатаимеет случайный цвет: вы можете видеть, что иногда он находит комнату, где ее нет, но я думаю, что это достойная отправная точка для вас.
Для этого я использовал скриншот изображения в вашем вопросе.Вы можете использовать возвращенные маски каждой комнаты, чтобы проиндексировать исходное изображение и обрезать его.Для обрезки просто используйте что-то вроде (не проверено, но должно работать по большей части):
for room in rooms:
crop = np.zeros_like(room).astype(np.uint8)
crop[room] = original_img[room] # Get the original image from somewhere
# if you need to crop the image into smaller parts as big as each room
r, c = np.nonzero(room)
min_r, max_r = r.argmin(), r.argmax()
min_c, max_c = c.argmin(), c.argmax()
crop = crop[min_r:max_r, min_c:max_c]
cv2.imshow("cropped room", crop)
cv2.waitKey()
cv2.destroyAllWindows()