import numpy as np import cv2 import copy def calcDistance(point_1, point_2): return pow(pow(point_1[0] - point_2[0], 2) + pow(point_1[1] - point_2[1], 2), 0.5) def calcLineDim(line, lineWidth=-1): if abs(line[0][0] - line[1][0]) > abs(line[0][1] - line[1][1]): if lineWidth < 0 or abs(line[0][1] - line[1][1]) <= lineWidth: return 0 elif abs(line[0][0] - line[1][0]) < abs(line[0][1] - line[1][1]): if lineWidth < 0 or abs(line[0][0] - line[1][0]) <= lineWidth: return 1 else: return -1 class Floorplan(): def __init__(self, filename): self.wallWidth = 0.005 self.filename = filename # Remove 3D material and model references self.wallMats = [] self.iconNodes = {} def read(self): floorplanFile = open(self.filename + '.txt', 'r') self.walls = [] self.doors = [] self.icons = [] self.wallsInt = [] for line in floorplanFile.readlines(): line = line.strip() values = line.split('\t') if len(values) == 2: self.width = float(values[0]) self.height = float(values[1]) self.maxDim = max(self.width, self.height) elif len(values) == 6: wall = [] for i in range(4): wall.append(float(values[i])) lineDim = calcLineDim(((wall[0], wall[1]), (wall[2], wall[3]))) wall[lineDim], wall[2 + lineDim] = min(wall[lineDim], wall[2 + lineDim]), max(wall[lineDim], wall[2 + lineDim]) wall[1 - lineDim] = wall[3 - lineDim] = (wall[1 - lineDim] + wall[3 - lineDim]) / 2 wall.append(int(values[4]) - 1) wall.append(int(values[5]) - 1) for pointIndex in range(2): wall[pointIndex * 2 + 0] /= self.maxDim wall[pointIndex * 2 + 1] /= self.maxDim self.walls.append(wall) wallInt = [] for i in range(4): wallInt.append(int(values[i])) wallInt[lineDim], wallInt[2 + lineDim] = min(wallInt[lineDim], wallInt[2 + lineDim]), max(wallInt[lineDim], wallInt[2 + lineDim]) self.wallsInt.append(wallInt) elif len(values) == 7: item = [] for i in range(4): item.append(float(values[i])) for pointIndex in range(2): item[pointIndex * 2 + 0] /= self.maxDim item[pointIndex * 2 + 1] /= self.maxDim if values[4] == 'door': self.doors.append(item) else: item.append(values[4]) self.icons.append(item) return def processFloorplan(self): # Process exterior walls exteriorWalls = [] for wall in self.walls: if wall[4] == 10 or wall[5] == 10: exteriorWalls.append(copy.deepcopy(wall)) # Process exterior openings (doors/windows) exteriorOpenings = [] for wall in exteriorWalls: lineDim = calcLineDim((wall[:2], wall[2:4])) for doorIndex, door in enumerate(self.doors): if calcLineDim((door[:2], door[2:4])) != lineDim: continue if (door[lineDim] >= wall[lineDim] and door[2 + lineDim] <= wall[2 + lineDim] and abs(door[1 - lineDim] - wall[1 - lineDim]) <= self.wallWidth): exteriorOpenings.append(doorIndex) # Find main entrance door minDistance = 10000 mainDoorIndex = -1 for icon in self.icons: if icon[4] == 'entrance': for doorIndex in exteriorOpenings: door = self.doors[doorIndex] distance = pow(pow((door[0] + door[2]) / 2 - (icon[0] + icon[2]) / 2, 2) + pow((door[1] + door[3]) / 2 - (icon[1] + icon[3]) / 2, 2), 0.5) if distance < minDistance: minDistance = distance mainDoorIndex = doorIndex break # Separate doors and windows newDoors = [] self.windows = [] for doorIndex, door in enumerate(self.doors): if doorIndex == mainDoorIndex or doorIndex not in exteriorOpenings: newDoors.append(door) else: self.windows.append(door) self.doors = newDoors # Find exterior wall loops exteriorWallLoops = [] visitedMask = {} gap = 5.0 / self.maxDim for wallIndex, wall in enumerate(exteriorWalls): if wallIndex in visitedMask: continue visitedMask[wallIndex] = True exteriorWallLoop = [] exteriorWallLoop.append(wall) for loopWall in exteriorWallLoop: for neighborWallIndex, neighborWall in enumerate(exteriorWalls): if neighborWallIndex in visitedMask: continue if calcDistance(neighborWall[:2], loopWall[2:4]) < gap: exteriorWallLoop.append(neighborWall) visitedMask[neighborWallIndex] = True break elif calcDistance(neighborWall[2:4], loopWall[2:4]) < gap: neighborWall[0], neighborWall[2] = neighborWall[2], neighborWall[0] neighborWall[1], neighborWall[3] = neighborWall[3], neighborWall[1] exteriorWallLoop.append(neighborWall) visitedMask[neighborWallIndex] = True break exteriorWallLoops.append(exteriorWallLoop) return exteriorWallLoops def segmentRooms(self): wallMask = np.ones((int(self.height), int(self.width)), np.uint8) * 255 for wall in self.wallsInt: lineDim = calcLineDim(((wall[0], wall[1]), (wall[2], wall[3]))) if lineDim == 0: wallMask[wall[1], wall[0]:wall[2] + 1] = 0 else: wallMask[wall[1]:wall[3] + 1, wall[0]] = 0 cv2.imwrite('test/walls.png', wallMask) numLabels, labels, stats, centroids = cv2.connectedComponentsWithStats(wallMask, 4) print("Number of labels:", numLabels) print("Labels shape:", labels.shape) print("Stats shape:", stats.shape) print("Centroids shape:", centroids.shape) cv2.imwrite('test/rooms.png', labels) return labels def generateFloorplanImage(self): # Create a blank image img = np.ones((int(self.height), int(self.width), 3), dtype=np.uint8) * 255 # Draw walls for wall in self.wallsInt: cv2.line(img, (wall[0], wall[1]), (wall[2], wall[3]), (0, 0, 0), 2) # Draw doors for door in self.doors: door_int = [int(x * self.maxDim) for x in door[:4]] cv2.line(img, (door_int[0], door_int[1]), (door_int[2], door_int[3]), (0, 255, 0), 2) # Draw windows for window in self.windows: window_int = [int(x * self.maxDim) for x in window[:4]] cv2.line(img, (window_int[0], window_int[1]), (window_int[2], window_int[3]), (0, 0, 255), 2) # Draw icons for icon in self.icons: icon_int = [int(x * self.maxDim) for x in icon[:4]] cv2.rectangle(img, (icon_int[0], icon_int[1]), (icon_int[2], icon_int[3]), (255, 0, 0), 2) if len(icon) > 4: cv2.putText(img, icon[4], (icon_int[0], icon_int[1]), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1) return img