Spaces:
Sleeping
Sleeping
| 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 |