Spaces:
Build error
Build error
Download util/lua/floorplan_utils.lua from rawanessam/FloorPlanTransformation: direct link, hf CLI and curl.
- Browser
- Download file 376 kB
-
https://huggingface.co/spaces/rawanessam/FloorPlanTransformation/resolve/main/util/lua/floorplan_utils.lua
- Command line
-
hf download hf://spaces/rawanessam/FloorPlanTransformation/util/lua/floorplan_utils.lua
-
curl -L -o floorplan_utils.lua https://huggingface.co/spaces/rawanessam/FloorPlanTransformation/resolve/main/util/lua/floorplan_utils.lua
376 kB
| require 'csvigo' | |
| require 'image' | |
| local pl = require 'pl.import_into' () | |
| cv = require 'cv' | |
| require 'cv.imgproc' | |
| py = require('python') | |
| --paths.dofile('/home/chenliu/Projects/Floorplan/floorplan/InverseCAD/models/SpatialSymmetricPadding.lua') | |
| local utils = {} | |
| function utils.getSortedModes() | |
| return {'points', 'doors', 'icons', 'labels'} | |
| end | |
| function utils.getNumberMap() | |
| if not utils.numberMap then | |
| local numberMap = {} | |
| numberMap.labels = {} | |
| numberMap.labels['living_room'] = 1 | |
| numberMap.labels['kitchen'] = 2 | |
| numberMap.labels['bedroom'] = 3 | |
| numberMap.labels['bathroom'] = 4 | |
| numberMap.labels['restroom'] = 5 | |
| numberMap.labels['balcony'] = 6 | |
| numberMap.labels['closet'] = 7 | |
| numberMap.labels['corridor'] = 8 | |
| numberMap.labels['washing_room'] = 9 | |
| numberMap.labels['PS'] = 10 | |
| numberMap.icons = {} | |
| numberMap.icons['bathtub'] = {1} | |
| numberMap.icons['cooking_counter'] = {2, 2} | |
| numberMap.icons['toilet'] = {3} | |
| numberMap.icons['entrance'] = {4} | |
| numberMap.icons['washing_basin'] = {5, 6, 5, 5} | |
| numberMap.icons['special'] = {8, 8, 9} | |
| numberMap.icons['stairs'] = {10} | |
| numberMap.doors = {} | |
| numberMap.doors['door'] = 1 | |
| utils.numberMap = numberMap | |
| end | |
| return utils.numberMap | |
| end | |
| function utils.getNumber(mode, itemInfo) | |
| local numberMap = utils.getNumberMap() | |
| if mode == 'points' then | |
| return (itemInfo[2] - 1) * 4 + itemInfo[3] | |
| elseif mode == 'doors' then | |
| return itemInfo[2] | |
| elseif mode == 'icons' then | |
| return numberMap[mode][itemInfo[1]][itemInfo[2]] | |
| elseif mode == 'labels' then | |
| return numberMap[mode][itemInfo[1]] | |
| else | |
| assert(false) | |
| end | |
| end | |
| function utils.getNameMap() | |
| if not utils.nameMap then | |
| local nameMap = {} | |
| local numberMap = utils.getNumberMap() | |
| for mode, map in pairs(numberMap) do | |
| nameMap[mode] = {} | |
| for name, number in pairs(map) do | |
| if type(number) == "table" then | |
| for _, num in pairs(number) do | |
| if not nameMap[mode][num] then | |
| nameMap[mode][num] = {name, _, 1} | |
| end | |
| end | |
| else | |
| if not nameMap[mode][number] then | |
| nameMap[mode][number] = {name, 1, 1} | |
| end | |
| end | |
| end | |
| end | |
| utils.nameMap = nameMap | |
| end | |
| return utils.nameMap | |
| end | |
| function utils.getItemInfo(mode, number) | |
| local nameMap = utils.getNameMap() | |
| if mode == 'points' then | |
| return {'point', math.floor((number - 1) / 4) + 1, (number - 1) % 4 + 1} | |
| elseif mode == 'doors' then | |
| return {'door', number, 1} | |
| elseif mode == 'icons' then | |
| number = math.min(number, #nameMap[mode]) | |
| return nameMap[mode][number] | |
| elseif mode == 'labels' then | |
| number = math.min(number, #nameMap[mode]) | |
| return nameMap[mode][number] | |
| else | |
| assert(false) | |
| end | |
| end | |
| function utils.keyMap() | |
| local keyMap = {} | |
| keyMap['b'] = {'icons', 'bathtub'} | |
| keyMap['c'] = {'icons', 'cooking_counter'} | |
| keyMap['t'] = {'icons', 'toilet'} | |
| keyMap['q'] = {'icons', 'special'} | |
| keyMap['e'] = {'icons', 'entrance'} | |
| keyMap['w'] = {'icons', 'washing_basin'} | |
| keyMap['s'] = {'icons', 'stairs'} | |
| for i = 1, 10 do | |
| keyMap[tostring(i % 10)] = {'labels', utils.getItemInfo('labels', i)[1]} | |
| end | |
| keyMap['a'] = {'walls', 'wall'} | |
| keyMap['d'] = {'doors', 'door'} | |
| return keyMap | |
| end | |
| function utils.modeMap() | |
| local modeMap = {} | |
| local keyMap = utils.keyMap() | |
| for _, modeNamePair in pairs(keyMap) do | |
| modeMap[modeNamePair[2]] = modeNamePair[1] | |
| end | |
| modeMap['point'] = 'points' | |
| return modeMap | |
| end | |
| function utils.numItemsPerCell() | |
| local numItemsPerCell = {} | |
| numItemsPerCell.points = 2 | |
| numItemsPerCell.doors = 2 | |
| numItemsPerCell.icons = 2 | |
| numItemsPerCell.labels = 2 | |
| return numItemsPerCell | |
| end | |
| function utils.numItemsGlobal() | |
| local numItemsGlobal = {} | |
| numItemsGlobal['bathtub'] = 1 | |
| numItemsGlobal['cooking_counter'] = 1 | |
| numItemsGlobal['toilet'] = 2 | |
| numItemsGlobal['special'] = 4 | |
| numItemsGlobal['entrance'] = 1 | |
| numItemsGlobal['washing_basin'] = 4 | |
| numItemsGlobal['stairs'] = 4 | |
| numItemsGlobal['living_room'] = 1 | |
| numItemsGlobal['kitchen'] = 1 | |
| numItemsGlobal['bedroom'] = 4 | |
| numItemsGlobal['bathroom'] = 1 | |
| numItemsGlobal['restroom'] = 2 | |
| numItemsGlobal['balcony'] = 2 | |
| numItemsGlobal['closet'] = 4 | |
| numItemsGlobal['corridor'] = 1 | |
| numItemsGlobal['washing_room'] = 1 | |
| numItemsGlobal['PS'] = 2 | |
| return numItemsGlobal | |
| end | |
| function utils.numFeaturesPerItem() | |
| local numFeaturesPerItem = {} | |
| numFeaturesPerItem.points = 5 | |
| numFeaturesPerItem.doors = 5 | |
| numFeaturesPerItem.icons = 5 | |
| numFeaturesPerItem.labels = 5 | |
| return numFeaturesPerItem | |
| end | |
| function utils.offsetsBB() | |
| local numItemsPerCell = utils.numItemsPerCell() | |
| local numFeaturesPerItem = utils.numFeaturesPerItem() | |
| local offsetsBB = {} | |
| offsetsBB.points = 0 | |
| offsetsBB.doors = offsetsBB.points + numFeaturesPerItem.points * numItemsPerCell.points | |
| offsetsBB.icons = offsetsBB.doors + numFeaturesPerItem.doors * numItemsPerCell.doors | |
| offsetsBB.labels = offsetsBB.icons + numFeaturesPerItem.icons * numItemsPerCell.icons | |
| return offsetsBB | |
| end | |
| function utils.numFeaturesBB() | |
| local numItemsPerCell = utils.numItemsPerCell() | |
| local offsetsBB = utils.offsetsBB() | |
| local numFeaturesPerItem = utils.numFeaturesPerItem() | |
| local numFeaturesBB = offsetsBB.labels + numFeaturesPerItem.labels * numItemsPerCell.labels | |
| return numFeaturesBB | |
| end | |
| function utils.offsetsClass() | |
| local numItemsPerCell = utils.numItemsPerCell() | |
| local numFeaturesBB = utils.numFeaturesBB() | |
| local offsetsClass = {} | |
| offsetsClass.points = numFeaturesBB | |
| offsetsClass.doors = offsetsClass.points + numItemsPerCell.points | |
| offsetsClass.icons = offsetsClass.doors + numItemsPerCell.doors | |
| offsetsClass.labels = offsetsClass.icons + numItemsPerCell.icons | |
| return offsetsClass | |
| end | |
| function utils.numFeaturesClass() | |
| local numItemsPerCell = utils.numItemsPerCell() | |
| local numFeaturesClass = numItemsPerCell.points + numItemsPerCell.doors + numItemsPerCell.icons + numItemsPerCell.labels | |
| return numFeaturesClass | |
| end | |
| function utils.getAnchorBoxesMap(cellWidth, cellHeight) | |
| local junctionAnchorBoxes = {} | |
| --[[ | |
| table.insert(junctionAnchorBoxes, {{cellWidth / 6 * 1, cellHeight / 6 * 1}, {cellWidth / 3, cellHeight / 3}}) | |
| table.insert(junctionAnchorBoxes, {{cellWidth / 6 * 3, cellHeight / 6 * 1}, {cellWidth / 3, cellHeight / 3}}) | |
| table.insert(junctionAnchorBoxes, {{cellWidth / 6 * 5, cellHeight / 6 * 1}, {cellWidth / 3, cellHeight / 3}}) | |
| table.insert(junctionAnchorBoxes, {{cellWidth / 6 * 1, cellHeight / 6 * 3}, {cellWidth / 3, cellHeight / 3}}) | |
| table.insert(junctionAnchorBoxes, {{cellWidth / 6 * 3, cellHeight / 6 * 3}, {cellWidth / 3, cellHeight / 3}}) | |
| table.insert(junctionAnchorBoxes, {{cellWidth / 6 * 5, cellHeight / 6 * 3}, {cellWidth / 3, cellHeight / 3}}) | |
| table.insert(junctionAnchorBoxes, {{cellWidth / 6 * 1, cellHeight / 6 * 5}, {cellWidth / 3, cellHeight / 3}}) | |
| table.insert(junctionAnchorBoxes, {{cellWidth / 6 * 3, cellHeight / 6 * 5}, {cellWidth / 3, cellHeight / 3}}) | |
| table.insert(junctionAnchorBoxes, {{cellWidth / 6 * 5, cellHeight / 6 * 5}, {cellWidth / 3, cellHeight / 3}}) | |
| ]]-- | |
| for i = 1, 13 do | |
| table.insert(junctionAnchorBoxes, {{cellWidth / 2, cellHeight / 2}, {cellWidth, cellHeight}}) | |
| end | |
| local doorAnchorBoxes = {} | |
| table.insert(doorAnchorBoxes, {{cellWidth / 8 * 1, cellHeight / 2}, {cellWidth / 4, cellHeight}}) | |
| table.insert(doorAnchorBoxes, {{cellWidth / 8 * 3, cellHeight / 2}, {cellWidth / 4, cellHeight}}) | |
| table.insert(doorAnchorBoxes, {{cellWidth / 8 * 5, cellHeight / 2}, {cellWidth / 4, cellHeight}}) | |
| table.insert(doorAnchorBoxes, {{cellWidth / 8 * 7, cellHeight / 2}, {cellWidth / 4, cellHeight}}) | |
| table.insert(doorAnchorBoxes, {{cellWidth / 2, cellHeight / 8 * 1}, {cellWidth, cellHeight / 4}}) | |
| table.insert(doorAnchorBoxes, {{cellWidth / 2, cellHeight / 8 * 3}, {cellWidth, cellHeight / 4}}) | |
| table.insert(doorAnchorBoxes, {{cellWidth / 2, cellHeight / 8 * 5}, {cellWidth, cellHeight / 4}}) | |
| table.insert(doorAnchorBoxes, {{cellWidth / 2, cellHeight / 8 * 7}, {cellWidth, cellHeight / 4}}) | |
| table.insert(doorAnchorBoxes, {{cellWidth / 2, cellHeight / 2}, {cellWidth, cellHeight}}) | |
| table.insert(doorAnchorBoxes, {{cellWidth / 4 * 1, cellHeight / 2}, {cellWidth / 2, cellHeight * 2}}) | |
| table.insert(doorAnchorBoxes, {{cellWidth / 4 * 3, cellHeight / 2}, {cellWidth / 2, cellHeight * 2}}) | |
| table.insert(doorAnchorBoxes, {{cellWidth / 2, cellHeight / 4 * 1}, {cellWidth * 2, cellHeight / 2}}) | |
| table.insert(doorAnchorBoxes, {{cellWidth / 2, cellHeight / 4 * 3}, {cellWidth * 2, cellHeight / 2}}) | |
| local iconAnchorBoxes = {} | |
| table.insert(iconAnchorBoxes, {{cellWidth / 2, cellHeight / 2}, {cellWidth, cellHeight}}) | |
| table.insert(iconAnchorBoxes, {{cellWidth / 2, cellHeight / 2}, {cellWidth * 2, cellHeight}}) | |
| table.insert(iconAnchorBoxes, {{cellWidth / 2, cellHeight / 2}, {cellWidth, cellHeight * 2}}) | |
| table.insert(iconAnchorBoxes, {{cellWidth / 2, cellHeight / 2}, {cellWidth / 2, cellHeight / 2}}) | |
| table.insert(iconAnchorBoxes, {{cellWidth / 2, cellHeight / 2}, {cellWidth / 2, cellHeight}}) | |
| table.insert(iconAnchorBoxes, {{cellWidth / 2, cellHeight / 2}, {cellWidth, cellHeight / 2}}) | |
| table.insert(iconAnchorBoxes, {{cellWidth / 2, cellHeight / 2}, {cellWidth / 4, cellHeight / 4}}) | |
| table.insert(iconAnchorBoxes, {{cellWidth / 2, cellHeight / 2}, {cellWidth / 4, cellHeight / 2}}) | |
| table.insert(iconAnchorBoxes, {{cellWidth / 2, cellHeight / 2}, {cellWidth / 2, cellHeight / 4}}) | |
| table.insert(iconAnchorBoxes, {{cellWidth / 2, cellHeight / 4 * 1}, {cellWidth, cellHeight / 2}}) | |
| table.insert(iconAnchorBoxes, {{cellWidth / 2, cellHeight / 4 * 3}, {cellWidth, cellHeight / 2}}) | |
| table.insert(iconAnchorBoxes, {{cellWidth / 4 * 1, cellHeight / 2}, {cellWidth / 2, cellHeight}}) | |
| table.insert(iconAnchorBoxes, {{cellWidth / 4 * 3, cellHeight / 2}, {cellWidth / 2, cellHeight}}) | |
| local labelAnchorBoxes = {} | |
| table.insert(labelAnchorBoxes, {{cellWidth / 6 * 1, cellHeight / 6 * 1}, {cellWidth / 3, cellHeight / 3}}) | |
| table.insert(labelAnchorBoxes, {{cellWidth / 6 * 3, cellHeight / 6 * 1}, {cellWidth / 3, cellHeight / 3}}) | |
| table.insert(labelAnchorBoxes, {{cellWidth / 6 * 5, cellHeight / 6 * 1}, {cellWidth / 3, cellHeight / 3}}) | |
| table.insert(labelAnchorBoxes, {{cellWidth / 6 * 1, cellHeight / 6 * 3}, {cellWidth / 3, cellHeight / 3}}) | |
| table.insert(labelAnchorBoxes, {{cellWidth / 6 * 3, cellHeight / 6 * 3}, {cellWidth / 3, cellHeight / 3}}) | |
| table.insert(labelAnchorBoxes, {{cellWidth / 6 * 5, cellHeight / 6 * 3}, {cellWidth / 3, cellHeight / 3}}) | |
| table.insert(labelAnchorBoxes, {{cellWidth / 6 * 1, cellHeight / 6 * 5}, {cellWidth / 3, cellHeight / 3}}) | |
| table.insert(labelAnchorBoxes, {{cellWidth / 6 * 3, cellHeight / 6 * 5}, {cellWidth / 3, cellHeight / 3}}) | |
| table.insert(labelAnchorBoxes, {{cellWidth / 6 * 5, cellHeight / 6 * 5}, {cellWidth / 3, cellHeight / 3}}) | |
| table.insert(labelAnchorBoxes, {{cellWidth / 4 * 1, cellHeight / 4 * 1}, {cellWidth / 2, cellHeight / 2}}) | |
| table.insert(labelAnchorBoxes, {{cellWidth / 4 * 1, cellHeight / 4 * 3}, {cellWidth / 2, cellHeight / 2}}) | |
| table.insert(labelAnchorBoxes, {{cellWidth / 4 * 3, cellHeight / 4 * 1}, {cellWidth / 2, cellHeight / 2}}) | |
| table.insert(labelAnchorBoxes, {{cellWidth / 4 * 3, cellHeight / 4 * 3}, {cellWidth / 2, cellHeight / 2}}) | |
| local anchorBoxesMap = {} | |
| anchorBoxesMap['points'] = junctionAnchorBoxes | |
| anchorBoxesMap['doors'] = doorAnchorBoxes | |
| anchorBoxesMap['icons'] = iconAnchorBoxes | |
| anchorBoxesMap['labels'] = labelAnchorBoxes | |
| return anchorBoxesMap | |
| end | |
| function utils.getColorMap() | |
| local colorMap = {} | |
| colorMap[1] = {255, 0, 0} | |
| colorMap[2] = {0, 255, 0} | |
| colorMap[3] = {0, 0, 255} | |
| colorMap[4] = {255, 255, 0} | |
| colorMap[5] = {0, 255, 255} | |
| colorMap[6] = {255, 0, 255} | |
| colorMap[7] = {128, 0, 0} | |
| colorMap[8] = {0, 128, 0} | |
| colorMap[9] = {0, 0, 128} | |
| colorMap[10] = {128, 128, 0} | |
| colorMap[11] = {0, 128, 128} | |
| colorMap[12] = {128, 0, 128} | |
| colorMap[13] = {128, 128, 128} | |
| return colorMap | |
| end | |
| function utils.findConnectedComponents(mask) | |
| local maskByte = (mask * 255):byte() | |
| local components = torch.IntTensor(maskByte:size()) | |
| local numComponents = cv.connectedComponents{maskByte, components} | |
| return components, numComponents | |
| end | |
| function utils.drawSegmentation(floorplanComponent, numComponents, denotedColorMap) | |
| local numComponents = numComponents or torch.max(floorplanComponent) | |
| local colorMap = denotedColorMap | |
| if colorMap == nil then | |
| colorMap = {} | |
| for i = 1, numComponents do | |
| colorMap[i] = torch.rand(3) | |
| end | |
| end | |
| colorMap[0] = torch.zeros(3) | |
| colorMap[-1] = torch.ones(3) | |
| local floorplanLabels = floorplanComponent:repeatTensor(3, 1, 1):double() | |
| for c = 1, 3 do | |
| floorplanLabels[c]:apply(function(x) return colorMap[x][c] end) | |
| end | |
| return floorplanLabels | |
| end | |
| function loadIcon(iconName, iconStyle, iconOrientation, extension) | |
| local icon = image.load('../icons/' .. iconName .. '_' .. iconStyle .. '.' .. extension, 1) | |
| assert(icon ~= nil) | |
| return rotateIcon(icon, iconOrientation) | |
| end | |
| function rotateIcon(iconOriginal, iconOrientation) | |
| local iconRotated | |
| if iconOrientation == 1 then | |
| iconRotated = iconOriginal:clone() | |
| elseif iconOrientation == 2 then | |
| if iconOriginal:dim() == 2 then | |
| iconRotated = image.hflip(iconOriginal:transpose(1, 2):contiguous()) | |
| else | |
| iconRotated = image.hflip(iconOriginal:transpose(2, 3):contiguous()) | |
| end | |
| elseif iconOrientation == 3 then | |
| iconRotated = image.vflip(image.hflip(iconOriginal)) | |
| else | |
| if iconOriginal:dim() == 2 then | |
| iconRotated = image.vflip(iconOriginal:transpose(1, 2):contiguous()) | |
| else | |
| iconRotated = image.vflip(iconOriginal:transpose(2, 3):contiguous()) | |
| end | |
| end | |
| return iconRotated | |
| end | |
| function utils.loadIconImages() | |
| if utils.iconImages == nil then | |
| utils.iconImages = {} | |
| local numberMap = utils.getNumberMap() | |
| for mode, map in pairs(numberMap) do | |
| local extension = 'jpg' | |
| if mode == 'labels' then | |
| extension = 'png' | |
| end | |
| for name, number in pairs(map) do | |
| utils.iconImages[name] = {} | |
| for style = 1, 13 do | |
| local iconExists, icon = pcall(function() | |
| return loadIcon(name, style, 1, extension) | |
| end) | |
| if iconExists == false then | |
| break | |
| end | |
| if icon:dim() == 3 then | |
| icon = icon[1] | |
| end | |
| utils.iconImages[name][style] = icon | |
| end | |
| end | |
| end | |
| end | |
| return utils.iconImages | |
| end | |
| function utils.lineDim(line, lineWidth) | |
| local lineWidth = lineWidth or 1 | |
| if math.abs(line[1][1] - line[2][1]) > math.abs(line[1][2] - line[2][2]) and math.abs(line[1][2] - line[2][2]) <= lineWidth then | |
| return 1 | |
| elseif math.abs(line[1][2] - line[2][2]) > math.abs(line[1][1] - line[2][1]) and math.abs(line[1][1] - line[2][1]) <= lineWidth then | |
| return 2 | |
| else | |
| return 0 | |
| end | |
| end | |
| function utils.cutoffRange(range, max) | |
| local lowerBound = math.max(range[1], 1) | |
| local upperBound = math.min(range[2], max) | |
| if lowerBound > upperBound then | |
| return {} | |
| else | |
| return {lowerBound, upperBound} | |
| end | |
| end | |
| function utils.drawLineMask(width, height, lines, lineWidth, indexed, lineExtentionLength, denotedLineDim) | |
| local lineWidth = lineWidth or 5 | |
| local indexed = indexed or false | |
| local lineExtentionLength = lineExtentionLength or lineWidth | |
| local lineMask = torch.zeros(height, width) | |
| local size = {width, height} | |
| local index = 1 | |
| for _, line in pairs(lines) do | |
| local lineDim = utils.lineDim(line) | |
| if (not denotedLineDim and lineDim > 0) or lineDim == denotedLineDim then | |
| local fixedDim = 3 - lineDim | |
| local fixedValue = (line[1][fixedDim] + line[2][fixedDim]) / 2 | |
| local fixedRange = utils.cutoffRange({fixedValue - lineWidth, fixedValue + lineWidth}, size[fixedDim]) | |
| if #fixedRange > 0 then | |
| local lineRange = utils.cutoffRange({line[1][lineDim] - lineExtentionLength, line[2][lineDim] + lineExtentionLength}, size[lineDim]) | |
| if lineDim == 1 then | |
| lineMask[{{fixedRange[1], fixedRange[2]}, {lineRange[1], lineRange[2]}}] = index | |
| else | |
| lineMask[{{lineRange[1], lineRange[2]}, {fixedRange[1], fixedRange[2]}}] = index | |
| end | |
| end | |
| elseif not denotedLineDim and lineDim == 0 then | |
| if math.abs(line[1][1] - line[2][1]) > math.abs(line[1][2] - line[2][2]) then | |
| lineDim = 1 | |
| else | |
| lineDim = 2 | |
| end | |
| for lineDimValue = math.min(line[1][lineDim], line[2][lineDim]), math.max(line[1][lineDim], line[2][lineDim]) do | |
| local orthogonalValue = line[1][3 - lineDim] + (lineDimValue - line[1][lineDim]) / (line[2][lineDim] - line[1][lineDim]) * (line[2][3 - lineDim] - line[1][3 - lineDim]) | |
| local orthogonalRange = utils.cutoffRange({orthogonalValue - lineWidth, orthogonalValue + lineWidth}, size[3 - lineDim]) | |
| if lineDim == 1 then | |
| lineMask[{{orthogonalRange[1], orthogonalRange[2]}, lineDimValue}] = index | |
| else | |
| lineMask[{lineDimValue, {orthogonalRange[1], orthogonalRange[2]}}] = index | |
| end | |
| end | |
| end | |
| if indexed then | |
| index = index + 1 | |
| end | |
| end | |
| return lineMask | |
| end | |
| function utils.sortLines(lines) | |
| for lineIndex, line in pairs(lines) do | |
| local lineDim = utils.lineDim(line) | |
| if lineDim > 0 and line[1][lineDim] > line[2][lineDim] then | |
| local temp = lines[lineIndex][1][lineDim] | |
| lines[lineIndex][1][lineDim] = lines[lineIndex][2][lineDim] | |
| lines[lineIndex][2][lineDim] = temp | |
| end | |
| end | |
| return lines | |
| end | |
| function utils.calcDistance(point_1, point_2) | |
| return math.sqrt(math.pow(point_1[1] - point_2[1], 2) + math.pow(point_1[2] - point_2[2], 2)) | |
| end | |
| function utils.findNearestJunctionPair(line_1, line_2, gap, styleSensitive) | |
| local nearestPair | |
| local minDistance | |
| for index_1 = 1, 2 do | |
| for index_2 = 1, 2 do | |
| local distance = utils.calcDistance(line_1[index_1], line_2[index_2]) | |
| if minDistance == nil or distance < minDistance then | |
| nearestPair = {index_1, index_2} | |
| minDistance = distance | |
| end | |
| end | |
| end | |
| if minDistance > gap then | |
| local lineDim_1 = utils.lineDim(line_1) | |
| local lineDim_2 = utils.lineDim(line_2) | |
| if lineDim_1 + lineDim_2 == 3 then | |
| local fixedValue_1 = (line_1[1][3 - lineDim_1] + line_1[2][3 - lineDim_1]) / 2 | |
| local fixedValue_2 = (line_2[1][3 - lineDim_2] + line_2[2][3 - lineDim_2]) / 2 | |
| if line_2[1][lineDim_2] < fixedValue_1 and line_2[2][lineDim_2] > fixedValue_1 then | |
| for index = 1, 2 do | |
| local distance = math.abs(line_1[index][lineDim_1] - fixedValue_2) | |
| if distance < minDistance then | |
| nearestPair = {index, 0} | |
| minDistance = distance | |
| end | |
| end | |
| end | |
| if line_1[1][lineDim_1] < fixedValue_2 and line_1[2][lineDim_1] > fixedValue_2 then | |
| for index = 1, 2 do | |
| local distance = math.abs(line_2[index][lineDim_2] - fixedValue_1) | |
| if distance < minDistance then | |
| nearestPair = {0, index} | |
| minDistance = distance | |
| end | |
| end | |
| end | |
| end | |
| end | |
| if styleSensitive ~= false then | |
| if #line_1 >= 3 and line_1[3][2] == 2 and line_2[3][2] == 1 then | |
| minDistance = gap + 1 | |
| nearestPair[2] = 0 | |
| end | |
| if #line_2 >= 3 and line_2[3][2] == 2 and line_1[3][2] == 1 then | |
| minDistance = gap + 1 | |
| nearestPair[1] = 0 | |
| end | |
| end | |
| return nearestPair, minDistance | |
| end | |
| function utils.stitchLines(lines, gap) | |
| for lineIndex_1, line_1 in pairs(lines) do | |
| local lineDim_1 = utils.lineDim(line_1) | |
| if lineDim_1 > 0 then | |
| local fixedValue_1 = (line_1[1][3 - lineDim_1] + line_1[2][3 - lineDim_1]) / 2 | |
| for lineIndex_2, line_2 in pairs(lines) do | |
| if lineIndex_2 > lineIndex_1 then | |
| local lineDim_2 = utils.lineDim(line_2) | |
| if lineDim_2 > 0 then | |
| local fixedValue_2 = (line_2[1][3 - lineDim_2] + line_2[2][3 - lineDim_2]) / 2 | |
| local nearestPair, minDistance = utils.findNearestJunctionPair(line_1, line_2, gap) | |
| --print(minDistance .. ' ' .. lineDim_1 .. ' ' .. lineDim_2) | |
| if minDistance <= gap and lineDim_1 + lineDim_2 == 3 then | |
| local pointIndex_1 = nearestPair[1] | |
| local pointIndex_2 = nearestPair[2] | |
| --print(lineIndex_1 .. ' ' .. lineIndex_2) | |
| if pointIndex_1 > 0 and pointIndex_2 > 0 then | |
| lines[lineIndex_1][pointIndex_1][lineDim_1] = fixedValue_2 | |
| lines[lineIndex_2][pointIndex_2][lineDim_2] = fixedValue_1 | |
| elseif pointIndex_1 > 0 and pointIndex_2 == 0 then | |
| lines[lineIndex_1][pointIndex_1][lineDim_1] = fixedValue_2 | |
| elseif pointIndex_1 == 0 and pointIndex_2 > 0 then | |
| lines[lineIndex_2][pointIndex_2][lineDim_2] = fixedValue_1 | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| return lines | |
| end | |
| function utils.mergeLines(lines, gap) | |
| local mergedLines = {} | |
| local mergedLineMap = {} | |
| while true do | |
| local hasChange = false | |
| for lineIndex_1, line_1 in pairs(lines) do | |
| if not mergedLines[lineIndex_1] and line_1[3][2] == 2 then | |
| local lineDim_1 = utils.lineDim(line_1) | |
| if lineDim_1 > 0 then | |
| local fixedValue_1 = (line_1[1][3 - lineDim_1] + line_1[2][3 - lineDim_1]) / 2 | |
| for lineIndex_2, line_2 in pairs(lines) do | |
| if lineIndex_2 ~= lineIndex_1 then | |
| local lineDim_2 = utils.lineDim(line_2) | |
| if lineDim_2 == lineDim_1 and line_2[3][2] == 1 then | |
| local nearestPair, minDistance = utils.findNearestJunctionPair(line_1, line_2, gap, false) | |
| if minDistance <= gap then | |
| local fixedValue_2 = (line_2[1][3 - lineDim_2] + line_2[2][3 - lineDim_2]) / 2 | |
| local pointIndex_1 = nearestPair[1] | |
| local pointIndex_2 = nearestPair[2] | |
| --print(lineIndex_1 .. ' ' .. lineIndex_2) | |
| if pointIndex_1 > 0 and pointIndex_2 > 0 then | |
| --local fixedValue = (fixedValue_1 + fixedValue_2) / 2 | |
| local fixedValue = fixedValue_2 | |
| for pointIndex = 1, 2 do | |
| --line_1[pointIndex][lineDim_1] = fixedValue | |
| line_1[pointIndex][3 - lineDim_1] = fixedValue | |
| --line_1[3][2] = 1 | |
| hasChange = true | |
| mergedLines[lineIndex_1] = true | |
| line_2[pointIndex_2][lineDim_2] = line_1[3 - pointIndex_1][lineDim_1] | |
| mergedLineMap[lineIndex_1] = lineIndex_2 | |
| end | |
| end | |
| end | |
| elseif lineDim_2 + lineDim_1 == 3 and line_2[3][2] == 2 then | |
| local nearestPair, minDistance = utils.findNearestJunctionPair(line_1, line_2, gap, false) | |
| if minDistance <= gap then | |
| local pointIndex_1 = nearestPair[1] | |
| local pointIndex_2 = nearestPair[2] | |
| --print(lineIndex_1 .. ' ' .. lineIndex_2) | |
| if pointIndex_1 > 0 and pointIndex_2 > 0 then | |
| local fixedValue_2 = (line_2[1][3 - lineDim_2] + line_2[2][3 - lineDim_2]) / 2 | |
| --local fixedValue = (fixedValue_1 + fixedValue_2) / 2 | |
| --local fixedValue = fixedValue_1 | |
| line_1[pointIndex_1][lineDim_1] = fixedValue_2 | |
| --line_2[pointIndex_2][lineDim_2] = fixedValue_1 | |
| mergedLines[lineIndex_1] = true | |
| --line_2[3][2] = 1 | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| if hasChange == false then | |
| break | |
| end | |
| end | |
| local newLines = {} | |
| for lineIndex, line in pairs(lines) do | |
| line[3][2] = 1 | |
| if not mergedLineMap[lineIndex] then | |
| table.insert(newLines, line) | |
| end | |
| end | |
| return newLines | |
| end | |
| function utils.fixedDoors(lines, walls, gap) | |
| for lineIndex_1, line_1 in pairs(lines) do | |
| local lineDim = utils.lineDim(line_1) | |
| if lineDim > 0 then | |
| local fixedValue_1 = (line_1[1][3 - lineDim] + line_1[2][3 - lineDim]) / 2 | |
| local doorFixed = false | |
| for lineIndex_2, line_2 in pairs(walls) do | |
| if utils.lineDim(line_2) == lineDim then | |
| local fixedValue_2 = (line_2[1][3 - lineDim] + line_2[2][3 - lineDim]) / 2 | |
| if math.abs(fixedValue_2 - fixedValue_1) <= gap and math.min(line_1[1][lineDim], line_1[2][lineDim]) >= math.min(line_2[1][lineDim], line_2[2][lineDim]) - gap and math.max(line_1[1][lineDim], line_1[2][lineDim]) <= math.max(line_2[1][lineDim], line_2[2][lineDim]) + gap then | |
| for c = 1, 2 do | |
| line_1[c][3 - lineDim] = fixedValue_2 | |
| end | |
| doorFixed = true | |
| end | |
| end | |
| end | |
| if doorFixed == false then | |
| print('door: ' .. lineIndex_1 .. ' not fixed') | |
| end | |
| end | |
| end | |
| return lines | |
| end | |
| function utils.gridPoint(width, height, gridWidth, gridHeight, point) | |
| local cellWidth = width / gridWidth | |
| local cellHeight = height / gridHeight | |
| local gridX = math.floor((point[1] - 1) / cellWidth) + 1 | |
| local gridY = math.floor((point[2] - 1) / cellHeight) + 1 | |
| local cellX = ((point[1] - 1) - (gridX - 1) * cellWidth) / (cellWidth - 1) | |
| local cellY = ((point[2] - 1) - (gridY - 1) * cellHeight) / (cellHeight - 1) | |
| return {{gridX, gridY}, {cellX, cellY}} | |
| end | |
| function utils.gridRectangle(width, height, gridWidth, gridHeight, rectangle) | |
| --[[ | |
| if point_2 == nil then | |
| local rectangle = utils.gridPoint(width, height, gridWidth, gridHeight, point_1) | |
| table.insert(rectangle, {0, 0}) | |
| return rectangle | |
| end | |
| ]]-- | |
| local point_1 = rectangle[1] | |
| local point_2 = rectangle[2] | |
| local center = utils.gridPoint(width, height, gridWidth, gridHeight, {(point_1[1] + point_2[1]) / 2, (point_1[2] + point_2[2]) / 2}) | |
| local rectangle = center | |
| table.insert(rectangle, {(point_2[1] - point_1[1] + 1) / width, (point_2[2] - point_1[2] + 1) / height}) | |
| return rectangle | |
| end | |
| function utils.imageRectangle(width, height, gridWidth, gridHeight, rectangle) | |
| local cellWidth = width / gridWidth | |
| local cellHeight = height / gridHeight | |
| local centerX = (rectangle[1][1] - 1) * cellWidth + rectangle[2][1] * (cellWidth - 1) + 1 | |
| local centerY = (rectangle[1][2] - 1) * cellHeight + rectangle[2][2] * (cellHeight - 1) + 1 | |
| return {{math.max(torch.round(centerX - (rectangle[3][1] * (width - 1)) / 2), 1), math.max(torch.round(centerY - (rectangle[3][2] * (height - 1)) / 2), 1)}, {math.min(torch.round(centerX + (rectangle[3][1] * (width - 1)) / 2), width), math.min(torch.round(centerY + (rectangle[3][2] * (height - 1)) / 2), height)}} | |
| end | |
| function utils.pointsToLines(width, height, points, lineWidth, divideLines, minLineLength, ordered, pointsDuplicated) | |
| local lineWidth = lineWidth or 5 | |
| --local minLineLength = minLineLength or math.max(width, height) | |
| local minLineLength = minLineLength or lineWidth * 2 | |
| local divideLines = divideLines or false | |
| local usedPointLineMask = {} | |
| local pointOrientations = {} | |
| for pointIndex, point in pairs(points) do | |
| usedPointLineMask[pointIndex] = {true, true, true, true} | |
| local orientations = {} | |
| local orientation = point[3][3] | |
| if point[3][2] == 1 then | |
| table.insert(orientations, (orientation + 2 - 1) % 4 + 1) | |
| elseif point[3][2] == 2 then | |
| table.insert(orientations, orientation) | |
| table.insert(orientations, (orientation + 3 - 1) % 4 + 1) | |
| elseif point[3][2] == 3 then | |
| if divideLines then | |
| for i = 1, 4 do | |
| if i ~= orientation then | |
| table.insert(orientations, i) | |
| end | |
| end | |
| else | |
| table.insert(orientations, (orientation + 2 - 1) % 4 + 1) | |
| end | |
| else | |
| if divideLines then | |
| for i = 1, 4 do | |
| table.insert(orientations, i) | |
| end | |
| end | |
| end | |
| pointOrientations[pointIndex] = orientations | |
| for _, orientation in pairs(orientations) do | |
| usedPointLineMask[pointIndex][orientation] = false | |
| end | |
| end | |
| local lines = {} | |
| local lineJunctionsMap = {} | |
| for pointIndex, point in pairs(points) do | |
| local orientations = pointOrientations[pointIndex] | |
| local x = point[1][1] | |
| local y = point[1][2] | |
| for _, orientation in pairs(orientations) do | |
| if usedPointLineMask[pointIndex][orientation] == false then | |
| local lineDim | |
| local fixedValue | |
| local junction | |
| local startPoint | |
| if orientation == 1 or orientation == 3 then | |
| lineDim = 2 | |
| fixedValue = x | |
| startPoint = y | |
| if orientation == 1 then | |
| junction = 1 | |
| else | |
| junction = height | |
| end | |
| else | |
| lineDim = 1 | |
| fixedValue = y | |
| startPoint = x | |
| if orientation == 4 then | |
| junction = 1 | |
| else | |
| junction = width | |
| end | |
| end | |
| local orientationOpposite = (orientation + 2 - 1) % 4 + 1 | |
| local selectedOtherPointIndex | |
| for otherPointIndex, otherPoint in pairs(points) do | |
| if otherPointIndex ~= pointIndex and usedPointLineMask[otherPointIndex][orientationOpposite] == false and (not ordered or otherPointIndex == pointIndex + 1) then | |
| local otherXY = otherPoint[1] | |
| if otherXY[lineDim] > math.min(junction, startPoint) and otherXY[lineDim] < math.max(junction, startPoint) and otherXY[3 - lineDim] >= fixedValue - lineWidth and otherXY[3 - lineDim] <= fixedValue + lineWidth then | |
| junction = otherXY[lineDim] | |
| selectedOtherPointIndex = otherPointIndex | |
| end | |
| end | |
| end | |
| local point_1 = {} | |
| point_1[lineDim] = math.min(startPoint, junction) | |
| point_1[3 - lineDim] = fixedValue | |
| local point_2 = {} | |
| point_2[lineDim] = math.max(startPoint, junction) | |
| point_2[3 - lineDim] = fixedValue | |
| if point_1[1] ~= 1 and point_1[1] ~= width and point_1[2] ~= 1 and point_1[2] ~= height and point_2[1] ~= 1 and point_2[1] ~= width and point_2[2] ~= 1 and point_2[2] ~= height then | |
| if utils.calcDistance(point_1, point_2) > minLineLength then | |
| table.insert(lines, {point_1, point_2, {"wall", 1, 1}}) | |
| if startPoint < junction then | |
| table.insert(lineJunctionsMap, {pointIndex, selectedOtherPointIndex}) | |
| else | |
| table.insert(lineJunctionsMap, {selectedOtherPointIndex, pointIndex}) | |
| end | |
| if not pointsDuplicated then | |
| if selectedOtherPointIndex ~= nil then | |
| usedPointLineMask[selectedOtherPointIndex][orientationOpposite] = true | |
| end | |
| usedPointLineMask[pointIndex][orientation] = true | |
| end | |
| --print(pointIndex .. ' ' .. orientation .. ' ' .. selectedOtherPointIndex) | |
| end | |
| end | |
| end | |
| end | |
| end | |
| if pointsDuplicated then | |
| local newLines = {} | |
| local newLineJunctionsMap = {} | |
| local pointNeighborMap = {} | |
| for _, line in pairs(lines) do | |
| local junctions = lineJunctionsMap[_] | |
| if not pointNeighborMap[junctions[1]] then | |
| pointNeighborMap[junctions[1]] = {} | |
| end | |
| if not pointNeighborMap[junctions[1]][junctions[2]] then | |
| table.insert(newLines, line) | |
| table.insert(newLineJunctionsMap, junctions) | |
| pointNeighborMap[junctions[1]][junctions[2]] = true | |
| end | |
| end | |
| lines = newLines | |
| lineJunctionsMap = newLineJunctionsMap | |
| end | |
| return lines, lineJunctionsMap | |
| end | |
| function utils.linesToPoints(width, height, lines, lineWidth) | |
| local lineWidth = lineWidth or 5 | |
| local points = {} | |
| local usedLinePointMask = {} | |
| for lineIndex, line in pairs(lines) do | |
| usedLinePointMask[lineIndex] = {false, false} | |
| end | |
| for lineIndex_1, line_1 in pairs(lines) do | |
| local lineDim_1 = utils.lineDim(line_1) | |
| if lineDim_1 > 0 then | |
| local fixedValue_1 = (line_1[1][3 - lineDim_1] + line_1[2][3 - lineDim_1]) / 2 | |
| for lineIndex_2, line_2 in pairs(lines) do | |
| if lineIndex_2 > lineIndex_1 then | |
| local lineDim_2 = utils.lineDim(line_2) | |
| if lineDim_2 + lineDim_1 == 3 then | |
| local fixedValue_2 = (line_2[1][3 - lineDim_2] + line_2[2][3 - lineDim_2]) / 2 | |
| local nearestPair, minDistance = utils.findNearestJunctionPair(line_1, line_2, lineWidth) | |
| if minDistance <= lineWidth then | |
| local pointIndex_1 = nearestPair[1] | |
| local pointIndex_2 = nearestPair[2] | |
| if pointIndex_1 > 0 and pointIndex_2 > 0 then | |
| local point = {} | |
| point[lineDim_1] = fixedValue_2 | |
| point[lineDim_2] = fixedValue_1 | |
| local side = {} | |
| side[lineDim_1] = line_1[3 - pointIndex_1][lineDim_1] - fixedValue_2 | |
| side[lineDim_2] = line_2[3 - pointIndex_2][lineDim_2] - fixedValue_1 | |
| if side[1] < 0 and side[2] < 0 then | |
| table.insert(points, {point, point, {'point', 2, 1}}) | |
| elseif side[1] > 0 and side[2] < 0 then | |
| table.insert(points, {point, point, {'point', 2, 2}}) | |
| elseif side[1] > 0 and side[2] > 0 then | |
| table.insert(points, {point, point, {'point', 2, 3}}) | |
| elseif side[1] < 0 and side[2] > 0 then | |
| table.insert(points, {point, point, {'point', 2, 4}}) | |
| end | |
| usedLinePointMask[lineIndex_1][pointIndex_1] = true | |
| usedLinePointMask[lineIndex_2][pointIndex_2] = true | |
| elseif (pointIndex_1 > 0 and pointIndex_2 == 0) or (pointIndex_1 == 0 and pointIndex_2 > 0) then | |
| local lineDim | |
| local pointIndex | |
| local fixedValue | |
| local pointValue | |
| if pointIndex_1 > 0 then | |
| lineDim = lineDim_1 | |
| pointIndex = pointIndex_1 | |
| fixedValue = fixedValue_2 | |
| pointValue = line_1[pointIndex_1][3 - lineDim_1] | |
| usedLinePointMask[lineIndex_1][pointIndex_1] = true | |
| else | |
| lineDim = lineDim_2 | |
| pointIndex = pointIndex_2 | |
| fixedValue = fixedValue_1 | |
| pointValue = line_2[pointIndex_2][3 - lineDim_2] | |
| usedLinePointMask[lineIndex_2][pointIndex_2] = true | |
| end | |
| local point = {} | |
| point[lineDim] = fixedValue | |
| point[3 - lineDim] = pointValue | |
| if pointIndex == 1 then | |
| if lineDim == 1 then | |
| table.insert(points, {point, point, {'point', 3, 4}}) | |
| else | |
| table.insert(points, {point, point, {'point', 3, 1}}) | |
| end | |
| else | |
| if lineDim == 1 then | |
| table.insert(points, {point, point, {'point', 3, 2}}) | |
| else | |
| table.insert(points, {point, point, {'point', 3, 3}}) | |
| end | |
| end | |
| end | |
| elseif line_1[1][lineDim_1] < fixedValue_2 and line_1[2][lineDim_1] > fixedValue_2 and line_2[1][lineDim_2] < fixedValue_1 and line_2[2][lineDim_2] > fixedValue_1 then | |
| local point = {} | |
| point[lineDim_1] = fixedValue_2 | |
| point[lineDim_2] = fixedValue_1 | |
| table.insert(points, {point, point, {'point', 4, 1}}) | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| for lineIndex, pointMask in pairs(usedLinePointMask) do | |
| local lineDim = utils.lineDim(lines[lineIndex]) | |
| for pointIndex = 1, 2 do | |
| if pointMask[pointIndex] == false then | |
| local point = {lines[lineIndex][pointIndex][1], lines[lineIndex][pointIndex][2]} | |
| if pointIndex == 1 then | |
| if lineDim == 1 then | |
| table.insert(points, {point, point, {'point', 1, 4}}) | |
| elseif lineDim == 2 then | |
| table.insert(points, {point, point, {'point', 1, 1}}) | |
| end | |
| else | |
| if lineDim == 1 then | |
| table.insert(points, {point, point, {'point', 1, 2}}) | |
| elseif lineDim == 2 then | |
| table.insert(points, {point, point, {'point', 1, 3}}) | |
| end | |
| end | |
| end | |
| end | |
| end | |
| return points | |
| end | |
| function utils.convertRepresentationToGeneral(width, height, representation, representationType, lineWidth) | |
| local representationGeneral = {} | |
| if representationType == 'P' then | |
| representationGeneral.walls = utils.pointsToLines(width, height, representation.points, lineWidth) | |
| local lineMask = torch.zeros(height, width) | |
| for _, line in pairs(representationGeneral.walls) do | |
| local lineDim = utils.lineDim(line) | |
| if lineDim > 0 then | |
| local maxSize = {width, height} | |
| local rectangle = {{}, {}} | |
| rectangle[1][lineDim] = math.min(line[1][lineDim], line[2][lineDim]) | |
| rectangle[2][lineDim] = math.max(line[1][lineDim], line[2][lineDim]) | |
| rectangle[1][3 - lineDim] = math.max((line[1][3 - lineDim] + line[2][3 - lineDim]) / 2 - lineWidth, 1) | |
| rectangle[2][3 - lineDim] = math.min((line[1][3 - lineDim] + line[2][3 - lineDim]) / 2 + lineWidth, maxSize[3 - lineDim]) | |
| --[[ | |
| for pointIndex = 1, 2 do | |
| for c = 1, 2 do | |
| print(rectangle[pointIndex][c]) | |
| end | |
| end | |
| print(#lineMask) | |
| ]]-- | |
| --representationImage[{{}, {rectangle[1][2], rectangle[2][2]}, {rectangle[1][1], rectangle[2][1]}}] = 0 | |
| lineMask[{{rectangle[1][2], rectangle[2][2]}, {rectangle[1][1], rectangle[2][1]}}] = 1 | |
| --success, newRepresentationImage = pcall(function() return image.drawRect(representationImage, rectangle[1][1], rectangle[1][2], rectangle[2][1], rectangle[2][2], {lineWidth = lineWidth, color = {0, 0, 0}}) end) | |
| --if success then | |
| --representationImage = newRepresentationImage | |
| --end | |
| end | |
| --representationImage[{{}, {minX, maxX}, {minY, maxY}}] = 0 | |
| end | |
| for mode, items in pairs(representation) do | |
| if mode == 'doors' or mode == 'icons' or mode == 'labels' then | |
| for __, item in pairs(items) do | |
| --print(mode) | |
| --print(item) | |
| --local iconImage = iconImages[item[3][1]][item[3][2]] | |
| if mode == 'icons' then | |
| local orientation = 1 | |
| local rectangle = {item[1], item[2]} | |
| if (math.abs(rectangle[2][1] - rectangle[1][1]) >= math.abs(rectangle[2][2] - rectangle[1][2]) and (item[3][1] ~= 'toilet' and item[3][1] ~= 'stairs')) or (math.abs(rectangle[2][1] - rectangle[1][1]) <= math.abs(rectangle[2][2] - rectangle[1][2]) and (item[3][1] == 'toilet' or item[3][1] == 'stairs')) then | |
| local min = math.min(rectangle[1][2], rectangle[2][2]) | |
| local max = math.max(rectangle[1][2], rectangle[2][2]) | |
| local center = (rectangle[1][1] + rectangle[2][1]) / 2 | |
| local deltaMin = 0 | |
| for delta = 1, min - 1 do | |
| deltaMin = delta | |
| if lineMask[min - delta][center] == 1 then | |
| break | |
| end | |
| end | |
| local deltaMax = 0 | |
| for delta = 1, height - max do | |
| deltaMax = delta | |
| if lineMask[max + delta][center] == 1 then | |
| break | |
| end | |
| end | |
| if deltaMin > deltaMax then | |
| orientation = 1 | |
| else | |
| orientation = 3 | |
| end | |
| else | |
| local min = math.min(rectangle[1][1], rectangle[2][1]) | |
| local max = math.max(rectangle[1][1], rectangle[2][1]) | |
| local center = (rectangle[1][2] + rectangle[2][2]) / 2 | |
| local deltaMin = 0 | |
| for delta = 1, min - 1 do | |
| deltaMin = delta | |
| if lineMask[center][min - delta] == 1 then | |
| break | |
| end | |
| end | |
| local deltaMax = 0 | |
| for delta = 1, width - max do | |
| deltaMax = delta | |
| if lineMask[center][max + delta] == 1 then | |
| break | |
| end | |
| end | |
| if deltaMin > deltaMax then | |
| orientation = 4 | |
| else | |
| orientation = 2 | |
| end | |
| end | |
| --[[ | |
| print(item[3][1]) | |
| print(math.abs(rectangle[2][1] - rectangle[1][1])) | |
| print(math.abs(rectangle[2][2] - rectangle[1][2])) | |
| print(orientation) | |
| ]]-- | |
| item[3][3] = orientation | |
| elseif mode == 'doors' then | |
| local lineDim | |
| if math.abs(item[2][1] - item[1][1]) > math.abs(item[2][2] - item[1][2]) then | |
| item[3][3] = 1 | |
| lineDim = 1 | |
| else | |
| item[3][3] = 2 | |
| lineDim = 2 | |
| end | |
| local fixedValue = (item[1][3 - lineDim] + item[2][3 - lineDim]) / 2 | |
| item[1][3 - lineDim] = fixedValue | |
| item[2][3 - lineDim] = fixedValue | |
| else | |
| item[3][3] = 1 | |
| end | |
| if mode ~= 'doors' then | |
| item[2][1] = math.max(item[2][1] - item[1][1], 30) + item[1][1] | |
| item[2][2] = math.max(item[2][2] - item[1][2], 30) + item[1][2] | |
| end | |
| --[[ | |
| iconImage = rotateIcon(iconImage, orientation) | |
| local icon = image.scale(iconImage, rectangle[2][1] - rectangle[1][1] + 1, rectangle[2][2] - rectangle[1][2] + 1) | |
| if icon:dim() == 2 then | |
| icon = icon:repeatTensor(3, 1, 1) | |
| end | |
| ]]-- | |
| --print(rectangle) | |
| --print(#icon) | |
| --print(#representationImage[{{}, {rectangle[1][2], rectangle[2][2]}, {rectangle[1][1], rectangle[2][1]}}]) | |
| --representationImage[{{}, {rectangle[1][2], rectangle[2][2]}, {rectangle[1][1], rectangle[2][1]}}] = icon | |
| end | |
| representationGeneral[mode] = items | |
| end | |
| end | |
| end | |
| return representationGeneral | |
| end | |
| function utils.convertRepresentation(width, height, representationGeneral, representationType, lineWidth) | |
| local representation = {} | |
| if representationType == 'P' then | |
| representationGeneral.walls = utils.mergeLines(representationGeneral.walls, lineWidth) | |
| representation.points = utils.linesToPoints(width, height, representationGeneral.walls, lineWidth) | |
| representation.doors = representationGeneral.doors | |
| representation.icons = representationGeneral.icons | |
| representation.labels = representationGeneral.labels | |
| end | |
| return representation | |
| end | |
| function utils.globalToGrid(width, height, gridWidth, gridHeight, representationGlobal) | |
| local representation = {} | |
| for mode, items in pairs(representationGlobal) do | |
| representation[mode] = {} | |
| for _, item in pairs(items) do | |
| local newRectangle = utils.gridRectangle(width, height, gridWidth, gridHeight, item) | |
| table.insert(newRectangle, item[3]) | |
| if newRectangle[1][1] >= 1 and newRectangle[1][1] <= gridWidth and newRectangle[1][2] >= 1 and newRectangle[1][2] <= gridHeight then | |
| table.insert(representation[mode], newRectangle) | |
| end | |
| end | |
| end | |
| return representation | |
| end | |
| function utils.gridToGlobal(width, height, gridWidth, gridHeight, representation) | |
| local representationGlobal = {} | |
| for mode, items in pairs(representation) do | |
| representationGlobal[mode] = {} | |
| for _, item in pairs(items) do | |
| local newRectangle = utils.imageRectangle(width, height, gridWidth, gridHeight, item) | |
| table.insert(newRectangle, item[4]) | |
| --print(mode) | |
| --print(item) | |
| --print(newRectangle) | |
| table.insert(representationGlobal[mode], newRectangle) | |
| end | |
| end | |
| --os.exit(1) | |
| return representationGlobal | |
| end | |
| function utils.convertRepresentationToTensor(width, height, gridWidth, gridHeight, representationGlobal) | |
| local representation = utils.globalToGrid(width, height, gridWidth, gridHeight, representationGlobal) | |
| local offsetsBB = utils.offsetsBB() | |
| local offsetsClass = utils.offsetsClass() | |
| local numFeaturesBB = utils.numFeaturesBB() | |
| local numFeaturesClass = utils.numFeaturesClass() | |
| local numItemsPerCell = utils.numItemsPerCell() | |
| local numFeaturesPerItem = utils.numFeaturesPerItem() | |
| local representationTensor = torch.zeros(numFeaturesBB + numFeaturesClass, gridHeight, gridWidth) | |
| local gridItems = {} | |
| for y = 1, gridHeight do | |
| gridItems[y] = {} | |
| for x = 1, gridWidth do | |
| gridItems[y][x] = {} | |
| for mode, offset in pairs(offsetsBB) do | |
| gridItems[y][x][mode] = {} | |
| end | |
| end | |
| end | |
| for mode, items in pairs(representation) do | |
| for __, item in pairs(items) do | |
| --if gridItems[item[1][2]][item[1][1]][mode] ~= nil then | |
| table.insert(gridItems[item[1][2]][item[1][1]][mode], item) | |
| --end | |
| end | |
| end | |
| for x = 1, gridWidth do | |
| for y = 1, gridHeight do | |
| for mode, items in pairs(gridItems[y][x]) do | |
| local numFeatures = numFeaturesPerItem[mode] | |
| local sortedItems = {} | |
| for _, item in pairs(items) do | |
| table.insert(sortedItems, {item[2][1] + item[2][2], item}) | |
| end | |
| table.sort(sortedItems, function(a, b) return a[1] < b[1] end) | |
| for i = numItemsPerCell[mode] + 1, #sortedItems do | |
| table.remove(sortedItems) | |
| end | |
| local offsetBB = offsetsBB[mode] | |
| for itemIndex, item in pairs(sortedItems) do | |
| item = item[2] | |
| local itemTensor = representationTensor[{{offsetBB + (itemIndex - 1) * numFeatures + 1, offsetBB + (itemIndex - 1) * numFeatures + numFeatures}, y, x}] | |
| itemTensor[1] = item[2][1] | |
| itemTensor[2] = item[2][2] | |
| itemTensor[3] = item[3][1] | |
| itemTensor[4] = item[3][2] | |
| itemTensor[5] = 1 | |
| end | |
| local offsetClass = offsetsClass[mode] | |
| for itemIndex, item in pairs(sortedItems) do | |
| item = item[2] | |
| local class = utils.getNumber(mode, item[4]) | |
| representationTensor[{offsetClass + itemIndex, y, x}] = class | |
| end | |
| end | |
| end | |
| end | |
| return representationTensor | |
| end | |
| function utils.convertRepresentationToTensorAnchor(width, height, gridWidth, gridHeight, representationGlobal, confidenceThreshold) | |
| local confidenceThreshold = confidenceThreshold or 0.7 | |
| local cellWidth = width / gridWidth | |
| local cellHeight = height / gridHeight | |
| local anchorBoxesMap = utils.getAnchorBoxesMap(cellWidth, cellHeight) | |
| local anchorBoxesTensors = {} | |
| for mode, items in pairs(representationGlobal) do | |
| local anchorBoxes = anchorBoxesMap[mode] | |
| local numAnchorBoxes = #anchorBoxes | |
| local anchorBoxesTensor = torch.zeros(gridHeight, gridWidth, numAnchorBoxes, 6) | |
| for __, item in pairs(items) do | |
| local point_1 = item[1] | |
| local point_2 = item[2] | |
| local label = utils.getNumber(mode, item[3]) | |
| local center = {(point_1[1] + point_2[1]) / 2, (point_1[2] + point_2[2]) / 2} | |
| local gridX = math.floor((center[1] - 1) / cellWidth) + 1 | |
| local gridY = math.floor((center[2] - 1) / cellHeight) + 1 | |
| if gridX >= 1 and gridX <= gridWidth and gridY >= 1 and gridY <= gridHeight then | |
| local cellX = (center[1] - 1) - (gridX - 1) * cellWidth | |
| local cellY = (center[2] - 1) - (gridY - 1) * cellHeight | |
| local gridBox = {{cellX, cellY}, {math.abs(point_2[1] - point_1[1]), math.abs(point_2[2] - point_1[2])}} | |
| if mode == 'labels' then | |
| gridBox[2] = {0, 0} | |
| end | |
| for c = 1, 2 do | |
| gridBox[2][c] = math.max(gridBox[2][c], 1) | |
| end | |
| local gridAnchorBoxesTensor = torch.zeros(numAnchorBoxes, 5) | |
| for anchorBoxIndex, anchorBox in pairs(anchorBoxes) do | |
| if mode ~= 'points' or anchorBoxIndex == label then | |
| local IOU = utils.calcIOU({{anchorBox[1][1] - anchorBox[2][1] / 2, anchorBox[1][2] - anchorBox[2][2] / 2}, {anchorBox[1][1] + anchorBox[2][1] / 2, anchorBox[1][2] + anchorBox[2][2] / 2}}, {{gridBox[1][1] - gridBox[2][1] / 2, gridBox[1][2] - gridBox[2][2] / 2}, {gridBox[1][1] + gridBox[2][1] / 2, gridBox[1][2] + gridBox[2][2] / 2}}) | |
| --[[ | |
| print(IOU) | |
| print(gridY) | |
| for index = 1, 2 do | |
| for c = 1, 2 do | |
| print(gridBox[index][c]) | |
| end | |
| end | |
| for index = 1, 2 do | |
| for c = 1, 2 do | |
| print(anchorBox[index][c]) | |
| end | |
| end]]-- | |
| --assert(IOU >= 0 and IOU <= 1) | |
| if IOU >= 0 and IOU <= 1 then | |
| gridAnchorBoxesTensor[anchorBoxIndex][5] = IOU | |
| for c = 1, 2 do | |
| gridAnchorBoxesTensor[anchorBoxIndex][c] = (gridBox[1][c] - anchorBox[1][c]) / anchorBox[2][c] | |
| gridAnchorBoxesTensor[anchorBoxIndex][c + 2] = math.log(gridBox[2][c] / anchorBox[2][c]) | |
| end | |
| end | |
| end | |
| end | |
| --[[ | |
| print('done') | |
| print(gridAnchorBoxesTensor) | |
| print(gridBox) | |
| print(anchorBox) | |
| os.exit(1) | |
| ]]-- | |
| local maxIOU | |
| local maxIOUIndex | |
| for anchorBoxIndex = 1, gridAnchorBoxesTensor:size(1) do | |
| local IOU = gridAnchorBoxesTensor[anchorBoxIndex][5] | |
| if gridY <= 0 or gridY > 8 or gridX <= 0 or gridX > 8 or anchorBoxIndex <= 0 or anchorBoxIndex > numAnchorBoxes then | |
| print(gridY .. ' ' .. gridX .. ' ' .. anchorBoxIndex) | |
| for index = 1, 2 do | |
| for c = 1, 2 do | |
| print(item[index][c]) | |
| end | |
| end | |
| for c = 1, 3 do | |
| print(item[3][c]) | |
| end | |
| end | |
| if IOU > anchorBoxesTensor[gridY][gridX][anchorBoxIndex][5] then | |
| anchorBoxesTensor[gridY][gridX][anchorBoxIndex][{{1, 5}}] = gridAnchorBoxesTensor[anchorBoxIndex] | |
| if IOU > confidenceThreshold then | |
| anchorBoxesTensor[gridY][gridX][anchorBoxIndex][6] = label | |
| end | |
| end | |
| if not maxIOU or IOU > maxIOU then | |
| maxIOUIndex = anchorBoxIndex | |
| maxIOU = IOU | |
| end | |
| end | |
| if maxIOU < confidenceThreshold then | |
| if maxIOU < 0.3 and false then | |
| print(centerX) | |
| print(centerY) | |
| print(gridX) | |
| print(gridY) | |
| for index = 1, 2 do | |
| for c = 1, 2 do | |
| print(gridBox[index][c]) | |
| end | |
| end | |
| print(maxIOU) | |
| print(maxIOUIndex) | |
| print(gridWidth) | |
| print(gridHeight) | |
| os.exit(1) | |
| end | |
| anchorBoxesTensor[gridY][gridX][maxIOUIndex][6] = label | |
| end | |
| end | |
| end | |
| anchorBoxesTensors[mode] = anchorBoxesTensor | |
| end | |
| local representationTensor | |
| local sortedModes = utils.getSortedModes() | |
| for _, mode in pairs(sortedModes) do | |
| local anchorBoxesTensor = anchorBoxesTensors[mode] | |
| if not representationTensor then | |
| representationTensor = anchorBoxesTensor | |
| else | |
| representationTensor = torch.cat(representationTensor, anchorBoxesTensor, 3) | |
| end | |
| end | |
| return representationTensor | |
| end | |
| function utils.getHeatmapsStyless(width, height, representation, lineWidth, kernelSize) | |
| representation.walls = utils.mergeLines(representation.walls, lineWidth) | |
| representation.points = utils.linesToPoints(width, height, representation.walls, lineWidth) | |
| --local kernelSize = kernelSize or torch.round(7.0 * math.max(width, height) / 256) | |
| local kernelSize = kernelSize or 13 | |
| local heatmaps = torch.zeros(13 + 4 + 4, height, width) | |
| for mode, items in pairs(representation) do | |
| if mode == 'points' then | |
| for _, item in pairs(items) do | |
| local x = torch.round(item[1][1]) | |
| local y = torch.round(item[1][2]) | |
| local label = utils.getNumber(mode, item[3]) | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| heatmaps[label][y][x] = 1 | |
| end | |
| end | |
| elseif mode == 'doors' then | |
| for _, item in pairs(items) do | |
| local lineDim = utils.lineDim(item) | |
| if lineDim > 0 then | |
| local x = torch.round(math.min(item[1][1], item[2][1])) | |
| local y = torch.round(math.min(item[1][2], item[2][2])) | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| heatmaps[13 + (lineDim - 1) * 2 + 1][y][x] = 1 | |
| end | |
| local x = torch.round(math.max(item[1][1], item[2][1])) | |
| local y = torch.round(math.max(item[1][2], item[2][2])) | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| heatmaps[13 + (lineDim - 1) * 2 + 2][y][x] = 1 | |
| end | |
| end | |
| end | |
| elseif mode == 'icons' then | |
| for _, item in pairs(items) do | |
| local lineDim = utils.lineDim(item) | |
| local x = torch.round(math.min(item[1][1], item[2][1])) | |
| local y = torch.round(math.min(item[1][2], item[2][2])) | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| heatmaps[13 + 4 + 1][y][x] = 1 | |
| end | |
| local x = torch.round(math.max(item[1][1], item[2][1])) | |
| local y = torch.round(math.min(item[1][2], item[2][2])) | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| heatmaps[13 + 4 + 2][y][x] = 1 | |
| end | |
| local x = torch.round(math.min(item[1][1], item[2][1])) | |
| local y = torch.round(math.max(item[1][2], item[2][2])) | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| heatmaps[13 + 4 + 3][y][x] = 1 | |
| end | |
| local x = torch.round(math.max(item[1][1], item[2][1])) | |
| local y = torch.round(math.max(item[1][2], item[2][2])) | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| heatmaps[13 + 4 + 4][y][x] = 1 | |
| end | |
| end | |
| end | |
| end | |
| --kernelSize = torch.round(1.0 * kernelSize * math.max(width, height) / 256) | |
| local kernel = image.gaussian(kernelSize) | |
| --local kernel = image.gaussian(nil, nil, nil, nil, torch.round(1.0 * kernelSize * width / 256), torch.round(1.0 * kernelSize * height / 256)) | |
| for i = 1, heatmaps:size(1) do | |
| heatmaps[i] = image.convolve(heatmaps[i], kernel, 'same') | |
| end | |
| return heatmaps | |
| end | |
| function utils.getHeatmaps(width, height, representation, lineWidth, kernelSize) | |
| representation.points = utils.linesToPoints(width, height, representation.walls, lineWidth) | |
| local kernelSize = kernelSize or torch.round(7.0 * math.max(width, height) / 256) | |
| local heatmaps = torch.zeros(13 * (1 + 4 + 4), height, width) | |
| for mode, items in pairs(representation) do | |
| if mode == 'points' then | |
| for _, item in pairs(items) do | |
| local x = torch.round(item[1][1]) | |
| local y = torch.round(item[1][2]) | |
| local label = utils.getNumber(mode, item[3]) | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| heatmaps[label][y][x] = 1 | |
| end | |
| end | |
| elseif mode == 'doors' then | |
| for _, item in pairs(items) do | |
| local lineDim = utils.lineDim(item) | |
| if lineDim > 0 then | |
| local label = utils.getNumber(mode, item[3]) | |
| local x = torch.round(math.min(item[1][1], item[2][1])) | |
| local y = torch.round(math.min(item[1][2], item[2][2])) | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| heatmaps[13 + (label - 1) * 4 + (lineDim - 1) * 2 + 1][y][x] = 1 | |
| end | |
| local x = torch.round(math.max(item[1][1], item[2][1])) | |
| local y = torch.round(math.max(item[1][2], item[2][2])) | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| heatmaps[13 + (label - 1) * 4 + (lineDim - 1) * 2 + 2][y][x] = 1 | |
| end | |
| end | |
| end | |
| elseif mode == 'icons' then | |
| for _, item in pairs(items) do | |
| local lineDim = utils.lineDim(item) | |
| local label = utils.getNumber(mode, item[3]) | |
| local x = torch.round(math.min(item[1][1], item[2][1])) | |
| local y = torch.round(math.min(item[1][2], item[2][2])) | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| heatmaps[13 * 5 + (label - 1) * 4 + 1][y][x] = 1 | |
| end | |
| local x = torch.round(math.max(item[1][1], item[2][1])) | |
| local y = torch.round(math.min(item[1][2], item[2][2])) | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| heatmaps[13 * 5 + (label - 1) * 4 + 2][y][x] = 1 | |
| end | |
| local x = torch.round(math.min(item[1][1], item[2][1])) | |
| local y = torch.round(math.max(item[1][2], item[2][2])) | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| heatmaps[13 * 5 + (label - 1) * 4 + 3][y][x] = 1 | |
| end | |
| local x = torch.round(math.max(item[1][1], item[2][1])) | |
| local y = torch.round(math.max(item[1][2], item[2][2])) | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| heatmaps[13 * 5 + (label - 1) * 4 + 4][y][x] = 1 | |
| end | |
| end | |
| end | |
| end | |
| local kernel = image.gaussian(kernelSize) | |
| heatmaps = image.convolve(heatmaps, kernel, 'same') | |
| return heatmaps | |
| end | |
| function utils.getJunctionHeatmap(width, height, representationGlobal, kernelSize) | |
| local kernelSize = kernelSize or 7 | |
| local heatmaps = torch.zeros(13, height, width) | |
| for mode, items in pairs(representationGlobal) do | |
| if mode == 'points' then | |
| for _, item in pairs(representationGlobal.points) do | |
| local x = torch.round(item[1][1]) | |
| local y = torch.round(item[1][2]) | |
| local label = utils.getNumber('points', item[3]) | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| heatmaps[label][y][x] = 1 | |
| end | |
| end | |
| end | |
| end | |
| local kernel = image.gaussian(kernelSize) | |
| for i = 1, heatmaps:size(1) do | |
| heatmaps[i] = image.convolve(heatmaps[i], kernel, 'same') | |
| end | |
| return heatmaps | |
| end | |
| function utils.convertTensorAnchorToRepresentation(width, height, representationTensor, confidenceThreshold) | |
| local confidenceThreshold = confidenceThreshold or 0.7 | |
| local gridWidth = representationTensor:size(2) | |
| local gridHeight = representationTensor:size(1) | |
| local cellWidth = width / gridWidth | |
| local cellHeight = height / gridHeight | |
| local anchorBoxesMap = utils.getAnchorBoxesMap(cellWidth, cellHeight) | |
| local anchorBoxesTensors = {} | |
| local offset = 1 | |
| local sortedModes = utils.getSortedModes() | |
| for _, mode in pairs(sortedModes) do | |
| local anchorBoxes = anchorBoxesMap[mode] | |
| anchorBoxesTensors[mode] = representationTensor:narrow(3, offset, #anchorBoxes) | |
| offset = offset + #anchorBoxes | |
| end | |
| --print(anchorBoxesTensors['points']) | |
| local representation = {} | |
| for mode, anchorBoxesTensor in pairs(anchorBoxesTensors) do | |
| representation[mode] = {} | |
| local anchorBoxes = anchorBoxesMap[mode] | |
| local numAnchorBoxes = #anchorBoxes | |
| for gridX = 1, gridWidth do | |
| for gridY = 1, gridHeight do | |
| for anchorBoxIndex = 1, numAnchorBoxes do | |
| if anchorBoxesTensor[gridY][gridX][anchorBoxIndex][5] > confidenceThreshold then | |
| --print(mode) | |
| --print(anchorBoxesTensor[gridY][gridX][anchorBoxIndex]) | |
| local gridBox = anchorBoxesTensor[gridY][gridX][anchorBoxIndex][{{1, 4}}] | |
| local anchorBox = anchorBoxes[anchorBoxIndex] | |
| local centerX = anchorBox[1][1] + gridBox[1] * anchorBox[2][1] + (gridX - 1) * cellWidth | |
| local centerY = anchorBox[1][2] + gridBox[2] * anchorBox[2][2] + (gridY - 1) * cellHeight | |
| local boxWidth = math.exp(gridBox[3]) * anchorBox[2][1] | |
| local boxHeight = math.exp(gridBox[4]) * anchorBox[2][2] | |
| local x_1 = math.min(math.max(torch.round(centerX - boxWidth / 2), 1), width) | |
| local y_1 = math.min(math.max(torch.round(centerY - boxHeight / 2), 1), height) | |
| local x_2 = math.max(math.min(torch.round(centerX + boxWidth / 2), width), 1) | |
| local y_2 = math.max(math.min(torch.round(centerY + boxHeight / 2), height), 1) | |
| local label = anchorBoxesTensor[gridY][gridX][anchorBoxIndex][6] | |
| local itemInfo = utils.getItemInfo(mode, label) | |
| if not itemInfo then | |
| print('cannot find item info') | |
| print(anchorBoxesTensor[gridY][gridX][anchorBoxIndex]) | |
| print(mode) | |
| print(label) | |
| os.exit(1) | |
| end | |
| local item = {{x_1, y_1}, {x_2, y_2}, itemInfo} | |
| if mode == 'icons' or mode == 'doors' then | |
| table.insert(item, anchorBoxesTensor[gridY][gridX][anchorBoxIndex][5]) | |
| end | |
| table.insert(representation[mode], item) | |
| end | |
| end | |
| end | |
| end | |
| end | |
| representation.doors = utils.nms(representation.doors) | |
| representation.icons = utils.nms(representation.icons) | |
| return representation | |
| end | |
| function utils.nms(items, threshold) | |
| local nms = require 'nms' | |
| require 'cunn' | |
| --require 'cutorch' | |
| local threshold = threshold or 0 | |
| local nameItemsMap = {} | |
| local nameItemsTensorMap = {} | |
| for _, item in pairs(items) do | |
| local itemTensor = torch.DoubleTensor({item[1][1], item[1][2], item[2][1], item[2][2], item[4]}):repeatTensor(1, 1) | |
| local name = item[3][1] | |
| if not nameItemsMap[name] then | |
| nameItemsTensorMap[name] = itemTensor | |
| nameItemsMap[name] = {item} | |
| else | |
| nameItemsTensorMap[name] = torch.cat(nameItemsTensorMap[name], itemTensor, 1) | |
| table.insert(nameItemsMap[name], item) | |
| end | |
| end | |
| local newItems = {} | |
| for name, itemsTensor in pairs(nameItemsTensorMap) do | |
| --print(name) | |
| --print(itemsTensor) | |
| local validIndices = nms.gpu_nms(itemsTensor:cuda(), threshold) | |
| --print(validIndices) | |
| for i = 1, validIndices:size(1) do | |
| local itemTensor = itemsTensor[validIndices[i]] | |
| local item = nameItemsMap[name][validIndices[i]] | |
| table.insert(newItems, {{itemTensor[1], itemTensor[2]}, {itemTensor[3], itemTensor[4]}, item[3]}) | |
| end | |
| end | |
| return newItems | |
| end | |
| function utils.compressRepresentation(width, height, representationTensor) | |
| local offsetsBB = utils.offsetsBB() | |
| local offsetsClass = utils.offsetsClass() | |
| local numItemsPerCell = utils.numItemsPerCell() | |
| local numItemsGlobal = utils.numItemsGlobal() | |
| local numFeaturesPerItem = utils.numFeaturesPerItem() | |
| local numFeaturesBB = utils.numFeaturesBB() | |
| local gridWidth = representationTensor:size(3) | |
| local gridHeight = representationTensor:size(2) | |
| local confidenceThreshold = confidenceThreshold or 0.5 | |
| local items = {} | |
| for x = 1, gridWidth do | |
| for y = 1, gridHeight do | |
| for mode, numItems in pairs(numItemsPerCell) do | |
| if mode == 'icons' or mode == 'labels' then | |
| local offsetBB = offsetsBB[mode] | |
| local offsetClass = offsetsClass[mode] | |
| local numFeatures = numFeaturesPerItem[mode] | |
| for itemIndex = 1, numItems do | |
| local item = representationTensor[{{offsetBB + (itemIndex - 1) * numFeatures + 1, offsetBB + (itemIndex - 1) * numFeatures + numFeatures}, y, x}] | |
| if item[5] > confidenceThreshold then | |
| local class = representationTensor[{offsetClass + itemIndex, y, x}] | |
| local name = utils.getItemInfo(mode, class)[1] | |
| local newRectangle = utils.gridRectangle(width, height, 1, 1, utils.imageRectangle(width, height, gridWidth, gridHeight, {{x, y}, {item[1], item[2]}, {item[3], item[4]}})) | |
| if items[name] == nil then | |
| items[name] = {} | |
| end | |
| table.insert(items[name], torch.Tensor({newRectangle[2][1], newRectangle[2][2], newRectangle[3][1], newRectangle[3][2], item[5]})) | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| local totalNumItems = 0 | |
| for _, num in pairs(numItemsGlobal) do | |
| totalNumItems = totalNumItems + num | |
| end | |
| local iconLabelTensor = torch.zeros(totalNumItems * 5) | |
| local offset = 0 | |
| for name, numItems in pairs(numItemsGlobal) do | |
| if items[name] ~= nil then | |
| for itemIndex = 1, math.min(numItems, #items[name]) do | |
| iconLabelTensor[{{offset + (itemIndex - 1) * 5 + 1, offset + (itemIndex - 1) * 5 + 5}}] = items[name][itemIndex] | |
| end | |
| end | |
| offset = offset + numItems * 5 | |
| end | |
| return {representationTensor[{{1, offsetsBB.icons}}], representationTensor[{{numFeaturesBB, offsetsClass.icons}}], iconLabelTensor:view(-1, 1, 1)} | |
| end | |
| function utils.uncompressRepresentation(width, height, representationTensors) | |
| local offsetsBB = utils.offsetsBB() | |
| local offsetsClass = utils.offsetsClass() | |
| local numItemsPerCell = utils.numItemsPerCell() | |
| local numItemsGlobal = utils.numItemsGlobal() | |
| local numFeaturesPerItem = utils.numFeaturesPerItem() | |
| local numFeaturesBB = utils.numFeaturesBB() | |
| local numFeaturesClass = utils.numFeaturesClass() | |
| local modeMap = utils.modeMap() | |
| local gridWidth = representationTensors[1]:size(3) | |
| local gridHeight = representationTensors[1]:size(2) | |
| local confidenceThreshold = confidenceThreshold or 0.5 | |
| local iconLabelTensor = representationTensors[3] | |
| --local iconLabelTensorBB = torch.zeros(numFeaturesBB - offsetsBB.icons, gridHeight, gridWidth) | |
| --local iconLabelTensorClass = torch.zeros(numFeaturesClass + numFeaturesBB - offsetsClass.icons, gridHeight, gridWidth) | |
| local representationTensor = torch.zeros(numFeaturesBB + numFeaturesClass, gridHeight, gridWidth) | |
| representationTensor[{{1, offsetsBB.icons}}] = representationTensors[1] | |
| representationTensor[{{numFeaturesBB + 1, offsetsClass.icons}}] = representationTensors[2] | |
| local gridItems = {} | |
| for y = 1, gridHeight do | |
| gridItems[y] = {} | |
| for x = 1, gridWidth do | |
| gridItems[y][x] = {} | |
| for mode, offset in pairs(offsetsBB) do | |
| gridItems[y][x][mode] = {} | |
| end | |
| end | |
| end | |
| local offset = 0 | |
| for name, numItems in pairs(numItemsGlobal) do | |
| local mode = modeMap[name] | |
| for itemIndex = 1, numItems do | |
| local item = iconLabelTensor[{{offset + (itemIndex - 1) * 5 + 1, offset + (itemIndex - 1) * 5 + 5}, 1, 1}] | |
| local newRectangle = utils.gridRectangle(width, height, gridWidth, gridWidth, utils.imageRectangle(width, height, 1, 1, {{1, 1}, {item[1], item[2]}, {item[3], item[4]}})) | |
| table.insert(newRectangle, {name, 1, 1}) | |
| --print(item) | |
| --print(newRectangle) | |
| --print(mode) | |
| --print(name) | |
| if newRectangle[1][1] >= 1 and newRectangle[1][1] <= gridWidth and newRectangle[1][2] >= 1 and newRectangle[1][2] <= gridHeight then | |
| table.insert(gridItems[newRectangle[1][2]][newRectangle[1][1]][mode], {item[5], newRectangle}) | |
| end | |
| end | |
| offset = offset + numItems * 5 | |
| end | |
| for x = 1, gridWidth do | |
| for y = 1, gridHeight do | |
| for mode, items in pairs(gridItems[y][x]) do | |
| local numFeatures = numFeaturesPerItem[mode] | |
| local sortedItems = items | |
| table.sort(sortedItems, function(a, b) return a[1] > b[1] end) | |
| for i = numItemsPerCell[mode] + 1, #sortedItems do | |
| table.remove(sortedItems) | |
| end | |
| local offsetBB = offsetsBB[mode] | |
| for itemIndex, item in pairs(sortedItems) do | |
| item = item[2] | |
| local itemTensor = representationTensor[{{offsetBB + (itemIndex - 1) * numFeatures + 1, offsetBB + (itemIndex - 1) * numFeatures + numFeatures}, y, x}] | |
| itemTensor[1] = item[2][1] | |
| itemTensor[2] = item[2][2] | |
| itemTensor[3] = item[3][1] | |
| itemTensor[4] = item[3][2] | |
| itemTensor[5] = 1 | |
| end | |
| local offsetClass = offsetsClass[mode] | |
| for itemIndex, item in pairs(sortedItems) do | |
| item = item[2] | |
| local class | |
| if mode == 'icons' then | |
| class = utils.getNumber(mode, item[4]) | |
| elseif mode == 'labels' then | |
| class = utils.getNumber(mode, item[4]) | |
| end | |
| representationTensor[{offsetClass + itemIndex, y, x}] = class | |
| end | |
| end | |
| end | |
| end | |
| return representationTensor | |
| end | |
| function utils.convertTensorToRepresentation(width, height, representationTensor, confidenceThreshold) | |
| local offsetsBB = utils.offsetsBB() | |
| local offsetsClass = utils.offsetsClass() | |
| local numItemsPerCell = utils.numItemsPerCell() | |
| local numFeaturesPerItem = utils.numFeaturesPerItem() | |
| local representation = {} | |
| representation.points = {} | |
| representation.doors = {} | |
| representation.icons = {} | |
| representation.labels = {} | |
| local gridWidth = representationTensor:size(3) | |
| local gridHeight = representationTensor:size(2) | |
| local confidenceThreshold = confidenceThreshold or 0.5 | |
| for x = 1, gridWidth do | |
| for y = 1, gridHeight do | |
| for mode, numItems in pairs(numItemsPerCell) do | |
| local offsetBB = offsetsBB[mode] | |
| local offsetClass = offsetsClass[mode] | |
| local numFeatures = numFeaturesPerItem[mode] | |
| for itemIndex = 1, numItems do | |
| local item = representationTensor[{{offsetBB + (itemIndex - 1) * numFeatures + 1, offsetBB + (itemIndex - 1) * numFeatures + numFeatures}, y, x}] | |
| local class = representationTensor[{offsetClass + itemIndex, y, x}] | |
| if item[5] > confidenceThreshold then | |
| table.insert(representation[mode], {{x, y}, {item[1], item[2]}, {item[3], item[4]}, utils.getItemInfo(mode, class)}) | |
| end | |
| end | |
| end | |
| end | |
| end | |
| local representationGlobal = utils.gridToGlobal(width, height, gridWidth, gridHeight, representation) | |
| return representationGlobal | |
| end | |
| function utils.cropRepresentation(representation, startX, startY, endX, endY) | |
| local newRepresentation = {} | |
| for mode, items in pairs(representation) do | |
| newRepresentation[mode] = {} | |
| for _, item in pairs(items) do | |
| if (item[1][1] + item[2][1]) / 2 <= endX and (item[1][2] + item[2][2]) / 2 <= endY and (item[1][1] + item[2][1]) / 2 > startX and (item[1][2] + item[2][2]) / 2 > startY then | |
| if mode == 'icons' or mode == 'labels' then | |
| assert(item[1][1] <= item[2][1] and item[1][2] <= item[2][2], mode .. ' ' .. item[1][1] .. ' ' .. item[2][1] .. ' ' .. item[1][2] .. ' ' .. item[2][2]) | |
| end | |
| for pointIndex = 1, 2 do | |
| item[pointIndex][1] = item[pointIndex][1] - startX | |
| item[pointIndex][2] = item[pointIndex][2] - startY | |
| end | |
| item[1][1] = math.max(item[1][1], 1) | |
| item[1][2] = math.max(item[1][2], 1) | |
| item[2][1] = math.min(item[2][1], endX - startX) | |
| item[2][2] = math.min(item[2][2], endY - startY) | |
| if mode == 'icons' or mode == 'labels' then | |
| if item[1][1] > item[2][1] or item[1][2] > item[2][2] then | |
| print('invalid rectangle') | |
| print(item[1][1] .. ' ' .. item[2][1] .. ' ' .. item[1][2] .. ' ' .. item[2][2]) | |
| print(startX .. ' ' .. startY) | |
| print(endX .. ' ' .. endY) | |
| os.exit(1) | |
| end | |
| end | |
| table.insert(newRepresentation[mode], item) | |
| end | |
| end | |
| end | |
| return newRepresentation | |
| end | |
| function utils.scaleRepresentation(representation, width, height, newWidth, newHeight) | |
| local newRepresentation = {} | |
| for mode, items in pairs(representation) do | |
| newRepresentation[mode] = {} | |
| for _, item in pairs(items) do | |
| for pointIndex = 1, 2 do | |
| item[pointIndex][1] = torch.round((item[pointIndex][1] - 1) / width * newWidth + 1) | |
| item[pointIndex][2] = torch.round((item[pointIndex][2] - 1) / height * newHeight + 1) | |
| end | |
| table.insert(newRepresentation[mode], item) | |
| end | |
| end | |
| return newRepresentation | |
| end | |
| function utils.scaleRepresentationByRatio(representation, ratio) | |
| local newRepresentation = {} | |
| for mode, items in pairs(representation) do | |
| newRepresentation[mode] = {} | |
| for _, item in pairs(items) do | |
| for pointIndex = 1, 2 do | |
| item[pointIndex][1] = torch.round((item[pointIndex][1] - 1) * ratio + 1) | |
| item[pointIndex][2] = torch.round((item[pointIndex][2] - 1) * ratio + 1) | |
| end | |
| table.insert(newRepresentation[mode], item) | |
| end | |
| end | |
| return newRepresentation | |
| end | |
| function utils.rotateRepresentation(representation, width, height, orientation) | |
| --[[ | |
| if orientation == 1 then | |
| return representation | |
| end | |
| ]]-- | |
| local newRepresentation = {} | |
| for mode, items in pairs(representation) do | |
| newRepresentation[mode] = {} | |
| for _, item in pairs(items) do | |
| for pointIndex = 1, 2 do | |
| local x = item[pointIndex][1] | |
| local y = item[pointIndex][2] | |
| if orientation == 2 then | |
| item[pointIndex][1] = height - y | |
| item[pointIndex][2] = x | |
| elseif orientation == 3 then | |
| item[pointIndex][1] = width - x | |
| item[pointIndex][2] = height - y | |
| elseif orientation == 4 then | |
| item[pointIndex][1] = y | |
| item[pointIndex][2] = width - x | |
| end | |
| end | |
| if mode == 'icons' or mode == 'labels' then | |
| item = {{math.min(item[1][1], item[2][1]), math.min(item[1][2], item[2][2])}, {math.max(item[1][1], item[2][1]), math.max(item[1][2], item[2][2])}, item[3]} | |
| end | |
| table.insert(newRepresentation[mode], item) | |
| end | |
| end | |
| newRepresentation.walls = utils.sortLines(newRepresentation.walls) | |
| newRepresentation.doors = utils.sortLines(newRepresentation.doors) | |
| return newRepresentation | |
| end | |
| function utils.drawRepresentationImage(floorplan, representation, representationType, renderingType, lineWidth) | |
| local width, height = floorplan:size(3), floorplan:size(2) | |
| local iconImages = utils.loadIconImages() | |
| local representationType = representationType or 'P' | |
| local renderingType = renderingType or 'L' | |
| local lineWidth = lineWidth or 5 | |
| local colorMap = {} | |
| colorMap[1] = {255, 0, 0} | |
| colorMap[2] = {0, 255, 0} | |
| colorMap[3] = {0, 0, 255} | |
| colorMap[4] = {255, 255, 0} | |
| colorMap[5] = {0, 255, 255} | |
| colorMap[6] = {255, 0, 255} | |
| colorMap[7] = {128, 0, 0} | |
| colorMap[8] = {0, 128, 0} | |
| colorMap[9] = {0, 0, 128} | |
| if representationType == 'P' and renderingType == 'P' then | |
| local representationImage = floorplan | |
| for mode, items in pairs(representation) do | |
| if mode == 'points' or mode == 'doors' or mode == 'icons' or mode == 'labels' then | |
| for __, item in pairs(items) do | |
| local rectangle = {item[1], item[2]} | |
| local center = {(rectangle[2][1] + rectangle[1][1]) / 2, (rectangle[2][2] + rectangle[1][2]) / 2} | |
| local rectangleWidth = math.max(rectangle[2][1] - rectangle[1][1] + 1, 10) | |
| local rectangleHeight = math.max(rectangle[2][2] - rectangle[1][2] + 1, 10) | |
| local color = colorMap[item[3][2]] | |
| local strokeWidth = 2 | |
| if mode == 'points' then | |
| strokeWidth = 10 | |
| end | |
| local success, newRepresentationImage = pcall(function() return image.drawRect(representationImage, math.max(center[1] - rectangleWidth / 2, 1), math.max(center[2] - rectangleHeight / 2, 1), math.min(center[1] + rectangleWidth / 2, width), math.min(center[2] + rectangleHeight / 2, height), {lineWidth = strokeWidth, color = color}) end) | |
| if success then | |
| representationImage = newRepresentationImage | |
| end | |
| end | |
| end | |
| end | |
| return representationImage | |
| elseif representationType == 'P' and renderingType == 'L' then | |
| if not representation.points or #representation.points == 0 then | |
| representation.points = utils.linesToPoints(width, height, representation.walls, lineWidth) | |
| end | |
| local representationImage = torch.ones(#floorplan) | |
| local lineWidth = 4 | |
| local lines = utils.pointsToLines(width, height, representation.points, lineWidth, true) | |
| local lineMask = utils.drawLineMask(width, height, lines, lineWidth) | |
| local segmentation = utils.getSegmentation(width, height, representation, lineWidth)[1] | |
| local colorMap = {{224, 255, 192}, {224, 255, 192}, {255, 224, 128}, {192, 255, 255}, {192, 255, 255}, {255, 224, 224}, {192, 192, 224}, {255, 160, 96}, {192, 255, 255}, {224, 224, 224}} | |
| local borderColorMap = {{128, 192, 64}, {128, 192, 64}, {192, 128, 64}, {0, 128, 192}, {0, 128, 192}, {192, 64, 64}, {128, 64, 160}, {192, 64, 0}, {0, 128, 192}, {255, 255, 255}} | |
| local representationImageVideo = representationImage | |
| image.save('test/representation_image_background.png', representationImageVideo) | |
| local representationImageVideo = torch.zeros(3, height, width) | |
| --local representationImageMask = torch.zeros(4, height, width) | |
| --representationImageMask[{{1, 3}}] = 1 | |
| for segmentIndex = 1, 10 do | |
| local segmentMask = segmentation:eq(segmentIndex) | |
| local dilatedSegmentMask = segmentMask | |
| for i = 1, lineWidth + 1 do | |
| dilatedSegmentMask = image.dilate(dilatedSegmentMask) | |
| end | |
| local borderMask = dilatedSegmentMask - segmentMask | |
| for c = 1, 3 do | |
| representationImage[c][borderMask] = borderColorMap[segmentIndex][c] / 256 | |
| representationImageVideo[c][borderMask] = borderColorMap[segmentIndex][c] / 256 | |
| end | |
| --representationImageVideo[4][borderMask] = 1 | |
| --representationImageMask[4][borderMask] = 1 | |
| end | |
| image.save('test/representation_image_walls.png', representationImageVideo) | |
| --local representationImageVideo = torch.ones(3, height, width) | |
| for segmentIndex = 1, 10 do | |
| local segmentMask = segmentation:eq(segmentIndex) | |
| for c = 1, 3 do | |
| representationImage[c][segmentMask] = colorMap[segmentIndex][c] / 256 | |
| representationImageVideo[c][segmentMask] = colorMap[segmentIndex][c] / 256 | |
| end | |
| --representationImageVideo[4][segmentMask] = 1 | |
| --representationImageMask[4][segmentMask] = 1 | |
| end | |
| image.save('test/representation_image_rooms.png', representationImageVideo) | |
| for __, item in pairs(representation.icons) do | |
| item[3] = utils.getItemInfo('icons', utils.getNumber('icons', item[3])) | |
| local iconImage = iconImages[item[3][1]][item[3][2]] | |
| --local rectangle = {item[1], item[2]} | |
| local rectangle = {{math.min(item[1][1], item[2][1]), math.min(item[1][2], item[2][2])}, {math.max(item[1][1], item[2][1]), math.max(item[1][2], item[2][2])}} | |
| for pointIndex = 1, 2 do | |
| rectangle[pointIndex][1] = torch.round(math.max(math.min(rectangle[pointIndex][1], width), 1)) | |
| rectangle[pointIndex][2] = torch.round(math.max(math.min(rectangle[pointIndex][2], height), 1)) | |
| end | |
| --print(mode) | |
| --print(rectangle) | |
| local orientation = 1 | |
| local minDistance = math.max(width, height) | |
| local minDistanceOrientation | |
| if (math.abs(rectangle[2][1] - rectangle[1][1]) > math.abs(rectangle[2][2] - rectangle[1][2]) and item[3][1] ~= 'toilet' and item[3][1] ~= 'stairs') or (math.abs(rectangle[2][1] - rectangle[1][1]) < math.abs(rectangle[2][2] - rectangle[1][2]) and (item[3][1] == 'toilet' or item[3][1] == 'stairs')) or item[3][1] == 'washing_basin' then | |
| local min = math.min(rectangle[1][2], rectangle[2][2]) | |
| local max = math.max(rectangle[1][2], rectangle[2][2]) | |
| local center = (rectangle[1][1] + rectangle[2][1]) / 2 | |
| local deltaMin = 0 | |
| for delta = 1, min - 1 do | |
| deltaMin = delta | |
| if lineMask[min - delta][center] == 1 then | |
| break | |
| end | |
| end | |
| local deltaMax = 0 | |
| for delta = 1, height - max do | |
| deltaMax = delta | |
| if lineMask[max + delta][center] == 1 then | |
| break | |
| end | |
| end | |
| if deltaMax < minDistance then | |
| minDistance = deltaMax | |
| minDistanceOrientation = 1 | |
| end | |
| if deltaMin < minDistance then | |
| minDistance = deltaMin | |
| minDistanceOrientation = 3 | |
| end | |
| end | |
| if (math.abs(rectangle[2][1] - rectangle[1][1]) < math.abs(rectangle[2][2] - rectangle[1][2]) and item[3][1] ~= 'toilet') or (math.abs(rectangle[2][1] - rectangle[1][1]) > math.abs(rectangle[2][2] - rectangle[1][2]) and item[3][1] == 'toilet') or item[3][1] == 'washing_basin' then | |
| local min = math.min(rectangle[1][1], rectangle[2][1]) | |
| local max = math.max(rectangle[1][1], rectangle[2][1]) | |
| local center = (rectangle[1][2] + rectangle[2][2]) / 2 | |
| local deltaMin = 0 | |
| for delta = 1, min - 1 do | |
| deltaMin = delta | |
| if lineMask[center][min - delta] == 1 then | |
| break | |
| end | |
| end | |
| local deltaMax = 0 | |
| for delta = 1, width - max do | |
| deltaMax = delta | |
| if lineMask[center][max + delta] == 1 then | |
| break | |
| end | |
| end | |
| if deltaMax < minDistance then | |
| minDistance = deltaMax | |
| minDistanceOrientation = 4 | |
| end | |
| if deltaMin < minDistance then | |
| minDistance = deltaMin | |
| minDistanceOrientation = 2 | |
| end | |
| end | |
| if iconImage == nil then | |
| print(item) | |
| end | |
| orientation = minDistanceOrientation | |
| --[[ | |
| if item[3][1] == 'stairs' and (rectangle[1][2] + rectangle[2][2]) / 2 > 277 then | |
| orientation = 1 | |
| end | |
| ]]-- | |
| iconImage = rotateIcon(iconImage, orientation) | |
| local icon = image.scale(iconImage, rectangle[2][1] - rectangle[1][1] + 1, rectangle[2][2] - rectangle[1][2] + 1) | |
| if icon:dim() == 2 or icon:size(1) == 1 then | |
| icon = icon:repeatTensor(3, 1, 1) | |
| end | |
| --print(#representationImage[{{}, {rectangle[1][2], rectangle[2][2]}, {rectangle[1][1], rectangle[2][1]}}]) | |
| --[[ | |
| print(#iconImage) | |
| print(#icon) | |
| print(rectangle[2][1] - rectangle[1][1] + 1) | |
| print(rectangle[2][2] - rectangle[1][2] + 1) | |
| print(item[3][1]) | |
| ]]-- | |
| representationImage[{{}, {rectangle[1][2], rectangle[2][2]}, {rectangle[1][1], rectangle[2][1]}}] = icon | |
| local representationImageVideo = torch.zeros(3, height, width) | |
| representationImageVideo[{{1, 3}, {rectangle[1][2], rectangle[2][2]}, {rectangle[1][1], rectangle[2][1]}}] = icon | |
| --representationImageVideo[{{4}, {rectangle[1][2], rectangle[2][2]}, {rectangle[1][1], rectangle[2][1]}}] = 1 | |
| --representationImageMask[{{4}, {rectangle[1][2], rectangle[2][2]}, {rectangle[1][1], rectangle[2][1]}}] = 1 | |
| image.save('test/representation_image_icon_' .. __ .. '.png', representationImageVideo) | |
| end | |
| local rooms, numRooms = utils.findConnectedComponents(1 - lineMask) | |
| --image.save('test/segmentation.png', utils.drawSegmentation(segmentation)) | |
| --os.exit(1) | |
| rooms = segmentation:int() * numRooms + rooms:int() | |
| local roomLabelMap = {} | |
| for __, item in pairs(representation.labels) do | |
| local labelImage = iconImages[item[3][1]][item[3][2]] | |
| --local rectangle = {item[1], item[2]} | |
| local cx = (item[1][1] + item[2][1]) / 2 | |
| local cy = (item[1][2] + item[2][2]) / 2 | |
| local labelWidth = 80 | |
| local labelHeight = 30 | |
| local roomIndex = rooms[torch.round(cy)][torch.round(cx)] | |
| if not roomLabelMap[roomIndex] then | |
| roomLabelMap[roomIndex] = {} | |
| end | |
| local number = utils.getNumber('labels', item[3]) | |
| if roomIndex > 0 and not roomLabelMap[roomIndex][number] then | |
| roomLabelMap[roomIndex][number] = true | |
| local roomIndices = rooms:eq(roomIndex):nonzero() | |
| local mins = torch.min(roomIndices, 1)[1] | |
| local maxs = torch.max(roomIndices, 1)[1] | |
| if maxs[2] - mins[2] + 1 < labelWidth * 0.8 then | |
| if math.min((maxs[2] - mins[2] + 1) / labelHeight, (maxs[1] - mins[1] + 1) / labelWidth) > (maxs[2] - mins[2] + 1) / labelWidth then | |
| labelImage = rotateIcon(labelImage, 2) | |
| labelWidth, labelHeight = labelHeight, labelWidth | |
| end | |
| if maxs[2] - mins[2] + 1 < labelWidth then | |
| local ratio = math.max((maxs[2] - mins[2] + 1) / labelWidth, 0.5) | |
| labelWidth, labelHeight = labelWidth * ratio, labelHeight * ratio | |
| end | |
| end | |
| if math.max(labelWidth, labelHeight) > 10 then | |
| local rectangle = {{cx - labelWidth / 2, cy - labelHeight / 2}, {cx + labelWidth / 2, cy + labelHeight / 2}} | |
| for pointIndex = 1, 2 do | |
| rectangle[pointIndex][1] = torch.round(math.max(math.min(rectangle[pointIndex][1], width), 1)) | |
| rectangle[pointIndex][2] = torch.round(math.max(math.min(rectangle[pointIndex][2], height), 1)) | |
| end | |
| --[[ | |
| if item[3][1] == 'bathroom' or item[3][1] == 'kitchen' then | |
| for pointIndex = 1, 2 do | |
| rectangle[pointIndex][2] = rectangle[pointIndex][2] - 20 | |
| end | |
| end | |
| if item[3][1] == 'bathroom' then | |
| for pointIndex = 1, 2 do | |
| rectangle[pointIndex][1] = rectangle[pointIndex][1] + 10 | |
| end | |
| end | |
| ]]-- | |
| local label = image.scale(labelImage, rectangle[2][1] - rectangle[1][1] + 1, rectangle[2][2] - rectangle[1][2] + 1, 'bicubic') | |
| label:narrow(1, 1, 1):fill(1) | |
| label:narrow(1, label:size(1), 1):fill(1) | |
| label:narrow(2, 1, 1):fill(1) | |
| label:narrow(2, label:size(2), 1):fill(1) | |
| if label:dim() == 2 or label:size(1) == 1 then | |
| label = label:repeatTensor(3, 1, 1) | |
| end | |
| for c = 1, 3 do | |
| representationImage[{c, {rectangle[1][2], rectangle[2][2]}, {rectangle[1][1], rectangle[2][1]}}][label[1]:lt(0.5)] = 0 | |
| end | |
| end | |
| --representationImage = image.drawText(representationImage, 'abs', (rectangle[1][1] + rectangle[2][1]) / 2, (rectangle[1][2] + rectangle[2][2]) / 2, {color = {0, 0, 0}, size = 2}) | |
| --representationImage[{{}, {rectangle[1][2], rectangle[2][2]}, {rectangle[1][1], rectangle[2][1]}}] = label | |
| end | |
| end | |
| local doorMask = torch.zeros(height, width) | |
| local endPointLineLength = 7 | |
| local dashLineStride = 10 | |
| for _, door in pairs(representation.doors) do | |
| local lineDim = utils.lineDim(door, lineWidth) | |
| for pointIndex = 1, 2 do | |
| for c = 1, 2 do | |
| door[pointIndex][c] = torch.round(door[pointIndex][c]) | |
| end | |
| end | |
| if lineDim > 0 then | |
| for pointIndex = 1, 2 do | |
| local point = door[pointIndex] | |
| doorMask:narrow(3 - lineDim, point[lineDim], 1):narrow(lineDim, math.max(point[3 - lineDim] - endPointLineLength, 1), math.min(endPointLineLength * 2 + 1, doorMask:size(lineDim) - (point[3 - lineDim] - endPointLineLength) + 1)):fill(1) | |
| end | |
| local fixedValue = (door[1][3 - lineDim] + door[2][3 - lineDim]) / 2 | |
| for lineValue = door[1][lineDim], door[2][lineDim], dashLineStride do | |
| doorMask:narrow(3 - lineDim, lineValue, 1):narrow(lineDim, fixedValue, 1):fill(1) | |
| end | |
| end | |
| end | |
| for i = 1, 2 do | |
| doorMask = image.dilate(doorMask) | |
| end | |
| doorMask = doorMask:byte() | |
| for c = 1, 3 do | |
| representationImage[c][doorMask] = 0 | |
| end | |
| local representationImageVideo = torch.ones(3, height, width) | |
| for c = 1, 3 do | |
| representationImageVideo[c][doorMask] = 0 | |
| end | |
| --representationImageVideo[4][doorMask] = 1 | |
| --representationImageMask[4][doorMask] = 1 | |
| image.save('test/representation_image_door.png', representationImageVideo) | |
| local pointMask = torch.zeros(height, width) | |
| for _, point in pairs(representation.points) do | |
| local x = torch.round(point[1][1]) | |
| local y = torch.round(point[1][2]) | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| pointMask[y][x] = 1 | |
| end | |
| end | |
| local largeDiskRadius = 8 | |
| local largeDiskKernel = torch.zeros(largeDiskRadius * 2 + 1, largeDiskRadius * 2 + 1) | |
| for y = 1, largeDiskRadius * 2 + 1 do | |
| for x = 1, largeDiskRadius * 2 + 1 do | |
| if ((x - largeDiskRadius - 1)^2 + (y - largeDiskRadius - 1)^2)^0.5 <= largeDiskRadius then | |
| largeDiskKernel[y][x] = 1 | |
| end | |
| end | |
| end | |
| local largeDiskMask = image.dilate(pointMask, largeDiskKernel):byte() | |
| representationImage[1][largeDiskMask] = 0.9 | |
| representationImage[2][largeDiskMask] = 0.3 | |
| representationImage[3][largeDiskMask] = 0.3 | |
| local smallDiskRadius = 4 | |
| local smallDiskKernel = torch.zeros(smallDiskRadius * 2 + 1, smallDiskRadius * 2 + 1) | |
| for y = 1, smallDiskRadius * 2 + 1 do | |
| for x = 1, smallDiskRadius * 2 + 1 do | |
| if ((x - smallDiskRadius - 1)^2 + (y - smallDiskRadius - 1)^2)^0.5 <= smallDiskRadius then | |
| smallDiskKernel[y][x] = 1 | |
| end | |
| end | |
| end | |
| local smallDiskMask = image.dilate(pointMask, smallDiskKernel):byte() | |
| for c = 1, 3 do | |
| representationImage[c][smallDiskMask] = 1 | |
| end | |
| local representationImageVideo = torch.ones(3, height, width) | |
| representationImageVideo[1][largeDiskMask] = 0.9 | |
| representationImageVideo[2][largeDiskMask] = 0.3 | |
| representationImageVideo[3][largeDiskMask] = 0.3 | |
| --representationImageVideo[4][largeDiskMask] = 1 | |
| --representationImageMask[4][largeDiskMask] = 1 | |
| for c = 1, 3 do | |
| representationImageVideo[c][smallDiskMask] = 0 | |
| end | |
| --representationImageVideo[4][smallDiskMask] = 1 | |
| --representationImageMask[4][smallDiskMask] = 1 | |
| image.save('test/representation_image_wall_junctions.png', representationImageVideo) | |
| --image.save('test/representation_image_mask.png', representationImageMask) | |
| if true then | |
| return representationImage | |
| end | |
| for c = 1, 3 do | |
| representationImage[c][lineMask:byte()] = 0 | |
| end | |
| local doorWidth = lineWidth - 2 | |
| for _, line in pairs(representation.doors) do | |
| local lineDim = utils.lineDim(line) | |
| if lineDim > 0 then | |
| local maxSize = {width, height} | |
| local rectangle = {{}, {}} | |
| rectangle[1][lineDim] = math.max(math.min(line[1][lineDim], line[2][lineDim], maxSize[lineDim]), 1) | |
| rectangle[2][lineDim] = math.min(math.max(line[1][lineDim], line[2][lineDim], 1), maxSize[lineDim]) | |
| rectangle[1][3 - lineDim] = math.max((line[1][3 - lineDim] + line[2][3 - lineDim]) / 2 - doorWidth, 1) | |
| rectangle[2][3 - lineDim] = math.min((line[1][3 - lineDim] + line[2][3 - lineDim]) / 2 + doorWidth, maxSize[3 - lineDim]) | |
| representationImage[{{}, {rectangle[1][2], rectangle[2][2]}, {rectangle[1][1], rectangle[2][1]}}] = 1 | |
| --success, newRepresentationImage = pcall(function() return image.drawRect(representationImage, rectangle[1][1], rectangle[1][2], rectangle[2][1], rectangle[2][2], {lineWidth = lineWidth, color = {0, 0, 0}}) end) | |
| --if success then | |
| --representationImage = newRepresentationImage | |
| --end | |
| end | |
| --representationImage[{{}, {minX, maxX}, {minY, maxY}}] = 0 | |
| end | |
| return representationImage | |
| end | |
| end | |
| function utils.getImageInfo(floorplan, representation, representationType, renderingType, lineWidth) | |
| local lineWidth = lineWidth or 5 | |
| for mode, items in pairs(representation) do | |
| for _, item in pairs(items) do | |
| for pointIndex = 1, 2 do | |
| for c = 1, 2 do | |
| item[pointIndex][c] = torch.round(item[pointIndex][c]) | |
| end | |
| end | |
| if mode == 'walls' or mode == 'doors' then | |
| local lineDim = utils.lineDim(item, lineWidth) | |
| local fixedValue = torch.round((item[1][3 - lineDim] + item[2][3 - lineDim]) / 2) | |
| item[1][3 - lineDim] = fixedValue | |
| item[2][3 - lineDim] = fixedValue | |
| end | |
| end | |
| end | |
| local width, height = floorplan:size(3), floorplan:size(2) | |
| local iconImages = utils.loadIconImages() | |
| local lineWidth = lineWidth or 4 | |
| if not representation.points or #representation.points == 0 then | |
| representation.points = utils.linesToPoints(width, height, representation.walls, lineWidth) | |
| end | |
| local newPoints = {} | |
| for _, point in pairs(representation.points) do | |
| if point[3][2] <= 2 then | |
| table.insert(newPoints, point) | |
| elseif point[3][2] == 3 then | |
| local orientation = point[3][3] | |
| if orientation == 1 then | |
| table.insert(newPoints, {{point[1][1], point[1][2]}, {point[2][1], point[2][2]}, {'point', 2, 3}}) | |
| table.insert(newPoints, {{point[1][1], point[1][2]}, {point[2][1], point[2][2]}, {'point', 2, 4}}) | |
| elseif orientation == 2 then | |
| table.insert(newPoints, {{point[1][1], point[1][2]}, {point[2][1], point[2][2]}, {'point', 2, 4}}) | |
| table.insert(newPoints, {{point[1][1], point[1][2]}, {point[2][1], point[2][2]}, {'point', 2, 1}}) | |
| elseif orientation == 3 then | |
| table.insert(newPoints, {{point[1][1], point[1][2]}, {point[2][1], point[2][2]}, {'point', 2, 1}}) | |
| table.insert(newPoints, {{point[1][1], point[1][2]}, {point[2][1], point[2][2]}, {'point', 2, 2}}) | |
| else | |
| table.insert(newPoints, {{point[1][1], point[1][2]}, {point[2][1], point[2][2]}, {'point', 2, 2}}) | |
| table.insert(newPoints, {{point[1][1], point[1][2]}, {point[2][1], point[2][2]}, {'point', 2, 3}}) | |
| end | |
| else | |
| for orientation = 1, 4 do | |
| table.insert(newPoints, {point[1], point[2], {'point', 2, orientation}}) | |
| end | |
| end | |
| end | |
| representation.points = newPoints | |
| local newWalls, wallPoints = utils.pointsToLines(width, height, representation.points, lineWidth, true, nil, nil, true) | |
| for _, wall in pairs(newWalls) do | |
| local lineDim = utils.lineDim(wall, lineWidth) | |
| local wallFixedValue = (wall[1][3 - lineDim] + wall[2][3 - lineDim]) / 2 | |
| local wallMinValue = math.min(wall[1][lineDim], wall[2][lineDim]) | |
| local wallMaxValue = math.max(wall[1][lineDim], wall[2][lineDim]) | |
| for _, longWall in pairs(representation.walls) do | |
| if utils.lineDim(longWall, lineWidth) == lineDim then | |
| local longWallFixedValue = (longWall[1][3 - lineDim] + longWall[2][3 - lineDim]) / 2 | |
| local longWallMinValue = math.min(longWall[1][lineDim], longWall[2][lineDim]) | |
| local longWallMaxValue = math.max(longWall[1][lineDim], longWall[2][lineDim]) | |
| if math.abs(longWallFixedValue - wallFixedValue) <= lineWidth and longWallMinValue - lineWidth <= wallMinValue and longWallMaxValue + lineWidth >= wallMaxValue then | |
| if #longWall == 4 then | |
| table.insert(newWalls[_], longWall[4]) | |
| break | |
| end | |
| end | |
| end | |
| end | |
| end | |
| local pointWalls = {} | |
| for wallIndex, points in pairs(wallPoints) do | |
| for pointIndex = 1, 2 do | |
| if not pointWalls[points[pointIndex]] then | |
| pointWalls[points[pointIndex]] = {} | |
| end | |
| table.insert(pointWalls[points[pointIndex]], {wallIndex, pointIndex}) | |
| end | |
| end | |
| local lineMask = utils.drawLineMask(width, height, representation.walls, lineWidth) | |
| local segmentation = utils.getSegmentation(width, height, representation, lineWidth)[1] | |
| local colorMap = {{224, 255, 192}, {224, 255, 192}, {255, 224, 128}, {192, 255, 255}, {192, 255, 255}, {255, 224, 224}, {192, 192, 224}, {255, 160, 96}, {192, 255, 255}, {224, 224, 224}} | |
| local borderColorMap = {{128, 192, 64}, {128, 192, 64}, {192, 128, 64}, {0, 128, 192}, {0, 128, 192}, {192, 64, 64}, {128, 64, 160}, {192, 64, 0}, {0, 128, 192}, {160, 160, 160}} | |
| table.insert(borderColorMap, {255, 255, 255}) | |
| for _, color in pairs(colorMap) do | |
| for c = 1, 3 do | |
| color[c] = color[c] / 255 | |
| end | |
| end | |
| for _, color in pairs(borderColorMap) do | |
| for c = 1, 3 do | |
| color[c] = color[c] / 255 | |
| end | |
| end | |
| local segmentImages = {} | |
| local allWallImage = torch.ones(3, height, width) | |
| local allSegmentMask = torch.zeros(height, width) | |
| for segmentIndex = 1, 10 do | |
| local segmentMask = segmentation:eq(segmentIndex) | |
| local dilatedSegmentMask = segmentMask | |
| for i = 1, lineWidth do | |
| dilatedSegmentMask = image.dilate(dilatedSegmentMask) | |
| end | |
| local borderMask = dilatedSegmentMask - segmentMask | |
| for c = 1, 3 do | |
| allWallImage[c][borderMask] = borderColorMap[segmentIndex][c] | |
| end | |
| --allWallImage[4][borderMask] = 1 | |
| --local segmentImage = torch.zeros(height, width) | |
| -- --local segmentMask = segmentation:eq(segmentIndex) | |
| --for c = 1, 3 do | |
| --segmentImage[c][segmentMask] = colorMap[segmentIndex][c] / 256 | |
| --end | |
| --segmentImage[segmentMask] = 1 | |
| table.insert(segmentImages, dilatedSegmentMask) | |
| allSegmentMask[segmentMask] = segmentIndex | |
| end | |
| --local wallRoomLabels = {} | |
| for _, wall in pairs(representation.walls) do | |
| local lineDim = utils.lineDim(wall, lineWidth) | |
| local x = (wall[1][1] + wall[2][1]) / 2 | |
| local y = (wall[1][2] + wall[2][2]) / 2 | |
| local deltas | |
| if lineDim == 1 then | |
| deltas = {0, 1} | |
| else | |
| deltas = {1, 0} | |
| end | |
| if #wall == 3 then | |
| local roomLabels = {11, 11} | |
| for direction = 1, 2 do | |
| local newX = x | |
| local newY = y | |
| for i = 1, 10 do | |
| newX = newX + deltas[1] * (direction * 2 - 3) | |
| newY = newY + deltas[2] * (direction * 2 - 3) | |
| if newX <= 0 or newX > width or newY <= 0 or newY > height then | |
| break | |
| end | |
| local segmentIndex = segmentation[newY][newX] | |
| if segmentIndex > 0 and segmentIndex <= 10 then | |
| roomLabels[direction] = segmentIndex | |
| break | |
| end | |
| end | |
| end | |
| --table.insert(wallRoomLabels, roomLabels) | |
| table.insert(representation.walls[_], roomLabels) | |
| end | |
| table.insert(representation.walls[_], wallPoints[_]) | |
| end | |
| for _, wall in pairs(representation.walls) do | |
| local lineDim = utils.lineDim(wall, lineWidth) | |
| local wallFixedValue = (wall[1][3 - lineDim] + wall[2][3 - lineDim]) / 2 | |
| local wallMinValue = math.min(wall[1][lineDim], wall[2][lineDim]) | |
| local wallMaxValue = math.max(wall[1][lineDim], wall[2][lineDim]) | |
| local doorIndices = {} | |
| for doorIndex, door in pairs(representation.doors) do | |
| if utils.lineDim(door, lineWidth) == lineDim then | |
| local doorFixedValue = (door[1][3 - lineDim] + door[2][3 - lineDim]) / 2 | |
| local doorMinValue = math.min(door[1][lineDim], door[2][lineDim]) | |
| local doorMaxValue = math.max(door[1][lineDim], door[2][lineDim]) | |
| if math.abs(doorFixedValue - wallFixedValue) <= lineWidth and math.max(doorMinValue, wallMinValue) < math.min(doorMaxValue, wallMaxValue) then | |
| table.insert(doorIndices, doorIndex) | |
| end | |
| end | |
| end | |
| table.insert(representation.walls[_], doorIndices) | |
| end | |
| -- local wallImages = {} | |
| -- local wallNeighborSegments = {} | |
| -- for _, wall in pairs(representation.walls) do | |
| -- local lineDim = utils.lineDim(wall, lineWidth) | |
| -- local fixedValue = (wall[1][3 - lineDim] + wall[2][3 - lineDim]) / 2 | |
| -- local minValue = math.min(wall[1][lineDim], wall[2][lineDim]) | |
| -- local maxValue = math.max(wall[1][lineDim], wall[2][lineDim]) | |
| -- table.insert(wallImages, allWallImage:narrow(4 - lineDim, minValue, maxValue - minValue + 1):narrow(1 + lineDim, fixedValue - lineWidth, lineWidth * 2 + 1)) | |
| -- wall[1][lineDim] = minValue | |
| -- wall[2][lineDim] = maxValue | |
| -- wall[1][3 - lineDim] = fixedValue | |
| -- wall[2][3 - lineDim] = fixedValue | |
| -- local segmentMasks = torch.zeros(2, 10, maxValue - minValue - 1) | |
| -- local deltas = {-1, 1} | |
| -- for value = minValue + 1, maxValue - 1 do | |
| -- for direction = 1, 2 do | |
| -- local segmentIndex | |
| -- for i = 1, 10 do | |
| -- if lineDim == 1 then | |
| -- if fixedValue + deltas[direction] * i <= 0 or fixedValue + deltas[direction] * i > height then | |
| -- break | |
| -- end | |
| -- segmentIndex = allSegmentMask[fixedValue + deltas[direction] * i][value] | |
| -- if segmentIndex > 0 then | |
| -- break | |
| -- end | |
| -- elseif lineDim == 2 then | |
| -- if fixedValue + deltas[direction] * i <= 0 or fixedValue + deltas[direction] * i > width then | |
| -- break | |
| -- end | |
| -- segmentIndex = allSegmentMask[value][fixedValue + deltas[direction] * i] | |
| -- if segmentIndex > 0 then | |
| -- break | |
| -- end | |
| -- end | |
| -- end | |
| -- if segmentIndex > 0 then | |
| -- segmentMasks[direction][segmentIndex][value - minValue] = 1 | |
| -- end | |
| -- end | |
| -- end | |
| -- local neighborSegments = {{}, {}} | |
| -- for direction = 1, 2 do | |
| -- for segmentIndex = 1, 10 do | |
| -- if segmentMasks[direction][segmentIndex]:sum() > 0 then | |
| -- local segmentMask = segmentMasks[direction][segmentIndex] | |
| -- segmentMask = segmentMask:repeatTensor(1, 1) | |
| -- if lineDim == 2 then | |
| -- segmentMask = segmentMask:transpose(1, 2) | |
| -- end | |
| -- table.insert(neighborSegments[direction], {segmentIndex, segmentMask:byte()}) | |
| -- end | |
| -- end | |
| -- end | |
| -- table.insert(wallNeighborSegments, neighborSegments) | |
| -- end | |
| -- for iconIndex, icon in pairs(representation.icons) do | |
| -- local neighborWalls = {} | |
| -- local center = {(icon[1][1] + icon[2][1]) / 2, (icon[1][2] + icon[2][2]) / 2} | |
| -- local size = {icon[2][1] - icon[1][1], icon[2][2] - icon[1][2]} | |
| -- for _, wall in pairs(representation.walls) do | |
| -- local lineDim = utils.lineDim(wall, lineWidth) | |
| -- local wallFixedValue = (wall[1][3 - lineDim] + wall[2][3 - lineDim]) / 2 | |
| -- local wallMinValue = math.min(wall[1][lineDim], wall[2][lineDim]) | |
| -- local wallMaxValue = math.max(wall[1][lineDim], wall[2][lineDim]) | |
| -- if wallMinValue < center[lineDim] + size[lineDim] / 2 and wallMaxValue > center[lineDim] - size[lineDim] / 2 then | |
| -- end | |
| -- end | |
| -- end | |
| for __, item in pairs(representation.icons) do | |
| item[3] = utils.getItemInfo('icons', utils.getNumber('icons', item[3])) | |
| local rectangle = {{math.min(item[1][1], item[2][1]), math.min(item[1][2], item[2][2])}, {math.max(item[1][1], item[2][1]), math.max(item[1][2], item[2][2])}} | |
| for pointIndex = 1, 2 do | |
| rectangle[pointIndex][1] = torch.round(math.max(math.min(rectangle[pointIndex][1], width), 1)) | |
| rectangle[pointIndex][2] = torch.round(math.max(math.min(rectangle[pointIndex][2], height), 1)) | |
| end | |
| local minDistance = math.max(width, height) | |
| local minDistanceOrientation | |
| if (math.abs(rectangle[2][1] - rectangle[1][1]) > math.abs(rectangle[2][2] - rectangle[1][2]) and item[3][1] ~= 'toilet' and item[3][1] ~= 'stairs') or (math.abs(rectangle[2][1] - rectangle[1][1]) < math.abs(rectangle[2][2] - rectangle[1][2]) and (item[3][1] == 'toilet' or item[3][1] == 'stairs')) or item[3][1] == 'washing_basin' then | |
| local min = math.min(rectangle[1][2], rectangle[2][2]) | |
| local max = math.max(rectangle[1][2], rectangle[2][2]) | |
| local center = (rectangle[1][1] + rectangle[2][1]) / 2 | |
| local deltaMin = 0 | |
| for delta = 1, min - 1 do | |
| deltaMin = delta | |
| if lineMask[min - delta][center] == 1 then | |
| break | |
| end | |
| end | |
| local deltaMax = 0 | |
| for delta = 1, height - max do | |
| deltaMax = delta | |
| if lineMask[max + delta][center] == 1 then | |
| break | |
| end | |
| end | |
| if deltaMax < minDistance then | |
| minDistance = deltaMax | |
| minDistanceOrientation = 1 | |
| end | |
| if deltaMin < minDistance then | |
| minDistance = deltaMin | |
| minDistanceOrientation = 3 | |
| end | |
| end | |
| if (math.abs(rectangle[2][1] - rectangle[1][1]) < math.abs(rectangle[2][2] - rectangle[1][2]) and item[3][1] ~= 'toilet') or (math.abs(rectangle[2][1] - rectangle[1][1]) > math.abs(rectangle[2][2] - rectangle[1][2]) and item[3][1] == 'toilet') or item[3][1] == 'washing_basin' then | |
| local min = math.min(rectangle[1][1], rectangle[2][1]) | |
| local max = math.max(rectangle[1][1], rectangle[2][1]) | |
| local center = (rectangle[1][2] + rectangle[2][2]) / 2 | |
| local deltaMin = 0 | |
| for delta = 1, min - 1 do | |
| deltaMin = delta | |
| if lineMask[center][min - delta] == 1 then | |
| break | |
| end | |
| end | |
| local deltaMax = 0 | |
| for delta = 1, width - max do | |
| deltaMax = delta | |
| if lineMask[center][max + delta] == 1 then | |
| break | |
| end | |
| end | |
| if deltaMax < minDistance then | |
| minDistance = deltaMax | |
| minDistanceOrientation = 4 | |
| end | |
| if deltaMin < minDistance then | |
| minDistance = deltaMin | |
| minDistanceOrientation = 2 | |
| end | |
| end | |
| item[3][3] = minDistanceOrientation | |
| end | |
| local rooms, numRooms = utils.findConnectedComponents(1 - lineMask) | |
| rooms = segmentation:int() * numRooms + rooms:int() | |
| local newLabels = {} | |
| local roomLabelMap = {} | |
| for __, item in pairs(representation.labels) do | |
| --local rectangle = {item[1], item[2]} | |
| local cx = (item[1][1] + item[2][1]) / 2 | |
| local cy = (item[1][2] + item[2][2]) / 2 | |
| local labelWidth = 80 | |
| local labelHeight = 30 | |
| local roomIndex = rooms[torch.round(cy)][torch.round(cx)] | |
| if not roomLabelMap[roomIndex] then | |
| roomLabelMap[roomIndex] = {} | |
| end | |
| local number = utils.getNumber('labels', item[3]) | |
| if roomIndex > 0 and not roomLabelMap[roomIndex][number] then | |
| roomLabelMap[roomIndex][number] = true | |
| local roomIndices = rooms:eq(roomIndex):nonzero() | |
| local mins = torch.min(roomIndices, 1)[1] | |
| local maxs = torch.max(roomIndices, 1)[1] | |
| local orientation = 1 | |
| if maxs[2] - mins[2] + 1 < labelWidth * 0.8 then | |
| if math.min((maxs[2] - mins[2] + 1) / labelHeight, (maxs[1] - mins[1] + 1) / labelWidth) > (maxs[2] - mins[2] + 1) / labelWidth then | |
| labelWidth, labelHeight = labelHeight, labelWidth | |
| orientation = 2 | |
| end | |
| if maxs[2] - mins[2] + 1 < labelWidth then | |
| local ratio = math.max((maxs[2] - mins[2] + 1) / labelWidth, 0.5) | |
| labelWidth, labelHeight = labelWidth * ratio, labelHeight * ratio | |
| end | |
| end | |
| if math.max(labelWidth, labelHeight) > 10 then | |
| local cx = (mins[2] + maxs[2]) / 2 | |
| local cy = (mins[1] + maxs[1]) / 2 | |
| local rectangle = {{cx - labelWidth / 2, cy - labelHeight / 2}, {cx + labelWidth / 2, cy + labelHeight / 2}} | |
| for pointIndex = 1, 2 do | |
| rectangle[pointIndex][1] = torch.round(math.max(math.min(rectangle[pointIndex][1], width), 1)) | |
| rectangle[pointIndex][2] = torch.round(math.max(math.min(rectangle[pointIndex][2], height), 1)) | |
| end | |
| local label = rectangle | |
| table.insert(label, {item[3][1], item[3][2], orientation}) | |
| table.insert(newLabels, label) | |
| end | |
| end | |
| end | |
| representation.labels = newLabels | |
| local doorMask = torch.zeros(height, width) | |
| local endPointLineLength = 7 | |
| local dashLineStride = 10 | |
| for _, door in pairs(representation.doors) do | |
| local lineDim = utils.lineDim(door, lineWidth) | |
| for pointIndex = 1, 2 do | |
| for c = 1, 2 do | |
| door[pointIndex][c] = torch.round(door[pointIndex][c]) | |
| end | |
| end | |
| for pointIndex = 1, 2 do | |
| local point = door[pointIndex] | |
| doorMask:narrow(3 - lineDim, point[lineDim], 1):narrow(lineDim, math.max(point[3 - lineDim] - endPointLineLength, 1), math.min(endPointLineLength * 2 + 1, doorMask:size(lineDim) - (point[3 - lineDim] - endPointLineLength) + 1)):fill(1) | |
| end | |
| local fixedValue = (door[1][3 - lineDim] + door[2][3 - lineDim]) / 2 | |
| for lineValue = door[1][lineDim], door[2][lineDim], dashLineStride do | |
| doorMask:narrow(3 - lineDim, lineValue, 1):narrow(lineDim, fixedValue, 1):fill(1) | |
| end | |
| end | |
| for i = 1, 2 do | |
| doorMask = image.dilate(doorMask) | |
| end | |
| doorMask = doorMask:byte() | |
| local allDoorImage = torch.zeros(4, height, width) | |
| for c = 1, 3 do | |
| allDoorImage[c][doorMask] = 0 | |
| end | |
| allDoorImage[4][doorMask] = 1 | |
| for _, door in pairs(representation.doors) do | |
| local lineDim = utils.lineDim(door, lineWidth) | |
| local fixedValue = (door[1][3 - lineDim] + door[2][3 - lineDim]) / 2 | |
| local minValue = math.min(door[1][lineDim], door[2][lineDim]) | |
| local maxValue = math.max(door[1][lineDim], door[2][lineDim]) | |
| --table.insert(doorImages, allDoorImage:narrow(4 - lineDim, minValue, maxValue - minValue + 1):narrow(1 + lineDim, fixedValue - endPointLineLength, endPointLineLength * 2 + 1)) | |
| door[1][lineDim] = minValue | |
| door[2][lineDim] = maxValue | |
| door[1][3 - lineDim] = fixedValue | |
| door[2][3 - lineDim] = fixedValue | |
| local doorImage = doorMask:narrow(3 - lineDim, minValue, maxValue - minValue + 1):narrow(lineDim, fixedValue - endPointLineLength, endPointLineLength * 2 + 1) | |
| table.insert(representation.doors[_], doorImage) | |
| end | |
| local pointMask = torch.zeros(height, width) | |
| for _, point in pairs(representation.points) do | |
| local x = torch.round(point[1][1]) | |
| local y = torch.round(point[1][2]) | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| pointMask[y][x] = 1 | |
| end | |
| end | |
| local largeDiskRadius = 8 | |
| local largeDiskKernel = torch.zeros(largeDiskRadius * 2 + 1, largeDiskRadius * 2 + 1) | |
| for y = 1, largeDiskRadius * 2 + 1 do | |
| for x = 1, largeDiskRadius * 2 + 1 do | |
| if ((x - largeDiskRadius - 1)^2 + (y - largeDiskRadius - 1)^2)^0.5 <= largeDiskRadius then | |
| largeDiskKernel[y][x] = 1 | |
| end | |
| end | |
| end | |
| local largeDiskMask = image.dilate(pointMask, largeDiskKernel):byte() | |
| local smallDiskRadius = 4 | |
| local smallDiskKernel = torch.zeros(smallDiskRadius * 2 + 1, smallDiskRadius * 2 + 1) | |
| for y = 1, smallDiskRadius * 2 + 1 do | |
| for x = 1, smallDiskRadius * 2 + 1 do | |
| if ((x - smallDiskRadius - 1)^2 + (y - smallDiskRadius - 1)^2)^0.5 <= smallDiskRadius then | |
| smallDiskKernel[y][x] = 1 | |
| end | |
| end | |
| end | |
| local smallDiskMask = image.dilate(pointMask, smallDiskKernel):byte() | |
| local allPointImage = torch.zeros(4, height, width) | |
| allPointImage[1][largeDiskMask] = 0.9 | |
| allPointImage[2][largeDiskMask] = 0.3 | |
| allPointImage[3][largeDiskMask] = 0.3 | |
| allPointImage[4][largeDiskMask] = 1 | |
| for c = 1, 3 do | |
| allPointImage[c][smallDiskMask] = 1 | |
| end | |
| local point = representation.points[1][1] | |
| local pointImage = allPointImage:narrow(2, point[2] - largeDiskRadius, largeDiskRadius * 2 + 1):narrow(3, point[1] - largeDiskRadius, largeDiskRadius * 2 + 1) | |
| -- local pointImages = {} | |
| -- for _, point in pairs(representation.points) do | |
| -- point = point[1] | |
| -- table.insert(pointImages, allPointImage:narrow(2, point[2] - largeDiskRadius, largeDiskRadius * 2 + 1):narrow(3, point[1] - largeDiskRadius, largeDiskRadius * 2 + 1)) | |
| -- end | |
| return {pointImage = pointImage, segmentImages = segmentImages, colorMap = colorMap, borderColorMap = borderColorMap, pointWalls = pointWalls} | |
| end | |
| function utils.loadRepresentation(filename) | |
| local representationExists, representationInfo = pcall(function() | |
| return csvigo.load({path=filename, mode="large", header=false, separator='\t', verbose=false}) | |
| end) | |
| if representationExists and representationInfo ~= nil then | |
| local representation = {} | |
| representation.walls = {} | |
| representation.doors = {} | |
| representation.icons = {} | |
| representation.labels = {} | |
| local modeMap = utils.modeMap() | |
| for _, itemInfo in pairs(representationInfo) do | |
| if itemInfo[5] ~= 'point' then | |
| local itemMode = modeMap[itemInfo[5]] | |
| local item = {{tonumber(itemInfo[1]), tonumber(itemInfo[2])}, {tonumber(itemInfo[3]), tonumber(itemInfo[4])}, {itemInfo[5], tonumber(itemInfo[6]), tonumber(itemInfo[7])}} | |
| if itemMode == 'icons' or itemMode == 'labels' or utils.lineDim(item) > 0 then | |
| item[1][1], item[2][1] = math.min(item[1][1], item[2][1]), math.max(item[1][1], item[2][1]) | |
| item[1][2], item[2][2] = math.min(item[1][2], item[2][2]), math.max(item[1][2], item[2][2]) | |
| end | |
| table.insert(representation[itemMode], item) | |
| end | |
| end | |
| return representation | |
| else | |
| return nil | |
| end | |
| end | |
| function utils.loadItems(filename) | |
| local representationExists, representationInfo = pcall(function() | |
| return csvigo.load({path=filename, mode="large", header=false, separator='\t', verbose=false}) | |
| --return io.open(filename, 'r') | |
| end) | |
| if representationExists and representationInfo ~= nil then | |
| --representationInfo = csvigo.load({path=filename, mode="large", header=false, separator='\t', verbose=false}) | |
| local items = {} | |
| for _, item in pairs(representationInfo) do | |
| local itemInfo = {{tonumber(item[1]), tonumber(item[2])}, {tonumber(item[3]), tonumber(item[4])}, {item[5], tonumber(item[6]), tonumber(item[7])}} | |
| --print(item[5]) | |
| --print(itemMode) | |
| table.insert(items, itemInfo) | |
| end | |
| return items | |
| else | |
| return {} | |
| end | |
| end | |
| function utils.finalizeRepresentation(representation) | |
| if representation.walls then | |
| representation.walls = utils.stitchLines(utils.sortLines(representation.walls), 5) | |
| if representation.doors then | |
| representation.doors = utils.fixedDoors(utils.sortLines(representation.doors), representation.walls, 5) | |
| end | |
| end | |
| return representation | |
| end | |
| function utils.saveRepresentation(filename, representation, ratio) | |
| --local representationExists, representationFile = pcall(function() | |
| --return io.open(filename, 'r') | |
| --end) | |
| --local override = false | |
| --if representationExists then | |
| --end | |
| print('save representation') | |
| if ratio ~= nil and ratio ~= 1 then | |
| representation = utils.scaleRepresentationByRatio(representation, ratio) | |
| end | |
| if filename:match('(.+)/(.+)') ~= nil then | |
| pl.dir.makepath(filename:match('(.+)/(.+)')) | |
| end | |
| local representationFile = io.open(filename, 'w') | |
| representation.points = nil | |
| for itemMode, items in pairs(representation) do | |
| for _, item in pairs(items) do | |
| for __, field in pairs(item) do | |
| if __ <= 3 then | |
| for ___, value in pairs(field) do | |
| representationFile:write(value .. '\t') | |
| end | |
| end | |
| end | |
| representationFile:write('\n') | |
| end | |
| end | |
| representationFile:close() | |
| end | |
| function utils.printRepresentation(representation) | |
| for itemMode, items in pairs(representation) do | |
| for _, item in pairs(items) do | |
| --for __, field in pairs(item) do | |
| --if __ <= 3 then | |
| --for ___, value in pairs(field) do | |
| --end | |
| --end | |
| --end | |
| print('(' .. item[1][1] .. ', ' .. item[1][2] .. ')\t(' .. item[2][1] .. ', ' .. item[2][2] .. ')\t' .. item[3][1] .. ' ' .. item[3][2] .. ' ' .. item[3][3]) | |
| print('\n') | |
| end | |
| end | |
| end | |
| function utils.predictRepresentation(modelPath, floorplan) | |
| if utils.model == nil then | |
| utils.model = torch.load(modelPath) | |
| end | |
| local width = floorplan:size(3) | |
| local height = floorplan:size(2) | |
| local sampleDim = 256 | |
| --local floorplanScaled = image.scale(floorplan, sampleDim, sampleDim) | |
| local lineWidth = lineWidth or 5 | |
| local model = utils.model | |
| model:evaluate() | |
| package.path = '../code/datasets/?.lua;' .. package.path | |
| package.path = '../code/?.lua;' .. package.path | |
| local dataset = require('floorplan-representation') | |
| local floorplanNormalized = dataset:preprocessResize(sampleDim, sampleDim)(floorplan) | |
| local input = floorplanNormalized:repeatTensor(1, 1, 1, 1):cuda() | |
| --local input = floorplanScaled:repeatTensor(1, 1, 1, 1):cuda() | |
| local output = model.modules[1]:forward(input) | |
| local prob, pred = torch.max(output[2]:double(), 2) | |
| pred = pred[{{}, 1}]:view(-1, 8, 8, 8):transpose(3, 4):transpose(2, 3):double() | |
| local outputRepresentation = torch.cat(output[1]:double(), pred, 2) | |
| local representationTensor = outputRepresentation[1] | |
| local representation = utils.convertTensorToRepresentation(sampleDim, sampleDim, representationTensor, 0.5) | |
| local representationUnscaled = utils.scaleRepresentation(representation, sampleDim, sampleDim, width, height) | |
| local representationGeneral = utils.convertRepresentationToGeneral(width, height, representationUnscaled, 'P', lineWidth) | |
| return representationGeneral | |
| end | |
| function utils.predictSegmentation(floorplan, walls, preview) | |
| if utils.model == nil then | |
| utils.model = torch.load('/home/chenliu/Projects/Floorplan/models/segmentation/model_best.t7') | |
| end | |
| local width = floorplan:size(3) | |
| local height = floorplan:size(2) | |
| local sampleDim = 256 | |
| --local floorplanScaled = image.scale(floorplan, sampleDim, sampleDim) | |
| local lineWidth = lineWidth or 5 | |
| local model = utils.model | |
| --model:evaluate() | |
| --local datasets = require 'datasets/init' | |
| --local opts = require '../../InverseCAD/opts' | |
| --local opt = opts.parse(arg) | |
| package.path = '../InverseCAD/datasets/?.lua;' .. package.path | |
| package.path = '../InverseCAD/?.lua;' .. package.path | |
| --local dataset = require('/home/chenliu/Projects/Floorplan/floorplan/InverseCAD/datasets/floorplan-representation')(nil, nil, 'val') | |
| --local dataset = require('floorplan-representation')(nil, nil, 'val') | |
| --local dataset = require 'floorplan-representation'({}, nil, 'val') | |
| local dataset = require('floorplan-representation') | |
| local floorplanNormalized = dataset:preprocessResize(sampleDim, sampleDim)(floorplan) | |
| local input = floorplanNormalized:repeatTensor(1, 1, 1, 1):cuda() | |
| --print(#input) | |
| local output = model:forward(input) | |
| --print(output[1]) | |
| local prob, pred = torch.max(output:double(), 2) | |
| local prediction = pred[{{}, 1}]:view((#input)[1], sampleDim, sampleDim)[1]:double() | |
| prediction = image.scale(prediction, width, height, 'simple') | |
| local segmentationImage = utils.drawSegmentation(prediction, torch.max(prediction)) | |
| --segmentationImage = image.scale(segmentationImage, width, height) | |
| if preview then | |
| return segmentationImage | |
| end | |
| local wallMask | |
| if walls then | |
| wallMask = utils.drawLineMask(width, height, walls, lineWidth) | |
| else | |
| wallMask = prediction:eq(25):double() | |
| end | |
| local rooms, numRooms = utils.findConnectedComponents(1 - wallMask) | |
| --image.save('test/rooms.png', utils.drawSegmentation(rooms)) | |
| local indexMask = (rooms * 100 + prediction:int()) | |
| local indices = indexMask:view(-1) | |
| local components = {} | |
| for i = 1, indices:size(1) do | |
| components[indices[i]] = true | |
| end | |
| local markers = torch.zeros(height, width) | |
| local markerIndex = 2 | |
| local markerLabelMap = {} | |
| for index, _ in pairs(components) do | |
| local label = index % 100 | |
| if index > 100 and label < 24 then | |
| local mask = indexMask:eq(index) | |
| local numErosions | |
| if label <= 10 then | |
| numErosions = 3 | |
| else | |
| numErosions = 3 | |
| end | |
| for i = 1, numErosions do | |
| mask = image.erode(mask:double()) | |
| end | |
| mask = mask:gt(0) | |
| if ##mask:nonzero() > 0 then | |
| markers[mask] = markerIndex | |
| --local indices = mask:nonzero() | |
| --local mean = torch.round(torch.mean(indices:double(), 1)[1]) | |
| --markers[mean[1]][mean[2]] = markerIndex | |
| --markers[{{mean[1] - 2, mean[1] + 2}, {mean[2] - 2, mean[2] + 2}}] = markerIndex | |
| markerLabelMap[markerIndex] = label | |
| markerIndex = markerIndex + 1 | |
| end | |
| end | |
| end | |
| --markers[rooms:eq(rooms[1][1])] = markerIndex | |
| local backgroundMask = rooms:eq(rooms[1][1]) | |
| wallMask = wallMask:byte() | |
| markers[backgroundMask] = -1 | |
| markers[wallMask] = -1 | |
| markers = markers:int() | |
| cv.watershed{(floorplan:transpose(1, 2):transpose(2, 3) * 255):byte(), markers} | |
| markers[backgroundMask] = 1 | |
| markers[wallMask] = 0 | |
| --image.save('test/markers.png', utils.drawSegmentation(markers)) | |
| --os.exit(1) | |
| return markers, markerLabelMap | |
| --return utils.drawSegmentation(markers:double()) | |
| --return segmentationImage | |
| end | |
| function utils.predictWalls(floorplan) | |
| if utils.model == nil then | |
| utils.model = torch.load('/home/chenliu/Projects/Floorplan/models/wall/model_best.t7') | |
| utils.model.modules[3] = nil | |
| collectgarbage() | |
| end | |
| local width = floorplan:size(3) | |
| local height = floorplan:size(2) | |
| local sampleDim = 256 | |
| --local floorplanScaled = image.scale(floorplan, sampleDim, sampleDim) | |
| local lineWidth = lineWidth or 5 | |
| local model = utils.model | |
| --model:evaluate() | |
| --local datasets = require 'datasets/init' | |
| --local opts = require '../../InverseCAD/opts' | |
| --local opt = opts.parse(arg) | |
| package.path = '../InverseCAD/datasets/?.lua;' .. package.path | |
| package.path = '../InverseCAD/?.lua;' .. package.path | |
| --local dataset = require('/home/chenliu/Projects/Floorplan/floorplan/InverseCAD/datasets/floorplan-representation')(nil, nil, 'val') | |
| --local dataset = require('floorplan-representation')(nil, nil, 'val') | |
| --local dataset = require 'floorplan-representation'({}, nil, 'val') | |
| local dataset = require('floorplan-representation') | |
| local floorplanNormalized = dataset:preprocessResize(sampleDim, sampleDim)(floorplan) | |
| local input = floorplanNormalized:repeatTensor(1, 1, 1, 1):cuda() | |
| --print(#input) | |
| local output = model:forward(input) | |
| print(#output[1]) | |
| local wallMask = output[1][1]:gt(0.5):double() | |
| --print(#wallMask) | |
| --image.save('test/wall.png', wallMask) | |
| wallMask = image.scale(wallMask, width, height, 'simple') | |
| --image.save('test/wall.png', wallMask) | |
| --os.exit(1) | |
| return wallMask:gt(0.5) | |
| end | |
| function utils.segmentFloorplan(floorplan, binaryThreshold, numOpenOperations, reverse, considerEdge) | |
| --local floorplanBinary = torch.ones((#floorplan)[2], (#floorplan)[3]) | |
| --for c = 1, 3 do | |
| --local mask = floorplan[c]:lt(binaryThreshold):double() | |
| --floorplanBinary = torch.cmul(floorplanBinary, mask) | |
| --end | |
| local floorplanBinary = floorplan:max(1)[1]:lt(binaryThreshold):double() | |
| if considerEdge then | |
| local floorplanGray = floorplan:mean(1)[1] | |
| local horizontalKernel = torch.zeros(3, 3) | |
| local verticalKernel = torch.zeros(3, 3) | |
| for c = 1, 3 do | |
| horizontalKernel[{2, c}] = 1 | |
| horizontalKernel[{1, c}] = -0.5 | |
| horizontalKernel[{3, c}] = -0.5 | |
| verticalKernel[{c, 2}] = 1 | |
| verticalKernel[{c, 1}] = -0.5 | |
| verticalKernel[{c, 3}] = -0.5 | |
| end | |
| local horizontalEdgeMap = image.convolve(floorplanGray, horizontalKernel, 'same'):lt(-0.1):double() | |
| local verticalEdgeMap = image.convolve(floorplanGray, verticalKernel, 'same'):lt(-0.1):double() | |
| --horizontalEdgeMap = image.erode(image.dilate(horizontalEdgeMap)) | |
| --verticalEdgeMap = image.erode(image.dilate(verticalEdgeMap)) | |
| --image.save('test/edge_1.png', horizontalEdgeMap) | |
| --image.save('test/edge_2.png', verticalEdgeMap) | |
| floorplanBinary = (floorplanBinary + horizontalEdgeMap + verticalEdgeMap):gt(0):double() | |
| end | |
| --floorplanBinary = image.dilate(floorplanBinary) | |
| local kernel = torch.ones(3, 3) | |
| if numOpenOperations > 0 then | |
| for i = 1, numOpenOperations do | |
| floorplanBinary = image.erode(floorplanBinary) | |
| floorplanBinary = image.dilate(floorplanBinary) | |
| end | |
| elseif numOpenOperations < 0 then | |
| for i = 1, -numOpenOperations do | |
| floorplanBinary = image.dilate(floorplanBinary) | |
| floorplanBinary = image.erode(floorplanBinary) | |
| end | |
| end | |
| floorplanBinaryByte = (floorplanBinary * 255):byte() | |
| if reverse ~= nil and reverse then | |
| floorplanBinaryByte = 255 - floorplanBinaryByte | |
| end | |
| local floorplanComponent = torch.IntTensor(floorplanBinaryByte:size()) | |
| local numComponents = cv.connectedComponents{255 - floorplanBinaryByte, floorplanComponent} | |
| floorplanComponent = floorplanComponent + 1 | |
| return floorplanComponent, numComponents, floorplanBinary | |
| end | |
| function utils.findWalls(floorplan, denotedWalls, lineWidth, lineLength) | |
| local width = floorplan:size(3) | |
| local height = floorplan:size(2) | |
| local lineWidth = lineWidth or 5 | |
| local binaryThreshold = 0.7 | |
| local minLength = lineWidth * 2 or 10 | |
| if denotedWalls == nil then | |
| local floorplanSegmentationReversed, numSegmentsReversed, floorplanBinaryReversed = utils.segmentFloorplan(floorplan, binaryThreshold, 0, true) | |
| local backgroundSegmentReversed = floorplanSegmentationReversed[1][1] | |
| local maxCount = 0 | |
| local borderSegmentReversed | |
| for segmentReversed = 1, numSegmentsReversed do | |
| if segmentReversed ~= backgroundSegmentReversed then | |
| local indices = floorplanSegmentationReversed:eq(segmentReversed):nonzero() | |
| if ##indices > 0 then | |
| local count = (#indices)[1] | |
| if count > maxCount then | |
| borderSegmentReversed = segmentReversed | |
| maxCount = count | |
| end | |
| end | |
| end | |
| end | |
| if borderSegmentReversed then | |
| local borderMask = floorplanSegmentationReversed:eq(borderSegmentReversed) | |
| binaryThreshold = torch.max(floorplan, 1)[1][borderMask]:max() | |
| end | |
| else | |
| local denotedWallMask = torch.zeros(height, width) | |
| for _, wall in pairs(denotedWalls) do | |
| local mins = {} | |
| local maxs = {} | |
| for c = 1, 2 do | |
| mins[c] = math.min(wall[1][c], wall[2][c]) | |
| maxs[c] = math.max(wall[1][c], wall[2][c]) | |
| end | |
| denotedWallMask[{{mins[2], max[2]}, {mins[1], max[1]}}] = 1 | |
| end | |
| binaryThreshold = torch.max(floorplan, 1)[1][denotedWallMask]:mean() + 0.1 | |
| end | |
| local floorplanSegmentation, numSegments, floorplanBinary = utils.segmentFloorplan(floorplan, binaryThreshold, 0) | |
| floorplanSegmentation = floorplanSegmentation - 1 | |
| numSegments = numSegments - 1 | |
| --[[ | |
| local floorplanSegmentationReversed, numSegmentsReversed, floorplanBinaryReversed = utils.segmentFloorplan(floorplan, binaryThreshold, 0, true) | |
| local backgroundSegmentReversed = floorplanSegmentationReversed[1][1] | |
| if partialWallMask ~= nil then | |
| floorplanSegmentationReversed = torch.cmul(floorplanSegmentationReversed, partialWallMask) | |
| end | |
| local maxCount = 0 | |
| local borderSegmentReversed | |
| for segmentReversed = 1, numSegmentsReversed do | |
| if segmentReversed ~= backgroundSegmentReversed then | |
| local indices = floorplanSegmentationReversed:eq(segmentReversed):nonzero() | |
| if ##indices > 0 then | |
| local count = (#indices)[1] | |
| if count > maxCount then | |
| borderSegmentReversed = segmentReversed | |
| maxCount = count | |
| end | |
| end | |
| end | |
| end | |
| if borderSegmentReversed == nil then | |
| return torch.zeros(#floorplan) | |
| end | |
| local borderMask = floorplanSegmentationReversed:eq(borderSegmentReversed) | |
| ]]-- | |
| --local borderMask = floorplanSegmentationReversed:gt(1) | |
| --local floorplanSegmentation = utils.predictSegmentation(floorplan) | |
| --local numSegments = torch.max(floorplanSegmentation) | |
| local borderMask = floorplanSegmentation:eq(0) | |
| --for i = 1, torch.round((math.max(width, height) / 256 - 1) * 5) do | |
| --for i = 1, 2 do | |
| --borderMask = image.erode(borderMask) | |
| --end | |
| --[[ | |
| local doorWidth = 5 | |
| for i = 1, doorWidth do | |
| borderMask = image.dilate(borderMask) | |
| end | |
| for i = 1, doorWidth do | |
| borderMask = image.erode(borderMask) | |
| end | |
| ]]-- | |
| local lineLengthMask = torch.zeros(4, height, width) | |
| for x = 1, width do | |
| local length = 0 | |
| for y = 1, height do | |
| if borderMask[y][x] == 1 then | |
| length = length + 1 | |
| lineLengthMask[1][y][x] = length | |
| else | |
| length = 0 | |
| end | |
| end | |
| length = 0 | |
| for y = height, 1, -1 do | |
| if borderMask[y][x] == 1 then | |
| length = length + 1 | |
| lineLengthMask[3][y][x] = length | |
| else | |
| length = 0 | |
| end | |
| end | |
| end | |
| for y = 1, height do | |
| local length = 0 | |
| for x = 1, width do | |
| if borderMask[y][x] == 1 then | |
| length = length + 1 | |
| lineLengthMask[4][y][x] = length | |
| else | |
| length = 0 | |
| end | |
| end | |
| length = 0 | |
| for x = width, 1, -1 do | |
| if borderMask[y][x] == 1 then | |
| length = length + 1 | |
| lineLengthMask[2][y][x] = length | |
| else | |
| length = 0 | |
| end | |
| end | |
| end | |
| local minLineWidth = 5 | |
| local maxLineWidth = 15 | |
| local minLineLength = 20 | |
| --local maxLineWidth = 15 | |
| --local minLineLength = 20 | |
| local horizontalLineLengthMask = lineLengthMask[2] + lineLengthMask[4] - 1 | |
| local verticalLineLengthMask = lineLengthMask[1] + lineLengthMask[3] - 1 | |
| local lineTypeMask = torch.zeros(height, width) | |
| lineTypeMask[torch.cmul(torch.cmul(verticalLineLengthMask:ge(minLineWidth), verticalLineLengthMask:le(maxLineWidth)), horizontalLineLengthMask:ge(minLineLength))] = 1 | |
| lineTypeMask[torch.cmul(torch.cmul(horizontalLineLengthMask:ge(minLineWidth), horizontalLineLengthMask:le(maxLineWidth)), verticalLineLengthMask:ge(minLineLength))] = 2 | |
| --local minLineWidth = 5 | |
| local maxDoorWidth = 8 | |
| local minDoorLength = 15 | |
| local backgroundSegment = floorplanSegmentation[1][1] | |
| for segment = 1, numSegments do | |
| if segment ~= backgroundSegment then | |
| local mask = floorplanSegmentation:eq(segment) | |
| local indices = floorplanSegmentation:eq(segment):nonzero() | |
| if ##indices > 0 then | |
| local mins = torch.min(indices, 1)[1] | |
| local maxs = torch.max(indices, 1)[1] | |
| for c = 1, 2 do | |
| if maxs[c] - mins[c] <= maxDoorWidth and maxs[3 - c] - mins[3 - c] >= minDoorLength then | |
| lineTypeMask[mask] = c | |
| borderMask[mask] = 1 | |
| end | |
| end | |
| end | |
| end | |
| end | |
| local pointLineTypeMask = torch.zeros(4, height, width) | |
| for x = 1, width do | |
| local lineExists = false | |
| for y = 1, height do | |
| if lineTypeMask[y][x] == 2 then | |
| lineExists = true | |
| end | |
| if borderMask[y][x] == 1 then | |
| if lineExists then | |
| pointLineTypeMask[1][y][x] = 1 | |
| end | |
| else | |
| lineExists = false | |
| end | |
| end | |
| lineExists = false | |
| for y = height, 1, -1 do | |
| if lineTypeMask[y][x] == 2 then | |
| lineExists = true | |
| end | |
| if borderMask[y][x] == 1 then | |
| if lineExists then | |
| pointLineTypeMask[3][y][x] = 1 | |
| end | |
| else | |
| lineExists = false | |
| end | |
| end | |
| end | |
| for y = 1, height do | |
| local lineExists = false | |
| for x = 1, width do | |
| if lineTypeMask[y][x] == 1 then | |
| lineExists = true | |
| end | |
| if borderMask[y][x] == 1 then | |
| if lineExists then | |
| pointLineTypeMask[4][y][x] = 1 | |
| end | |
| else | |
| lineExists = false | |
| end | |
| end | |
| lineExists = false | |
| for x = width, 1, -1 do | |
| if lineTypeMask[y][x] == 1 then | |
| lineExists = true | |
| end | |
| if borderMask[y][x] == 1 then | |
| if lineExists then | |
| pointLineTypeMask[2][y][x] = 1 | |
| end | |
| else | |
| lineExists = false | |
| end | |
| end | |
| end | |
| local pointMask = torch.cmul(pointLineTypeMask[1] + pointLineTypeMask[3], pointLineTypeMask[2] + pointLineTypeMask[4]):gt(0) | |
| pointMask = image.erode(pointMask) | |
| pointMask = image.dilate(pointMask) | |
| pointMask = image.dilate(pointMask) | |
| local pointIndexedMask, numPoints = utils.findConnectedComponents(pointMask) | |
| local points = {} | |
| for pointIndex = 1, numPoints do | |
| local indices = pointIndexedMask:eq(pointIndex):nonzero() | |
| if ##indices > 0 then | |
| local means = torch.mean(indices:double(), 1)[1] | |
| local x = torch.round(means[2]) | |
| local y = torch.round(means[1]) | |
| local lineTypes = pointLineTypeMask[{{}, y, x}] | |
| local pointType = lineTypes:sum() | |
| local pointOrientation | |
| if pointType == 2 then | |
| if lineTypes[1] == 1 then | |
| if lineTypes[2] == 1 then | |
| pointOrientation = 2 | |
| elseif lineTypes[4] == 1 then | |
| pointOrientation = 1 | |
| end | |
| elseif lineTypes[3] == 1 then | |
| if lineTypes[2] == 1 then | |
| pointOrientation = 3 | |
| elseif lineTypes[4] == 1 then | |
| pointOrientation = 4 | |
| end | |
| end | |
| elseif pointType == 3 then | |
| for c = 1, 4 do | |
| if lineTypes[c] == 0 then | |
| pointOrientation = c | |
| end | |
| end | |
| elseif pointType == 4 then | |
| pointOrientation = 1 | |
| end | |
| if pointOrientation then | |
| table.insert(points, {{x, y}, {x, y}, {'point', pointType, pointOrientation}}) | |
| --print('(' .. x .. ', ' .. y .. ')\t' .. pointType .. ' ' .. pointOrientation) | |
| end | |
| end | |
| end | |
| --[[ | |
| image.save('test/border.png', borderMask:double()) | |
| image.save('test/segmentation.png', utils.drawSegmentation(floorplanSegmentation)) | |
| for c = 1, 4 do | |
| image.save('test/points_' .. c .. '.png', pointLineTypeMask[c]:double()) | |
| end | |
| image.save('test/points.png', pointMask:double()) | |
| ]]-- | |
| local walls = utils.pointsToLines(width, height, points, lineWidth) | |
| return walls, points | |
| end | |
| function utils.invertFloorplanHeuristic(floorplan) | |
| pl.dir.makepath('test/') | |
| local width = floorplan:size(3) | |
| local height = floorplan:size(2) | |
| local lineWidth = lineWidth or 5 | |
| local binaryThreshold = 0.7 | |
| if denotedWalls == nil then | |
| local floorplanSegmentationReversed, numSegmentsReversed, floorplanBinaryReversed = utils.segmentFloorplan(floorplan, binaryThreshold, 0, true) | |
| local backgroundSegmentReversed = floorplanSegmentationReversed[1][1] | |
| local maxCount = 0 | |
| local borderSegmentReversed | |
| for segmentReversed = 1, numSegmentsReversed do | |
| if segmentReversed ~= backgroundSegmentReversed then | |
| local indices = floorplanSegmentationReversed:eq(segmentReversed):nonzero() | |
| if ##indices > 0 then | |
| local count = (#indices)[1] | |
| if count > maxCount then | |
| borderSegmentReversed = segmentReversed | |
| maxCount = count | |
| end | |
| end | |
| end | |
| end | |
| if borderSegmentReversed then | |
| local borderMask = floorplanSegmentationReversed:eq(borderSegmentReversed) | |
| binaryThreshold = torch.max(floorplan, 1)[1][borderMask]:max() | |
| end | |
| else | |
| local denotedWallMask = torch.zeros(height, width) | |
| for _, wall in pairs(denotedWalls) do | |
| local mins = {} | |
| local maxs = {} | |
| for c = 1, 2 do | |
| mins[c] = math.min(wall[1][c], wall[2][c]) | |
| maxs[c] = math.max(wall[1][c], wall[2][c]) | |
| end | |
| denotedWallMask[{{mins[2], max[2]}, {mins[1], max[1]}}] = 1 | |
| end | |
| binaryThreshold = torch.max(floorplan, 1)[1][denotedWallMask]:mean() + 0.1 | |
| end | |
| local floorplanSegmentation, numSegments, floorplanBinary = utils.segmentFloorplan(floorplan, binaryThreshold, 0) | |
| floorplanSegmentation = floorplanSegmentation - 1 | |
| numSegments = numSegments - 1 | |
| local borderMask = floorplanSegmentation:eq(0) | |
| local maxDoorWidth = 8 | |
| local minDoorLength = 15 | |
| local backgroundSegment = floorplanSegmentation[1][1] | |
| --local doorMask = torch.zeros(height, width) | |
| for segment = 1, numSegments do | |
| if segment ~= backgroundSegment then | |
| local mask = floorplanSegmentation:eq(segment) | |
| local indices = floorplanSegmentation:eq(segment):nonzero() | |
| if ##indices > 0 then | |
| local mins = torch.min(indices, 1)[1] | |
| local maxs = torch.max(indices, 1)[1] | |
| for c = 1, 2 do | |
| if maxs[c] - mins[c] <= maxDoorWidth and maxs[3 - c] - mins[3 - c] >= minDoorLength then | |
| local borderOverlap = torch.cmul(image.dilate(mask), borderMask):nonzero() | |
| if ##borderOverlap > 0 then | |
| borderMask[mask] = 1 | |
| --doorMask[mask] = 1 | |
| break | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| local lineLengthMask = torch.zeros(4, height, width) | |
| for x = 1, width do | |
| local length = 0 | |
| for y = 1, height do | |
| if borderMask[y][x] == 1 then | |
| length = length + 1 | |
| lineLengthMask[1][y][x] = length | |
| else | |
| length = 0 | |
| end | |
| end | |
| length = 0 | |
| for y = height, 1, -1 do | |
| if borderMask[y][x] == 1 then | |
| length = length + 1 | |
| lineLengthMask[3][y][x] = length | |
| else | |
| length = 0 | |
| end | |
| end | |
| end | |
| for y = 1, height do | |
| local length = 0 | |
| for x = 1, width do | |
| if borderMask[y][x] == 1 then | |
| length = length + 1 | |
| lineLengthMask[4][y][x] = length | |
| else | |
| length = 0 | |
| end | |
| end | |
| length = 0 | |
| for x = width, 1, -1 do | |
| if borderMask[y][x] == 1 then | |
| length = length + 1 | |
| lineLengthMask[2][y][x] = length | |
| else | |
| length = 0 | |
| end | |
| end | |
| end | |
| local minLineWidth = 5 | |
| local maxLineWidth = 15 | |
| local minLineLength = 20 | |
| --local maxLineWidth = 15 | |
| --local minLineLength = 20 | |
| local horizontalLineLengthMask = lineLengthMask[2] + lineLengthMask[4] - 1 | |
| local verticalLineLengthMask = lineLengthMask[1] + lineLengthMask[3] - 1 | |
| local lineTypeMask = torch.zeros(height, width) | |
| local verticalLineMask = torch.cmul(torch.cmul(verticalLineLengthMask:ge(minLineWidth), verticalLineLengthMask:le(maxLineWidth)), horizontalLineLengthMask:ge(minLineLength)) | |
| verticalLineMask = image.erode(verticalLineMask) | |
| verticalLineMask = image.dilate(verticalLineMask) | |
| lineTypeMask[verticalLineMask] = 1 | |
| local horizontalLineMask = torch.cmul(torch.cmul(horizontalLineLengthMask:ge(minLineWidth), horizontalLineLengthMask:le(maxLineWidth)), verticalLineLengthMask:ge(minLineLength)) | |
| horizontalLineMask = image.erode(horizontalLineMask) | |
| horizontalLineMask = image.dilate(horizontalLineMask) | |
| lineTypeMask[horizontalLineMask] = 2 | |
| local pointLineTypeMask = torch.zeros(4, height, width) | |
| for x = 1, width do | |
| local lineExists = false | |
| for y = 1, height do | |
| if lineTypeMask[y][x] == 2 then | |
| lineExists = true | |
| end | |
| if borderMask[y][x] == 1 then | |
| if lineExists then | |
| pointLineTypeMask[1][y][x] = 1 | |
| end | |
| else | |
| lineExists = false | |
| end | |
| end | |
| lineExists = false | |
| for y = height, 1, -1 do | |
| if lineTypeMask[y][x] == 2 then | |
| lineExists = true | |
| end | |
| if borderMask[y][x] == 1 then | |
| if lineExists then | |
| pointLineTypeMask[3][y][x] = 1 | |
| end | |
| else | |
| lineExists = false | |
| end | |
| end | |
| end | |
| for y = 1, height do | |
| local lineExists = false | |
| for x = 1, width do | |
| if lineTypeMask[y][x] == 1 then | |
| lineExists = true | |
| end | |
| if borderMask[y][x] == 1 then | |
| if lineExists then | |
| pointLineTypeMask[4][y][x] = 1 | |
| end | |
| else | |
| lineExists = false | |
| end | |
| end | |
| lineExists = false | |
| for x = width, 1, -1 do | |
| if lineTypeMask[y][x] == 1 then | |
| lineExists = true | |
| end | |
| if borderMask[y][x] == 1 then | |
| if lineExists then | |
| pointLineTypeMask[2][y][x] = 1 | |
| end | |
| else | |
| lineExists = false | |
| end | |
| end | |
| end | |
| local pointMask = torch.cmul(pointLineTypeMask[1] + pointLineTypeMask[3], pointLineTypeMask[2] + pointLineTypeMask[4]):gt(0) | |
| pointMask = image.erode(pointMask) | |
| pointMask = image.dilate(pointMask) | |
| pointMask = image.dilate(pointMask) | |
| local pointIndexedMask, numPoints = utils.findConnectedComponents(pointMask) | |
| local points = {} | |
| for pointIndex = 1, numPoints do | |
| local indices = pointIndexedMask:eq(pointIndex):nonzero() | |
| if ##indices > 0 then | |
| local means = torch.mean(indices:double(), 1)[1] | |
| local x = torch.round(means[2]) | |
| local y = torch.round(means[1]) | |
| local lineTypes = pointLineTypeMask[{{}, y, x}] | |
| local pointType = lineTypes:sum() | |
| local pointOrientation | |
| if pointType == 2 then | |
| if lineTypes[1] == 1 then | |
| if lineTypes[2] == 1 then | |
| pointOrientation = 2 | |
| elseif lineTypes[4] == 1 then | |
| pointOrientation = 1 | |
| end | |
| elseif lineTypes[3] == 1 then | |
| if lineTypes[2] == 1 then | |
| pointOrientation = 3 | |
| elseif lineTypes[4] == 1 then | |
| pointOrientation = 4 | |
| end | |
| end | |
| elseif pointType == 3 then | |
| for c = 1, 4 do | |
| if lineTypes[c] == 0 then | |
| pointOrientation = c | |
| end | |
| end | |
| elseif pointType == 4 then | |
| pointOrientation = 1 | |
| end | |
| if pointOrientation then | |
| table.insert(points, {{x, y}, {x, y}, {'point', pointType, pointOrientation}}) | |
| --print('(' .. x .. ', ' .. y .. ')\t' .. pointType .. ' ' .. pointOrientation) | |
| end | |
| end | |
| end | |
| --[[ | |
| image.save('test/border.png', borderMask:double()) | |
| image.save('test/segmentation.png', utils.drawSegmentation(floorplanSegmentation)) | |
| for c = 1, 4 do | |
| image.save('test/points_' .. c .. '.png', pointLineTypeMask[c]:double()) | |
| end | |
| image.save('test/points.png', pointMask:double()) | |
| ]]-- | |
| local walls, lineJunctionsMap = utils.pointsToLines(width, height, points, lineWidth) | |
| -- borderMask = image.erode(borderMask) | |
| -- borderMask = image.dilate(borderMask) | |
| -- borderMask = image.erode(borderMask) | |
| -- borderMask = image.dilate(borderMask) | |
| -- local minLineWidth = 1 | |
| -- local maxLineWidth = 10 | |
| -- local minLineLength = 20 | |
| -- -- local horizontalLineLengthMask = lineLengthMask[2] + lineLengthMask[4] - 1 | |
| -- -- local verticalLineLengthMask = lineLengthMask[1] + lineLengthMask[3] - 1 | |
| -- -- local lineTypeMask = torch.zeros(height, width) | |
| -- -- lineTypeMask[torch.cmul(torch.cmul(verticalLineLengthMask:ge(5), verticalLineLengthMask:le(15)), horizontalLineLengthMask:ge(minLineLength))] = 1 | |
| -- -- lineTypeMask[torch.cmul(torch.cmul(horizontalLineLengthMask:ge(5), horizontalLineLengthMask:le(15)), verticalLineLengthMask:ge(minLineLength))] = 2 | |
| -- local points = {} | |
| -- local pointOrientations = {{{3}, {4}, {1}, {2}}, {{4, 1}, {1, 2}, {2, 3}, {3, 4}}, {{2, 3, 4}, {3, 4, 1}, {4, 1, 2}, {1, 2, 3}}, {{1, 2, 3, 4}}} | |
| -- for pointType, orientationMaps in pairs(pointOrientations) do | |
| -- local numOrientations = 4 | |
| -- if pointType == 4 then | |
| -- numOrientations = 1 | |
| -- end | |
| -- for pointOrientation, orientations in pairs(orientationMaps) do | |
| -- local mask = torch.ones(height, width):byte() | |
| -- local orientationMask = {} | |
| -- for _, orientation in pairs(orientations) do | |
| -- mask:cmul(lineLengthMask[orientation]:ge(minLineLength)) | |
| -- orientationMask[orientation] = true | |
| -- end | |
| -- for orientation = 1, numOrientations do | |
| -- if not orientationMask[orientation] then | |
| -- mask:cmul(lineLengthMask[orientation]:le(maxLineWidth)) | |
| -- mask:cmul(lineLengthMask[orientation]:gt(minLineWidth)) | |
| -- end | |
| -- end | |
| -- --mask = image.dilate(mask) | |
| -- mask = image.dilate(mask) | |
| -- mask = image.erode(mask) | |
| -- mask = image.erode(mask) | |
| -- -- mask = image.erode(mask) | |
| -- -- mask = image.erode(mask) | |
| -- local pointMask, numPoints = utils.findConnectedComponents(mask) | |
| -- image.save('test/points_' .. pointType .. '_' .. pointOrientation .. '.png', pointMask) | |
| -- for pointIndex = 1, numPoints - 1 do | |
| -- local indices = pointMask:eq(pointIndex):nonzero() | |
| -- local means = indices:double():mean(1)[1] | |
| -- local maxs = indices:max(1)[1] | |
| -- local mins = indices:min(1)[1] | |
| -- if maxs[1] - mins[1] <= maxLineWidth * 2 and maxs[2] - mins[2] <= maxLineWidth * 2 then | |
| -- table.insert(points, {{means[2], means[1]}, {means[2], means[1]}, {'point', pointType, pointOrientation}}) | |
| -- else | |
| -- if maxs[1] - mins[1] > maxLineWidth * 2 then | |
| -- table.insert(points, {{means[2], mins[1] + lineWidth}, {means[2], mins[1] + lineWidth}, {'point', pointType, pointOrientation}}) | |
| -- table.insert(points, {{means[2], maxs[1] - lineWidth}, {means[2], maxs[1] - lineWidth}, {'point', pointType, pointOrientation}}) | |
| -- end | |
| -- if maxs[2] - mins[2] > maxLineWidth * 2 then | |
| -- table.insert(points, {{mins[2] + lineWidth, means[1]}, {mins[2] + lineWidth, means[1]}, {'point', pointType, pointOrientation}}) | |
| -- table.insert(points, {{maxs[2] - lineWidth, means[1]}, {maxs[2] - lineWidth, means[1]}, {'point', pointType, pointOrientation}}) | |
| -- end | |
| -- end | |
| -- end | |
| -- end | |
| -- end | |
| -- local mask = torch.ones(height, width):byte() | |
| -- mask:cmul(lineLengthMask[1]:ge(minLineLength)) | |
| -- mask:cmul(lineLengthMask[2]:ge(minLineLength)) | |
| -- image.save('test/mask_1.png', mask:double()) | |
| -- mask:cmul(lineLengthMask[3]:le(maxLineWidth)) | |
| -- mask:cmul(lineLengthMask[4]:le(maxLineWidth)) | |
| -- image.save('test/mask_2.png', mask:double()) | |
| -- mask:cmul(lineLengthMask[3]:ge(minLineWidth)) | |
| -- mask:cmul(lineLengthMask[4]:ge(minLineWidth)) | |
| -- image.save('test/mask_3.png', mask:double()) | |
| -- image.save('test/floorplan.png', floorplan) | |
| -- local walls, lineJunctionsMap = utils.pointsToLines(width, height, points, lineWidth, true) | |
| -- local maxPooling = nn.SpatialMaxPooling(lineWidth * 2 + 1, lineWidth * 2 + 1, 1, 1, lineWidth, lineWidth) | |
| -- lineLengthMask = maxPooling:forward(lineLengthMask) | |
| -- local newWalls = {} | |
| -- local invalidPointMask = {} | |
| -- for wallIndex, wall in pairs(walls) do | |
| -- local lineDim = utils.lineDim(wall) | |
| -- local endPoints = {} | |
| -- if lineDim == 1 then | |
| -- endPoints = {2, 4} | |
| -- else | |
| -- endPoints = {3, 1} | |
| -- end | |
| -- local lineLength = wall[2][lineDim] - wall[1][lineDim] | |
| -- local invalid = false | |
| -- for i = 1, 2 do | |
| -- if lineLengthMask[endPoints[i]][wall[i][2]][wall[i][1]] < lineLength - lineWidth then | |
| -- invalid = true | |
| -- invalidPointMask[lineJunctionsMap[wallIndex][1]] = true | |
| -- invalidPointMask[lineJunctionsMap[wallIndex][2]] = true | |
| -- --print(wall) | |
| -- --print(lineLength) | |
| -- --print(lineLengthMask[endPoints[i]][wall[i][2]][wall[i][1]]) | |
| -- break | |
| -- end | |
| -- end | |
| -- if not invalid then | |
| -- table.insert(newWalls, wall) | |
| -- end | |
| -- end | |
| local wallMask = utils.drawLineMask(width, height, walls, lineWidth) | |
| local segments, numSegments = utils.findConnectedComponents(wallMask) | |
| local rooms, numRooms = utils.findConnectedComponents(1 - wallMask) | |
| local backgroundRoomIndex | |
| local imageCorners = {{1, 1}, {width, 1}, {width, height}, {1, height}} | |
| for _, imageCorner in pairs(imageCorners) do | |
| local roomIndex = rooms[imageCorner[2]][imageCorner[1]] | |
| if roomIndex > 0 then | |
| if not backgroundRoomIndex then | |
| backgroundRoomIndex = roomIndex | |
| elseif roomIndex ~= backgroundRoomIndex then | |
| rooms[rooms:eq(roomIndex)] = backgroundRoomIndex | |
| end | |
| end | |
| end | |
| if backgroundRoomIndex then | |
| local backgroundMask = rooms:eq(backgroundRoomIndex) | |
| backgroundMask = image.dilate(backgroundMask) | |
| local backgroundSegmentIndices = torch.cmul(segments, backgroundMask:int()):nonzero() | |
| if ##backgroundSegmentIndices > 0 then | |
| local backgroundSegments = {} | |
| for i = 1, backgroundSegmentIndices:size(1) do | |
| backgroundSegments[segments[backgroundSegmentIndices[i][1]][backgroundSegmentIndices[i][2]]] = true | |
| end | |
| local lineMask = torch.zeros(height, width) | |
| for segmentIndex, _ in pairs(backgroundSegments) do | |
| lineMask[segments:eq(segmentIndex)] = 1 | |
| end | |
| local newPoints = {} | |
| for pointIndex, point in pairs(points) do | |
| if lineMask[point[1][2]][point[1][1]] > 0 then | |
| table.insert(newPoints, point) | |
| end | |
| end | |
| points = newPoints | |
| end | |
| end | |
| local doorMask = torch.cmul((borderMask - floorplanSegmentation:eq(0)):double(), wallMask) | |
| doorMask = image.dilate(doorMask) | |
| doorMask = image.dilate(doorMask) | |
| doorMask = image.erode(doorMask) | |
| doorMask = image.erode(doorMask) | |
| local doorsMask, numDoors = utils.findConnectedComponents(doorMask) | |
| local wallMaskIndexed = utils.drawLineMask(width, height, walls, lineWidth, true) | |
| local doors = {} | |
| for doorIndex = 1, numDoors - 1 do | |
| local doorIndices = doorsMask:eq(doorIndex):nonzero() | |
| local wallPoints = {} | |
| for i = 1, doorIndices:size(1) do | |
| local x = doorIndices[i][2] | |
| local y = doorIndices[i][1] | |
| local wallIndex = wallMaskIndexed[y][x] | |
| if not wallPoints[wallIndex] then | |
| wallPoints[wallIndex] = {} | |
| end | |
| table.insert(wallPoints[wallIndex], {x, y}) | |
| end | |
| for wallIndex, doorPoints in pairs(wallPoints) do | |
| local minX, maxX, minY, maxY | |
| for _, point in pairs(doorPoints) do | |
| local x = point[1] | |
| local y = point[2] | |
| if not minX or x < minX then | |
| minX = x | |
| end | |
| if not maxX or x > maxX then | |
| maxX = x | |
| end | |
| if not minY or y < minY then | |
| minY = y | |
| end | |
| if not maxY or y > maxY then | |
| maxY = y | |
| end | |
| end | |
| if maxX - minX > math.max(maxY - minY, minDoorLength) then | |
| table.insert(doors, {{minX, (minY + maxY) / 2}, {maxX, (minY + maxY) / 2}, {'door', 1, 1}}) | |
| elseif maxY - minY > math.max(maxX - minX, minDoorLength) then | |
| table.insert(doors, {{(minX + maxX) / 2, minY}, {(minX + maxX) / 2, maxY}, {'door', 1, 2}}) | |
| end | |
| end | |
| end | |
| --image.save('test/walls.png', utils.drawLineMask(width, height, walls, lineWidth)) | |
| --image.save('test/walls_new.png', utils.drawLineMask(width, height, newWalls, lineWidth)) | |
| --walls = newWalls | |
| --points = utils.linesToPoints(width, height, walls, lineWidth) | |
| --walls = utils.pointsToLines(width, height, walls, lineWidth) | |
| --walls = utils.linesToPoints(width, height, walls, lineWidth) | |
| --points = utils.linesToPoints(width, height, walls, lineWidth) | |
| --print(points) | |
| --os.exit(1) | |
| for _, door in pairs(doors) do | |
| local lineDim = fp_ut.lineDim(door, lineWidth) | |
| local doorFixedValue = (door[1][3 - lineDim] + door[2][3 - lineDim]) / 2 | |
| local doorMinValue = math.min(door[1][lineDim], door[2][lineDim]) | |
| local doorMaxValue = math.max(door[1][lineDim], door[2][lineDim]) | |
| for _, wall in pairs(walls) do | |
| if fp_ut.lineDim(wall, lineWidth) == lineDim then | |
| local wallFixedValue = (wall[1][3 - lineDim] + wall[2][3 - lineDim]) / 2 | |
| local wallMinValue = math.min(wall[1][lineDim], wall[2][lineDim]) | |
| local wallMaxValue = math.max(wall[1][lineDim], wall[2][lineDim]) | |
| if math.abs(doorFixedValue - wallFixedValue) <= lineWidth and math.max(doorMinValue, wallMinValue) < math.min(doorMaxValue, wallMaxValue) then | |
| door[1][3 - lineDim] = wallFixedValue | |
| door[2][3 - lineDim] = wallFixedValue | |
| end | |
| end | |
| end | |
| end | |
| local representation = {} | |
| representation.walls = walls | |
| representation.points = points | |
| representation.doors = doors | |
| representation.icons = {} | |
| representation.labels = {} | |
| return representation | |
| end | |
| function utils.filterPoints(width, height, points, pointsConfidence, lineWidth, gap) | |
| local lineWidth = lineWidth or 5 | |
| local gap = gap or lineWidth * 2 | |
| local pointJunctionMap = {} | |
| local pointOrientations = {} | |
| for pointIndex, point in pairs(points) do | |
| pointJunctionMap[pointIndex] = {-1, -1, -1, -1} | |
| local orientations = {} | |
| local orientation = point[3][3] | |
| if point[3][2] == 1 then | |
| table.insert(orientations, (orientation + 2 - 1) % 4 + 1) | |
| elseif point[3][2] == 2 then | |
| table.insert(orientations, orientation) | |
| table.insert(orientations, (orientation + 3 - 1) % 4 + 1) | |
| elseif point[3][2] == 3 then | |
| table.insert(orientations, (orientation + 2 - 1) % 4 + 1) | |
| --[[ | |
| for i = 1, 4 do | |
| if i ~= orientation then | |
| table.insert(orientations, i) | |
| end | |
| end | |
| ]]-- | |
| else | |
| --[[ | |
| for i = 1, 4 do | |
| table.insert(orientations, i) | |
| end | |
| ]]-- | |
| end | |
| pointOrientations[pointIndex] = orientations | |
| for _, orientation in pairs(orientations) do | |
| pointJunctionMap[pointIndex][orientation] = 0 | |
| end | |
| end | |
| for pointIndex, point in pairs(points) do | |
| local orientations = pointOrientations[pointIndex] | |
| local x = point[1][1] | |
| local y = point[1][2] | |
| for _, orientation in pairs(orientations) do | |
| if pointJunctionMap[pointIndex][orientation] == 0 then | |
| local lineDim | |
| local fixedValue | |
| local junction | |
| local startPoint | |
| if orientation == 1 or orientation == 3 then | |
| lineDim = 2 | |
| fixedValue = x | |
| startPoint = y | |
| if orientation == 1 then | |
| junction = 1 | |
| else | |
| junction = height | |
| end | |
| else | |
| lineDim = 1 | |
| fixedValue = y | |
| startPoint = x | |
| if orientation == 4 then | |
| junction = 1 | |
| else | |
| junction = width | |
| end | |
| end | |
| local orientationOpposite = (orientation + 2 - 1) % 4 + 1 | |
| local selectedOtherPointIndex | |
| for otherPointIndex, otherPoint in pairs(points) do | |
| if otherPointIndex ~= pointIndex and pointJunctionMap[otherPointIndex][orientationOpposite] == 0 then | |
| local otherXY = otherPoint[1] | |
| if otherXY[lineDim] > math.min(junction, startPoint) and otherXY[lineDim] < math.max(junction, startPoint) and otherXY[3 - lineDim] >= fixedValue - lineWidth and otherXY[3 - lineDim] <= fixedValue + lineWidth then | |
| junction = otherXY[lineDim] | |
| selectedOtherPointIndex = otherPointIndex | |
| end | |
| end | |
| end | |
| local point_1 = {} | |
| point_1[lineDim] = math.min(startPoint, junction) | |
| point_1[3 - lineDim] = fixedValue | |
| local point_2 = {} | |
| point_2[lineDim] = math.max(startPoint, junction) | |
| point_2[3 - lineDim] = fixedValue | |
| if point_1[1] ~= 1 and point_1[1] ~= width and point_1[2] ~= 1 and point_1[2] ~= height and point_2[1] ~= 1 and point_2[1] ~= width and point_2[2] ~= 1 and point_2[2] ~= height then | |
| --table.insert(lines, {point_1, point_2, {"wall", 1, 1}}) | |
| if selectedOtherPointIndex ~= nil then | |
| pointJunctionMap[selectedOtherPointIndex][orientationOpposite] = pointIndex | |
| end | |
| pointJunctionMap[pointIndex][orientation] = selectedOtherPointIndex | |
| --print(pointIndex .. ' ' .. orientation .. ' ' .. selectedOtherPointIndex) | |
| end | |
| end | |
| end | |
| end | |
| local removedPoints = {} | |
| for pointIndex_1, point_1 in pairs(points) do | |
| if not removedPoints[pointIndex_1] then | |
| for pointIndex_2, point_2 in pairs(points) do | |
| if pointIndex_2 > pointIndex_1 and not removedPoints[pointIndex_2] then | |
| local distance = utils.calcDistance(point_1[1], point_2[1]) | |
| if distance < gap then | |
| local confidence_1 = pointsConfidence[pointIndex_1] | |
| local confidence_2 = pointsConfidence[pointIndex_2] | |
| local isNeighbor = false | |
| for orientation, neighborPointIndex in pairs(pointJunctionMap[pointIndex_1]) do | |
| if neighborPointIndex > 0 then | |
| if neighborPointIndex == pointIndex_2 then | |
| isNeighbor = true | |
| end | |
| confidence_1 = confidence_1 + pointsConfidence[neighborPointIndex] | |
| end | |
| end | |
| for orientation, neighborPointIndex in pairs(pointJunctionMap[pointIndex_2]) do | |
| if neighborPointIndex > 0 then | |
| confidence_2 = confidence_2 + pointsConfidence[neighborPointIndex] | |
| end | |
| end | |
| if isNeighbor == false then | |
| if confidence_1 < confidence_2 then | |
| removedPoints[pointIndex_1] = true | |
| else | |
| removedPoints[pointIndex_2] = true | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| local newPoints = {} | |
| for pointIndex, point in pairs(points) do | |
| if not removedPoints[pointIndex] then | |
| table.insert(newPoints, point) | |
| end | |
| end | |
| return newPoints | |
| end | |
| function utils.invertFloorplan(model, floorplan, withoutQP, relaxedQP, useStack) | |
| local lineWidth = lineWidth or 3 | |
| local useStack = useStack or false | |
| local oriWidth = floorplan:size(3) | |
| local oriHeight = floorplan:size(2) | |
| local width = 256 | |
| local height = 256 | |
| local floorplanOri = floorplan:clone() | |
| image.save('test/floorplan.png', floorplanOri) | |
| local floorplan = image.scale(floorplan, width, height) | |
| local junctionHeatmaps, doorHeatmaps, iconHeatmaps, segmentations = utils.estimateHeatmaps(model, floorplanOri:clone(), 'single', useStack) | |
| --local junctionHeatmaps, doorHeatmaps, iconHeatmaps, _ = utils.estimateHeatmaps(floorplanOri:clone(), 'single', true) | |
| pl.dir.makepath('test/heatmaps/') | |
| for i = 1, 13 do | |
| image.save('test/heatmaps/junction_heatmap_' .. i .. '.png', junctionHeatmaps[i]) | |
| end | |
| for i = 1, 4 do | |
| image.save('test/heatmaps/door_heatmap_' .. i .. '.png', doorHeatmaps[i]) | |
| image.save('test/heatmaps/icon_heatmap_' .. i .. '.png', iconHeatmaps[i]) | |
| end | |
| pl.dir.makepath('test/segmentation/') | |
| for segmentIndex = 1, 30 do | |
| image.save('test/segmentation/segment_' .. segmentIndex .. '.png', segmentations[segmentIndex]) | |
| end | |
| local representation = {} | |
| py.execute('import os') | |
| if withoutQP then | |
| py.execute('os.system("python PostProcessing/QP.py 1")') | |
| representation.points = utils.loadItems('test/points_out.txt') | |
| representation.doors = utils.loadItems('test/doors_out.txt') | |
| representation.icons = utils.loadItems('test/icons_out.txt') | |
| representation.walls = utils.pointsToLines(oriWidth, oriHeight, representation.points, lineWidth, true) | |
| local wallMask = utils.drawLineMask(oriWidth, oriHeight, representation.walls, lineWidth) | |
| --image.save('test/floorplan.png', floorplan) | |
| --image.save('test/wall_mask.png', wallMask) | |
| local rooms, numRooms = utils.findConnectedComponents(1 - wallMask) | |
| local backgroundRoomIndex | |
| local imageCorners = {{1, 1}, {oriWidth, 1}, {oriWidth, oriHeight}, {1, oriHeight}} | |
| for _, imageCorner in pairs(imageCorners) do | |
| local roomIndex = rooms[imageCorner[2]][imageCorner[1]] | |
| if roomIndex > 0 then | |
| if not backgroundRoomIndex then | |
| backgroundRoomIndex = roomIndex | |
| elseif roomIndex ~= backgroundRoomIndex then | |
| rooms[rooms:eq(roomIndex)] = backgroundRoomIndex | |
| end | |
| end | |
| end | |
| if not backgroundRoomIndex then | |
| backgroundRoomIndex = numRooms | |
| end | |
| representation.labels = {} | |
| for roomIndex = 1, numRooms - 1 do | |
| if roomIndex ~= backgroundRoomIndex then | |
| local roomMask = rooms:eq(roomIndex) | |
| if ##roomMask:nonzero() > 0 then | |
| local means = roomMask:nonzero():double():mean(1)[1] | |
| local y = means[1] | |
| local x = means[2] | |
| local maxSum | |
| local maxSumSegmentIndex | |
| for segmentIndex = 1, 10 do | |
| local sum = segmentations[segmentIndex][roomMask]:sum() | |
| if not maxSum or sum > maxSum then | |
| maxSum = sum | |
| maxSumSegmentIndex = segmentIndex | |
| end | |
| end | |
| if maxSumSegmentIndex then | |
| table.insert(representation.labels, {{x - 20, y - 10}, {x + 20, y + 10}, utils.getItemInfo('labels', maxSumSegmentIndex)}) | |
| end | |
| end | |
| end | |
| end | |
| return representation | |
| end | |
| py.execute('os.system("python PostProcessing/QP.py")') | |
| --os.exit(1) | |
| local points = utils.loadItems('test/points_out.txt') | |
| local pointLabelsFile = csvigo.load({path='test/point_labels.txt', mode="large", header=false, separator='\t', verbose=false}) | |
| local pointLabels = {} | |
| for _, labels in pairs(pointLabelsFile) do | |
| table.insert(pointLabels, {tonumber(labels[1]), tonumber(labels[2]), tonumber(labels[3]), tonumber(labels[4])}) | |
| end | |
| representation.points = points | |
| representation.walls, wallJunctionsMap = utils.pointsToLines(oriWidth, oriHeight, points, lineWidth, true) | |
| for wallIndex, junctions in pairs(wallJunctionsMap) do | |
| local lineDim = utils.lineDim(representation.walls[wallIndex]) | |
| local labels_1 = pointLabels[junctions[1]] | |
| local labels_2 = pointLabels[junctions[2]] | |
| local label_1, label_2 | |
| if lineDim == 1 then | |
| label_1 = math.min(labels_1[1], labels_2[4]) | |
| label_2 = math.min(labels_1[2], labels_2[3]) | |
| elseif lineDim == 2 then | |
| label_1 = math.min(labels_1[3], labels_2[4]) | |
| label_2 = math.min(labels_1[2], labels_2[1]) | |
| end | |
| if label_1 == 0 then | |
| label_1 = 11 | |
| end | |
| if label_2 == 0 then | |
| label_2 = 11 | |
| end | |
| table.insert(representation.walls[wallIndex], {label_1, label_2}) | |
| end | |
| local wallMask = utils.drawLineMask(oriWidth, oriHeight, representation.walls, lineWidth) | |
| local rooms, numRooms = utils.findConnectedComponents(1 - wallMask) | |
| --image.save('test/walls_backup.png', wallMask) | |
| --image.save('test/rooms_backup.png', utils.drawSegmentation(rooms)) | |
| local deltas = {{1, -1}, {1, 1}, {-1, 1}, {-1, -1}} | |
| local roomLabelsMap = {} | |
| for pointIndex, point in pairs(points) do | |
| local roomLabels = {} | |
| for orientation, delta in pairs(deltas) do | |
| local label = pointLabels[pointIndex][orientation] | |
| if label >= 1 then | |
| local x = torch.round(point[1][1]) | |
| local y = torch.round(point[1][2]) | |
| --x = math.max(math.min(x, width), 1) | |
| --y = math.max(math.min(y, height), 1) | |
| for i = 1, 10 do | |
| if x < 1 or x > oriWidth or y < 1 or y > oriHeight then | |
| break | |
| end | |
| local roomIndex = rooms[y][x] | |
| if roomIndex > 0 then | |
| if not roomLabels[roomIndex] then | |
| roomLabels[roomIndex] = {} | |
| end | |
| roomLabels[roomIndex][label] = true | |
| break | |
| end | |
| x = x + delta[1] | |
| y = y + delta[2] | |
| end | |
| end | |
| end | |
| for roomIndex, labels in pairs(roomLabels) do | |
| if not roomLabelsMap[roomIndex] then | |
| roomLabelsMap[roomIndex] = {} | |
| end | |
| for label, _ in pairs(labels) do | |
| if not roomLabelsMap[roomIndex][label] then | |
| roomLabelsMap[roomIndex][label] = {} | |
| end | |
| local x = point[1][1] | |
| local y = point[1][2] | |
| table.insert(roomLabelsMap[roomIndex][label], {x, y}) | |
| end | |
| end | |
| end | |
| representation.labels = {} | |
| local labelWidth = 50 | |
| local labelHeight = 20 | |
| for roomIndex, labels in pairs(roomLabelsMap) do | |
| local numLabels = 0 | |
| for label, locations in pairs(labels) do | |
| numLabels = numLabels + 1 | |
| end | |
| if numLabels == 1 then | |
| local means = rooms:eq(roomIndex):nonzero():double():mean(1)[1] | |
| local y = means[1] | |
| local x = means[2] | |
| for label, locations in pairs(labels) do | |
| table.insert(representation.labels, {{x - labelWidth / 2, y - labelHeight / 2}, {x + labelWidth / 2, y + labelHeight / 2}, utils.getItemInfo('labels', label)}) | |
| end | |
| else | |
| for label, locations in pairs(labels) do | |
| local locationGroupMap = {} | |
| local groupIndex = 1 | |
| for locationIndex, location in pairs(locations) do | |
| if not locationGroupMap[locationIndex] then | |
| locationGroupMap[locationIndex] = groupIndex | |
| local groupLocations = {location} | |
| while true do | |
| local hasChange = false | |
| for _, groupLocation in pairs(groupLocations) do | |
| for neighborLocationIndex, neighborLocation in pairs(locations) do | |
| if not locationGroupMap[neighborLocationIndex] and (math.abs(neighborLocation[1] - groupLocation[1]) < lineWidth or math.abs(neighborLocation[2] - groupLocation[2]) < lineWidth) then | |
| local lineDim = utils.lineDim({groupLocation, neighborLocation}, lineWidth) | |
| local fixedValue = torch.round((groupLocation[3 - lineDim] + neighborLocation[3 - lineDim]) / 2) | |
| local minValue = torch.round(math.min(groupLocation[lineDim], neighborLocation[lineDim])) | |
| local maxValue = torch.round(math.max(groupLocation[lineDim], neighborLocation[lineDim])) | |
| local onWall = true | |
| for value = minValue, maxValue do | |
| if (lineDim == 1 and wallMask[fixedValue][value] == 0) or (lineDim == 2 and wallMask[value][fixedValue] == 0) then | |
| onWall = false | |
| end | |
| end | |
| if onWall then | |
| table.insert(groupLocations, neighborLocation) | |
| locationGroupMap[neighborLocationIndex] = groupIndex | |
| hasChange = true | |
| end | |
| end | |
| end | |
| end | |
| if not hasChange then | |
| break | |
| end | |
| end | |
| groupIndex = groupIndex + 1 | |
| end | |
| end | |
| for groupIndex = 1, groupIndex - 1 do | |
| local groupLocations = {} | |
| for locationIndex, index in pairs(locationGroupMap) do | |
| if index == groupIndex then | |
| table.insert(groupLocations, locations[locationIndex]) | |
| end | |
| end | |
| if #groupLocations > 2 then | |
| local x = 0 | |
| local y = 0 | |
| for _, location in pairs(groupLocations) do | |
| x = x + location[1] | |
| y = y + location[2] | |
| end | |
| x = x / #groupLocations | |
| y = y / #groupLocations | |
| table.insert(representation.labels, {{x - labelWidth / 2, y - labelHeight / 2}, {x + labelWidth / 2, y + labelHeight / 2}, utils.getItemInfo('labels', label)}) | |
| end | |
| end | |
| end | |
| end | |
| end | |
| representation.doors = utils.loadItems('test/doors_out.txt') | |
| representation.icons = utils.loadItems('test/icons_out.txt') | |
| return representation | |
| end | |
| function utils.postprocessRepresentation(width, height, representation) | |
| local lineWidth = lineWidth or 5 | |
| local lineMask | |
| if representation.walls then | |
| lineMask = utils.drawLineMask(width, height, representation.walls, lineWidth) | |
| else | |
| local lines = utils.pointsToLines(width, height, representation.points, lineWidth, true) | |
| local lineMask = utils.drawLineMask(width, height, lines, lineWidth) | |
| end | |
| for __, item in pairs(representation.icons) do | |
| item[3] = utils.getItemInfo('icons', utils.getNumber('icons', item[3])) | |
| local rectangle = {{math.min(item[1][1], item[2][1]), math.min(item[1][2], item[2][2])}, {math.max(item[1][1], item[2][1]), math.max(item[1][2], item[2][2])}} | |
| for pointIndex = 1, 2 do | |
| rectangle[pointIndex][1] = torch.round(math.max(math.min(rectangle[pointIndex][1], width), 1)) | |
| rectangle[pointIndex][2] = torch.round(math.max(math.min(rectangle[pointIndex][2], height), 1)) | |
| end | |
| local minDistance = math.max(width, height) | |
| local minDistanceOrientation = 1 | |
| if (math.abs(rectangle[2][1] - rectangle[1][1]) > math.abs(rectangle[2][2] - rectangle[1][2]) and item[3][1] ~= 'toilet' and item[3][1] ~= 'stairs') or (math.abs(rectangle[2][1] - rectangle[1][1]) < math.abs(rectangle[2][2] - rectangle[1][2]) and (item[3][1] == 'toilet' or item[3][1] == 'stairs')) or item[3][1] == 'washing_basin' then | |
| local min = math.min(rectangle[1][2], rectangle[2][2]) | |
| local max = math.max(rectangle[1][2], rectangle[2][2]) | |
| local center = (rectangle[1][1] + rectangle[2][1]) / 2 | |
| local deltaMin = 0 | |
| for delta = 1, min - 1 do | |
| deltaMin = delta | |
| if lineMask[min - delta][center] == 1 then | |
| break | |
| end | |
| end | |
| local deltaMax = 0 | |
| for delta = 1, height - max do | |
| deltaMax = delta | |
| if lineMask[max + delta][center] == 1 then | |
| break | |
| end | |
| end | |
| if deltaMax < minDistance then | |
| minDistance = deltaMax | |
| minDistanceOrientation = 1 | |
| end | |
| if deltaMin < minDistance then | |
| minDistance = deltaMin | |
| minDistanceOrientation = 3 | |
| end | |
| end | |
| if (math.abs(rectangle[2][1] - rectangle[1][1]) < math.abs(rectangle[2][2] - rectangle[1][2]) and item[3][1] ~= 'toilet') or (math.abs(rectangle[2][1] - rectangle[1][1]) > math.abs(rectangle[2][2] - rectangle[1][2]) and item[3][1] == 'toilet') or item[3][1] == 'washing_basin' then | |
| local min = math.min(rectangle[1][1], rectangle[2][1]) | |
| local max = math.max(rectangle[1][1], rectangle[2][1]) | |
| local center = (rectangle[1][2] + rectangle[2][2]) / 2 | |
| local deltaMin = 0 | |
| for delta = 1, min - 1 do | |
| deltaMin = delta | |
| if lineMask[center][min - delta] == 1 then | |
| break | |
| end | |
| end | |
| local deltaMax = 0 | |
| for delta = 1, width - max do | |
| deltaMax = delta | |
| if lineMask[center][max + delta] == 1 then | |
| break | |
| end | |
| end | |
| if deltaMax < minDistance then | |
| minDistance = deltaMax | |
| minDistanceOrientation = 4 | |
| end | |
| if deltaMin < minDistance then | |
| minDistance = deltaMin | |
| minDistanceOrientation = 2 | |
| end | |
| end | |
| item[3][3] = minDistanceOrientation | |
| end | |
| end | |
| function utils.findWallsPrediction(floorplan) | |
| local lineWidth = lineWidth or 5 | |
| local oriWidth = floorplan:size(3) | |
| local oriHeight = floorplan:size(2) | |
| local width = 256 | |
| local height = 256 | |
| local floorplanOri = floorplan:clone() | |
| local floorplan = image.scale(floorplan, width, height) | |
| local useStack = false | |
| if utils.modelWall == nil then | |
| if useStack then | |
| utils.modelJunction = torch.load('/home/chenliu/Projects/Floorplan/models/junction-wall-stack-16/model_best.t7') | |
| else | |
| utils.modelJunction = torch.load('/home/chenliu/Projects/Floorplan/models/junction-heatmap/model_best.t7') | |
| end | |
| end | |
| local sampleDim = 256 | |
| local gridDim = 8 | |
| local numAnchorBoxes = 13 | |
| local nClasses = 13 | |
| --local floorplanScaled = image.scale(floorplan, sampleDim, sampleDim) | |
| --model:evaluate() | |
| --local datasets = require 'datasets/init' | |
| --local opts = require '../../InverseCAD/opts' | |
| --local opt = opts.parse(arg) | |
| package.path = '../InverseCAD/datasets/?.lua;' .. package.path | |
| package.path = '../InverseCAD/?.lua;' .. package.path | |
| --local dataset = require('/home/chenliu/Projects/Floorplan/floorplan/InverseCAD/datasets/floorplan-representation')(nil, nil, 'val') | |
| --local dataset = require('floorplan-representation')(nil, nil, 'val') | |
| --local dataset = require 'floorplan-representation'({}, nil, 'val') | |
| local dataset = require('floorplan-representation') | |
| local floorplanNormalized = dataset:preprocessResize(sampleDim, sampleDim)(floorplan) | |
| local input = floorplanNormalized:repeatTensor(1, 1, 1, 1):cuda() | |
| --print(#input) | |
| local output = utils.modelJunction:forward(input) | |
| local wallsResult | |
| local labelsResult | |
| image.save('test/floorplan.png', floorplanOri) | |
| for stackIndex = 1, 2 do | |
| local prediction | |
| if useStack then | |
| prediction = output[stackIndex][1] | |
| else | |
| prediction = output[1] | |
| end | |
| --print(torch.max(prediction)) | |
| --print(torch.min(prediction)) | |
| local points = {} | |
| local pointsConfidence = {} | |
| local pointLabels = {} | |
| for i = 1, 13 do | |
| --print(i) | |
| local pointMask = prediction[i]:double():gt(0.5) | |
| pointMask = image.erode(pointMask) | |
| --pointMask = image.erode(pointMask) | |
| if ##pointMask:nonzero() > 0 then | |
| local components, numComponents = utils.findConnectedComponents(pointMask) | |
| --image.save('test/mask_' .. i .. '.png', pointMask:double()) | |
| --image.save('test/mask_' .. i .. '.png', utils.drawSegmentation(components)) | |
| numComponents = numComponents - 1 | |
| --print(torch.max(components) .. ' ' .. numComponents) | |
| local itemInfo = utils.getItemInfo('points', i) | |
| for componentIndex = 1, numComponents do | |
| local indices = components:eq(componentIndex):nonzero() | |
| local mean = torch.mean(indices:double(), 1)[1] | |
| --print(mean) | |
| local x = mean[2] / width * oriWidth | |
| local y = mean[1] / height * oriHeight | |
| table.insert(points, {{x, y}, {x, y}, itemInfo}) | |
| x = math.max(math.min(torch.round(mean[2]), width), 1) | |
| y = math.max(math.min(torch.round(mean[1]), height), 1) | |
| table.insert(pointsConfidence, prediction[i][y][x]) | |
| end | |
| end | |
| end | |
| --points = utils.filterPoints(width, height, points, pointsConfidence, lineWidth) | |
| --os.exit(1) | |
| local walls = utils.pointsToLines(oriWidth, oriHeight, points, lineWidth) | |
| --wallsResult = walls | |
| local wallMask = utils.drawLineMask(oriWidth, oriHeight, walls, lineWidth) | |
| if true then | |
| wallMask:mul(0.3) | |
| wallMask = wallMask:repeatTensor(3, 1, 1) | |
| local junctionMasks = utils.drawJunctionMasks(oriWidth, oriHeight, points, #points) | |
| for i = 1, junctionMasks:size(1) do | |
| for c = 1, 3 do | |
| wallMask[c][junctionMasks[i]:gt(0)] = torch.uniform() | |
| end | |
| end | |
| end | |
| if useStack then | |
| image.save('test/junction_' .. stackIndex .. '.png', image.scale(wallMask, oriWidth, oriHeight)) | |
| else | |
| image.save('test/junction.png', image.scale(wallMask, oriWidth, oriHeight)) | |
| end | |
| if true then | |
| local wallMask = torch.ones(oriHeight, oriWidth) | |
| if useStack then | |
| wallMask = image.scale(prediction[14]:double(), oriWidth, oriHeight, 'simple') | |
| else | |
| if utils.modelWall == nil then | |
| utils.modelWall = torch.load('/home/chenliu/Projects/Floorplan/models/wall/model_best.t7') | |
| end | |
| wallMask = image.scale(utils.modelWall:forward(input)[1]:double(), oriWidth, oriHeight, 'simple') | |
| end | |
| image.save('test/wall.png', wallMask) | |
| utils.saveRepresentation('test/points.txt', {points = points}) | |
| local representation = {} | |
| representation.points = points | |
| representation.doors = {} | |
| representation.icons = {} | |
| representation.labels = {} | |
| if true then | |
| pl.dir.makepath('test/segmentation/') | |
| if utils.modelSegmentation == nil then | |
| utils.modelSegmentation = torch.load('/home/chenliu/Projects/Floorplan/models/segmentation/model_best.t7') | |
| end | |
| local outputSegmentation = utils.modelSegmentation:forward(input) | |
| local prob, predictionSegmentation = torch.max(outputSegmentation, 2) | |
| predictionSegmentation = predictionSegmentation[{{}, 1}]:view(sampleDim, sampleDim) | |
| for segmentIndex = 1, 10 do | |
| image.save('test/segmentation/segment_' .. segmentIndex .. '.png', image.scale(predictionSegmentation:eq(segmentIndex):double(), oriWidth, oriHeight, 'simple')) | |
| end | |
| local backgroundSegmentIndex = 26 | |
| image.save('test/segmentation/segment_0.png', image.scale(predictionSegmentation:eq(backgroundSegmentIndex):double(), oriWidth, oriHeight, 'simple')) | |
| end | |
| if false then | |
| if utils.modelDoor == nil then | |
| utils.modelDoor = torch.load('/home/chenliu/Projects/Floorplan/models/anchor-junction-door/model_best.t7') | |
| end | |
| local outputDoor = utils.modelDoor:forward(input) | |
| local coordinates = outputDoor[1]:double():transpose(2, 3):transpose(3, 4):contiguous() | |
| coordinates = coordinates:view(coordinates:size(1), coordinates:size(2), coordinates:size(3), -1, 4) | |
| local objectness = outputDoor[2]:double():transpose(2, 3):transpose(3, 4):contiguous() | |
| objectness = objectness:view(objectness:size(1), objectness:size(2), objectness:size(3), objectness:size(4), 1) | |
| local prob, pred = torch.max(outputDoor[3]:double(), 2) | |
| pred = pred[{{}, 1}]:view(-1, gridDim, gridDim, numAnchorBoxes, 1):double() | |
| pred[pred:eq(nClasses + 1)] = 0 | |
| local representationTensors = torch.cat(torch.cat(coordinates, pred:gt(0):double(), 5), pred, 5) | |
| local representationTensor = representationTensors[1] | |
| representationTensor = torch.cat(torch.cat(torch.zeros(representationTensor:size(1), representationTensor:size(2), nClasses, representationTensor:size(4)), representationTensor:double(), 3), torch.zeros(representationTensor:size(1), representationTensor:size(2), nClasses * 2, representationTensor:size(4)), 3) | |
| local representationPrediction = utils.convertTensorAnchorToRepresentation(sampleDim, sampleDim, representationTensor, 0.1) | |
| representationPrediction= utils.scaleRepresentation(representationPrediction, sampleDim, sampleDim, oriWidth, oriHeight) | |
| representation.doors = representationPrediction.doors | |
| end | |
| if false then | |
| if utils.modelIcon == nil then | |
| utils.modelIcon = torch.load('/home/chenliu/Projects/Floorplan/models/anchor-junction-icon-64/model_best.t7') | |
| end | |
| local outputIcon = utils.modelIcon:forward(input) | |
| local coordinates = outputIcon[1]:double():transpose(2, 3):transpose(3, 4):contiguous() | |
| coordinates = coordinates:view(coordinates:size(1), coordinates:size(2), coordinates:size(3), -1, 4) | |
| local objectness = outputIcon[2]:double():transpose(2, 3):transpose(3, 4):contiguous() | |
| objectness = objectness:view(objectness:size(1), objectness:size(2), objectness:size(3), objectness:size(4), 1) | |
| local prob, pred = torch.max(outputIcon[3]:double(), 2) | |
| pred = pred[{{}, 1}]:view(-1, gridDim, gridDim, numAnchorBoxes, 1):double() | |
| pred[pred:eq(nClasses + 1)] = 0 | |
| local representationTensors = torch.cat(torch.cat(coordinates, pred:gt(0):double(), 5), pred, 5) | |
| local representationTensor = representationTensors[1] | |
| representationTensor = torch.cat(torch.cat(torch.zeros(representationTensor:size(1), representationTensor:size(2), nClasses * 2, representationTensor:size(4)), representationTensor:double(), 3), torch.zeros(representationTensor:size(1), representationTensor:size(2), nClasses, representationTensor:size(4)), 3) | |
| local representationPrediction = utils.convertTensorAnchorToRepresentation(sampleDim, sampleDim, representationTensor, 0.1) | |
| representationPrediction= utils.scaleRepresentation(representationPrediction, sampleDim, sampleDim, oriWidth, oriHeight) | |
| representation.icons = representationPrediction.icons | |
| end | |
| utils.saveRepresentation('test/representation.txt', representation) | |
| --if false then | |
| --utils.filterJunctions(oriWidth, oriHeight, points, wallMask) | |
| --end | |
| py.execute('import os') | |
| py.execute('os.system("python ../InverseCAD/PostProcessing/QP_backup.py")') | |
| --os.exit(1) | |
| points = utils.loadItems('test/points_out.txt') | |
| if true then | |
| local pointLabelsFile = csvigo.load({path='test/point_labels.txt', mode="large", header=false, separator='\t', verbose=false}) | |
| pointLabels = {} | |
| for _, labels in pairs(pointLabelsFile) do | |
| table.insert(pointLabels, {tonumber(labels[1]), tonumber(labels[2]), tonumber(labels[3]), tonumber(labels[4])}) | |
| end | |
| end | |
| --os.exit(1) | |
| end | |
| --os.exit(1) | |
| local walls = utils.pointsToLines(oriWidth, oriHeight, points, lineWidth) | |
| wallsResult = walls | |
| local wallMask = utils.drawLineMask(oriWidth, oriHeight, walls, lineWidth) | |
| local rooms, numRooms = utils.findConnectedComponents(1 - wallMask) | |
| local deltas = {{1, -1}, {1, 1}, {-1, 1}, {-1, -1}} | |
| local roomLabelsMap = {} | |
| for pointIndex, point in pairs(points) do | |
| for orientation, delta in pairs(deltas) do | |
| local label = pointLabels[pointIndex][orientation] | |
| if label >= 1 then | |
| local x = torch.round(point[1][1]) | |
| local y = torch.round(point[1][2]) | |
| --x = math.max(math.min(x, width), 1) | |
| --y = math.max(math.min(y, height), 1) | |
| for i = 1, 10 do | |
| if x < 1 or x > oriWidth or y < 1 or y > oriHeight then | |
| break | |
| end | |
| local roomIndex = rooms[y][x] | |
| if roomIndex > 0 then | |
| if not roomLabelsMap[roomIndex] then | |
| roomLabelsMap[roomIndex] = {} | |
| end | |
| roomLabelsMap[roomIndex][label] = true | |
| break | |
| end | |
| x = x + delta[1] | |
| y = y + delta[2] | |
| end | |
| end | |
| end | |
| end | |
| labelsResult = {} | |
| for roomIndex, labels in pairs(roomLabelsMap) do | |
| local means = rooms:eq(roomIndex):nonzero():double():mean(1)[1] | |
| local y = means[1] | |
| local x = means[2] | |
| for label, _ in pairs(labels) do | |
| --print(roomIndex .. ' ' .. x .. ' ' .. y .. ' ' .. label) | |
| table.insert(labelsResult, {{x - 20, y - 10}, {x + 20, y + 10}, utils.getItemInfo('labels', label)}) | |
| end | |
| end | |
| --[[ | |
| image.save('test/floorplan.png', floorplanOri) | |
| floorplanOri[1][wallMask:gt(0)] = 1 | |
| floorplanOri[2][wallMask:gt(0)] = 0 | |
| floorplanOri[3][wallMask:gt(0)] = 0 | |
| image.save('test/walls.png', floorplanOri) | |
| ]]-- | |
| break | |
| -- if true then | |
| -- wallMask:mul(0.3) | |
| -- wallMask = wallMask:repeatTensor(3, 1, 1) | |
| -- local junctionMasks = utils.drawJunctionMasks(oriWidth, oriHeight, points, #points) | |
| -- for i = 1, junctionMasks:size(1) do | |
| -- for c = 1, 3 do | |
| -- wallMask[c][junctionMasks[i]:gt(0)] = torch.uniform() | |
| -- end | |
| -- end | |
| -- end | |
| -- if useStack then | |
| -- image.save('test/junction_' .. stackIndex .. '.png', image.scale(wallMask, oriWidth, oriHeight)) | |
| -- else | |
| -- image.save('test/junction.png', image.scale(wallMask, oriWidth, oriHeight)) | |
| -- break | |
| -- end | |
| end | |
| return wallsResult, labelsResult | |
| end | |
| function utils.saveJunctions(filename, points) | |
| local junctionFile = io.open(filename, 'w') | |
| for _, item in pairs(points) do | |
| for __, field in pairs(item) do | |
| if __ <= 3 then | |
| for ___, value in pairs(field) do | |
| junctionFile:write(value .. '\t') | |
| end | |
| end | |
| end | |
| junctionFile:write('\n') | |
| end | |
| junctionFile:close() | |
| end | |
| function utils.findDoors(floorplan, denotedDoors, walls, floorplanSegmentation) | |
| end | |
| function utils.extractRepresentation(floorplan) | |
| local representation = {} | |
| local denotedRepresentation = denotedRepresentation or {} | |
| representation.walls, representation.points = utils.findWalls(floorplan, denotedRepresentation.walls, 5) | |
| --representation.doors = utils.findDoors(floorplan, denotedRepresentation.doors, representation.walls, floorplanSegmentation) | |
| if true then | |
| representation.doors = {} | |
| representation.icons = {} | |
| representation.labels = {} | |
| return representation | |
| end | |
| representation.doors = {} | |
| local candidateRegions, maskImage, labels = utils.extractCandidateRegions(floorplan, floorplanSegmentation, representation.walls) | |
| for _, mode in pairs({'doors'}) do | |
| local numRegions = torch.max(candidateRegions[mode]) | |
| print(numRegions) | |
| for regionIndex = 1, numRegions do | |
| local mask = candidateRegions[mode]:eq(regionIndex) | |
| local indices = mask:nonzero() | |
| if ##indices > 0 then | |
| local mins = torch.min(indices, 1)[1] | |
| local maxs = torch.max(indices, 1)[1] | |
| local point_1, point_2 | |
| if mode == 'doors' then | |
| if maxs[1] - mins[1] > maxs[2] - mins[2] then | |
| local meanX = (maxs[2] + mins[2]) / 2 | |
| point_1 = {meanX, mins[1]} | |
| point_2 = {meanX, maxs[1]} | |
| else | |
| local meanY = (maxs[1] + mins[1]) / 2 | |
| point_1 = {mins[2], meanY} | |
| point_2 = {maxs[2], meanY} | |
| end | |
| else | |
| point_1 = {mins[2] - 1, mins[1] - 1} | |
| point_2 = {maxs[2] + 1, maxs[1] + 1} | |
| end | |
| if mode == 'doors' then | |
| table.insert(representation[mode], {point_1, point_2, {'door', 1, 1}}) | |
| else | |
| table.insert(representation[mode], {point_1, point_2, {'bathtub', 1, 1}}) | |
| end | |
| end | |
| end | |
| end | |
| --representation.icons = utils.detectItems(floorplan, 'icons') | |
| representation.icons = {} | |
| representation.labels = labels | |
| return representation | |
| end | |
| function utils.extractCandidateRegions(floorplan, floorplanSegmentation, walls, lineWidth) | |
| local width = floorplan:size(3) | |
| local height = floorplan:size(2) | |
| local candidateRegions = {} | |
| candidateRegions.doors = {} | |
| candidateRegions.icons = {} | |
| --candidateRegions.labels = {} | |
| local lineWidth = lineWidth or 5 | |
| local lineMask = utils.drawLineMask(width, height, walls, lineWidth) | |
| if floorplanSegmentation == nil then | |
| local binaryThreshold = torch.max(floorplan, 1)[1][lineMask:byte()]:mean() + 0.1 | |
| floorplanSegmentation, numSegments, floorplanBinary = utils.segmentFloorplan(floorplan, binaryThreshold, 0, false, true) | |
| end | |
| local doorMask = torch.cmul(floorplanSegmentation:gt(1):double(), image.dilate(lineMask)) | |
| for segment = 2, torch.max(floorplanSegmentation) do | |
| local segmentMask = floorplanSegmentation:eq(segment) | |
| local indices = doorMask[segmentMask]:nonzero() | |
| if ##indices > 0 then | |
| if (#indices)[1] < (#segmentMask:nonzero())[1] then | |
| doorMask[segmentMask] = 0 | |
| end | |
| end | |
| end | |
| doorMask = image.dilate(doorMask) | |
| local doors, numDoors = utils.findConnectedComponents(doorMask) | |
| local wallMask = utils.drawLineMask(width, height, walls, lineWidth, true) | |
| local doorIndex = 1 | |
| while doorIndex <= numDoors do | |
| local door = doors:eq(doorIndex) | |
| local indices = door:nonzero() | |
| if ##indices > 0 then | |
| local wallIndex = wallMask[indices[1]:totable()] | |
| local notInWall = wallMask:ne(wallIndex) | |
| local doorNotInWall = torch.cmul(door, notInWall) | |
| if ##doorNotInWall:nonzero() > 0 then | |
| numDoors = numDoors + 1 | |
| doors[doorNotInWall] = numDoors | |
| end | |
| end | |
| doorIndex = doorIndex + 1 | |
| end | |
| floorplanSegmentation = torch.cmul(floorplanSegmentation, (1 - lineMask):int()) | |
| local smallSegments = {} | |
| for segment = 2, torch.max(floorplanSegmentation) do | |
| local segmentMask = floorplanSegmentation:eq(segment) | |
| local indices = segmentMask:nonzero() | |
| if ##indices > 0 then | |
| local mins = torch.min(indices, 1)[1] | |
| local maxs = torch.max(indices, 1)[1] | |
| if (#indices)[1] < width * height * 0.03 and maxs[1] - mins[1] < height * 0.25 and maxs[2] - mins[2] < width * 0.25 then | |
| table.insert(smallSegments, segmentMask) | |
| end | |
| end | |
| end | |
| local rooms, numRooms = utils.findConnectedComponents(1 - lineMask) | |
| local backgroundRoomIndex | |
| local imageCorners = {{1, 1}, {width, 1}, {width, height}, {1, height}} | |
| for _, imageCorner in pairs(imageCorners) do | |
| local roomIndex = rooms[imageCorner[2]][imageCorner[1]] | |
| if roomIndex > 0 then | |
| if not backgroundRoomIndex then | |
| backgroundRoomIndex = roomIndex | |
| elseif roomIndex ~= backgroundRoomIndex then | |
| rooms[rooms:eq(roomIndex)] = backgroundRoomIndex | |
| end | |
| end | |
| end | |
| if not backgroundRoomIndex then | |
| backgroundRoomIndex = numRooms | |
| end | |
| local iconMask = torch.zeros(height, width) | |
| for roomIndex = 1, numRooms do | |
| if roomIndex ~= backgroundRoomIndex then | |
| local roomMask = rooms:eq(roomIndex) | |
| --local maxCount | |
| --local maxCountSegmentIndex | |
| if ##roomMask:nonzero() > 0 then | |
| local roomCount = (#roomMask:nonzero())[1] | |
| local iconSegments = {} | |
| for segmentIndex, segmentMask in pairs(smallSegments) do | |
| local indices = torch.cmul(segmentMask, roomMask):nonzero() | |
| if ##indices > 0 then | |
| local count = (#indices)[1] | |
| if count == (#segmentMask:nonzero())[1] and count < roomCount * 0.5 then | |
| table.insert(iconSegments, segmentIndex) | |
| end | |
| end | |
| end | |
| --table.insert(invalidSegments, maxCountSegmentIndex) | |
| for _, segmentIndex in pairs(iconSegments) do | |
| local indices = smallSegments[segmentIndex]:nonzero() | |
| local mins = torch.min(indices, 1)[1] | |
| local maxs = torch.max(indices, 1)[1] | |
| --local segmentMask = torch.zeros(height, width) | |
| --segmentMask[{{mins[1], maxs[1]}, {mins[2], maxs[2]}}] = 1 | |
| --iconMask[torch.cmul(segmentMask:byte(), roomMask)] = 0 | |
| iconMask[{{mins[1], maxs[1]}, {mins[2], maxs[2]}}] = 1 | |
| end | |
| end | |
| end | |
| end | |
| --image.save('test/mask_1.png', iconMask:double()) | |
| --iconMask = image.dilate(iconMask) | |
| local icons, numIcons = utils.findConnectedComponents(iconMask) | |
| local segmentation, labelMap = utils.predictSegmentation(floorplan, walls) | |
| --local icons = torch.zeros(#segmentation) | |
| local iconIndex = numIcons | |
| local labelSizeThresholds = {{5, 50}, {10, 100}} | |
| local roomLabels = {} | |
| for segmentIndex = 2, torch.max(segmentation) do | |
| local mask = segmentation:eq(segmentIndex) | |
| local label = labelMap[segmentIndex] | |
| if label <= 10 then | |
| local indices = mask:nonzero():double() | |
| if ##indices > 0 then | |
| local means = torch.mean(indices, 1)[1] | |
| local maxs = torch.max(indices, 1)[1] | |
| local mins = torch.min(indices, 1)[1] | |
| if maxs[1] - mins[1] >= labelSizeThresholds[1][1] and maxs[2] - mins[2] >= labelSizeThresholds[2][1] then | |
| local labelSize = {} | |
| for c = 1, 2 do | |
| labelSize[c] = math.min((maxs[c] - mins[c]) * 4 / 5, labelSizeThresholds[c][2]) | |
| end | |
| local roomIndex = rooms[indices[1][1]][indices[1][2]] | |
| if roomLabels[roomIndex] == nil then | |
| roomLabels[roomIndex] = {} | |
| end | |
| table.insert(roomLabels[roomIndex], {indices:size(1), {{means[2] - labelSize[2] / 2, means[1] - labelSize[1] / 2}, {means[2] + labelSize[2] / 2, means[1] + labelSize[1] / 2}, utils.getItemInfo('labels', label)}}) | |
| --table.insert(labels, {{means[2] - labelSize[2] / 2, means[1] - labelSize[1] / 2}, {means[2] + labelSize[2] / 2, means[1] + labelSize[1] / 2}, utils.getItemInfo('labels', label)}) | |
| end | |
| end | |
| else | |
| mask = torch.cmul(mask, 1 - iconMask:byte()) | |
| icons[mask] = iconIndex | |
| iconIndex = iconIndex + 1 | |
| end | |
| end | |
| local labels = {} | |
| for roomIndex, sortedLabels in pairs(roomLabels) do | |
| table.sort(sortedLabels, function(a, b) return a[1] > b[1] end) | |
| table.insert(labels, sortedLabels[1][2]) | |
| end | |
| iconMask = icons:gt(0) | |
| --icons = image.convolve(icons:double(), torch.ones(3, 3), 'same'):int() | |
| --image.save('test/binary.png', floorplanBinary:double()) | |
| --image.save('test/floorplan.png', floorplan) | |
| --image.save('test/segmentation.png', utils.drawSegmentation(floorplanSegmentation, torch.max(floorplanSegmentation))) | |
| --image.save('test/mask.png', utils.drawSegmentation(icons)) | |
| --image.save('test/mask.png', (1 - torch.cmul(doors:eq(0), icons:eq(0))):double()) | |
| --os.exit(1) | |
| --[[ | |
| if true then | |
| return 1 - torch.cmul(doors:eq(0), icons:eq(0)) | |
| end | |
| ]]-- | |
| local candidateRegions = {} | |
| candidateRegions.doors = doors | |
| candidateRegions.icons = icons | |
| --return candidateRegions, torch.clamp(floorplan * 0.2 + torch.cat(torch.cat(doors:eq(0):repeatTensor(1, 1, 1):double(), icons:eq(0):repeatTensor(1, 1, 1):double(), 1), torch.ones(1, height, width), 1), 0, 1) | |
| return candidateRegions, torch.cat(torch.cat(doors:eq(0):repeatTensor(1, 1, 1):double(), icons:eq(0):repeatTensor(1, 1, 1):double(), 1), torch.ones(1, height, width), 1), labels | |
| end | |
| function utils.getWallPlacements(width, height, representation, lineWidth, representationType) | |
| local lineWidth = lineWidth or 5 | |
| local representationType = representationType or 'L' | |
| local points | |
| if representationType == 'L' then | |
| representation.walls = utils.mergeLines(representation.walls, lineWidth) | |
| points = utils.linesToPoints(width, height, representation.walls, lineWidth) | |
| else | |
| points = representation.points | |
| end | |
| local walls, wallJunctionsMap = utils.pointsToLines(width, height, points, lineWidth, true) | |
| local imageSize = {width, height} | |
| --[[ | |
| local wallIndices = {6, 32} | |
| for _, wallIndex in pairs(wallIndices) do | |
| print(wallIndex) | |
| for index = 1, 2 do | |
| for c = 1, 2 do | |
| print(walls[wallIndex][index][c]) | |
| end | |
| end | |
| print(wallJunctionsMap[wallIndex][1] .. ' ' .. wallJunctionsMap[wallIndex][2]) | |
| end | |
| local pointIndices = {3, 20} | |
| for _, pointIndex in pairs(pointIndices) do | |
| print(pointIndex) | |
| for index = 1, 3, 2 do | |
| for c = 1, 3 do | |
| print(points[pointIndex][index][c]) | |
| end | |
| end | |
| --print(wallJunctionsMap[wallIndex][1] .. ' ' .. wallJunctionsMap[wallIndex][2]) | |
| end | |
| os.exit(1) | |
| for wallIndex, wall in pairs(representation.walls) do | |
| print(wallIndex) | |
| print(wall[1][1] .. ' ' .. wall[1][2] .. ' ' .. wall[2][1] .. ' ' .. wall[2][2]) | |
| end | |
| for pointIndex, point in pairs(points) do | |
| print(pointIndex) | |
| print(point[1][1] .. ' ' .. point[1][2] .. ' ' .. point[3][2] .. ' ' .. point[3][3]) | |
| end | |
| --os.exit(1) | |
| ]]-- | |
| local wallOrderMap = {} | |
| local orderPlacementTypeMap = {} | |
| orderPlacementTypeMap[0] = 0 | |
| local wallMask = utils.drawLineMask(width, height, walls, lineWidth, true, 1) | |
| local rooms, numRooms = utils.findConnectedComponents(1 - wallMask:gt(0)) | |
| local backgroundRoomIndex | |
| local imageCorners = {{1, 1}, {width, 1}, {width, height}, {1, height}} | |
| for _, imageCorner in pairs(imageCorners) do | |
| local roomIndex = rooms[imageCorner[2]][imageCorner[1]] | |
| if roomIndex > 0 then | |
| if not backgroundRoomIndex then | |
| backgroundRoomIndex = roomIndex | |
| elseif roomIndex ~= backgroundRoomIndex then | |
| rooms[rooms:eq(roomIndex)] = backgroundRoomIndex | |
| end | |
| end | |
| end | |
| if not backgroundRoomIndex then | |
| backgroundRoomIndex = numRooms | |
| end | |
| local backgroundMask = rooms:eq(backgroundRoomIndex) | |
| local backgroundMaskDiff = image.dilate(backgroundMask) - backgroundMask | |
| backgroundMaskDiff[{1, {}}] = 1 | |
| backgroundMaskDiff[{height, {}}] = 1 | |
| backgroundMaskDiff[{{}, 1}] = 1 | |
| backgroundMaskDiff[{{}, width}] = 1 | |
| --image.save('test/walls.png', utils.drawSegmentation(wallMask)) | |
| --image.save('test/rooms.png', utils.drawSegmentation(rooms)) | |
| --image.save('test/background.png', backgroundMask:double()) | |
| --image.save('test/background_diff.png', backgroundMaskDiff:double()) | |
| local wallMaskDiff = torch.cmul(wallMask, backgroundMaskDiff:double()) | |
| local outmostWalls = wallMask[backgroundMaskDiff] | |
| for i = 1, outmostWalls:size(1) do | |
| local wallIndex = outmostWalls[i] | |
| if wallIndex > 0 and not wallOrderMap[wallIndex] then | |
| wallOrderMap[wallIndex] = 0 | |
| --print(0 .. ' ' .. wallIndex) | |
| --[[ | |
| local lineDim = utils.lineDim(walls[wallIndex]) | |
| local indices = wallMaskDiff:eq(wallIndex):nonzero() | |
| local mins = torch.min(indices, 1)[1] | |
| local maxs = torch.max(indices, 1)[1] | |
| local diffLineDim = 0 | |
| if maxs[1] - mins[1] <= maxs[2] - mins[2] and maxs[1] - mins[1] <= 1 then | |
| diffLineDim = 1 | |
| elseif maxs[1] - mins[1] >= maxs[2] - mins[2] and maxs[2] - mins[2] <= 1 then | |
| diffLineDim = 2 | |
| end | |
| if lineDim == 0 or lineDim == diffLineDim then | |
| wallOrderMap[wallIndex] = 0 | |
| print(0 .. ' ' .. wallIndex) | |
| end | |
| ]]-- | |
| end | |
| end | |
| local junctionWallMap = {} | |
| for wallIndex, junctions in pairs(wallJunctionsMap) do | |
| for c = 1, 2 do | |
| if junctionWallMap[junctions[c]] == nil then | |
| junctionWallMap[junctions[c]] = {} | |
| end | |
| table.insert(junctionWallMap[junctions[c]], wallIndex) | |
| end | |
| end | |
| local wallLengthDimMap = {} | |
| for wallIndex, wall in pairs(walls) do | |
| wallLengthDimMap[wallIndex] = {utils.calcDistance(wall[1], wall[2]), utils.lineDim(wall)} | |
| end | |
| --[[ | |
| local wallNeighbors = {} | |
| for wallIndex, wall in pairs(walls) do | |
| if wallNeighbors[wallIndex] == nil then | |
| wallNeighbors[wallIndex] = {} | |
| end | |
| end | |
| for junctionIndex, wallIndices in pairs(junctionWallMap) do | |
| for _, wallIndex in pairs(wallIndices) do | |
| wallNeighbors[wallIndex][junctionIndex] = {} | |
| end | |
| end | |
| for junctionIndex, wallIndices in pairs(junctionWallMap) do | |
| for _, wallIndex_1 in pairs(wallIndices) do | |
| for _, wallIndex_2 in pairs(wallIndices) do | |
| if wallIndex_1 < wallIndex_2 then | |
| local lineDim_1 = utils.lineDim(walls[wallIndex_1]) | |
| local lineDim_2 = utils.lineDim(walls[wallIndex_2]) | |
| local neighborType | |
| if lineDim_1 == 0 or lineDim_2 == 0 then | |
| neighborType = 0 | |
| elseif lineDim_1 == lineDim_2 then | |
| neighborType = 1 | |
| else | |
| neighborType = 2 | |
| end | |
| wallNeighbors[wallIndex_1][junctionIndex][wallIndex_2] = neighborType | |
| wallNeighbors[wallIndex_2][junctionIndex][wallIndex_1] = neighborType | |
| end | |
| end | |
| end | |
| end | |
| ]]-- | |
| local order = 1 | |
| while true do | |
| local junctionTypeMap = {} | |
| for wallIndex, junctions in pairs(wallJunctionsMap) do | |
| if wallOrderMap[wallIndex] and wallOrderMap[wallIndex] < order then | |
| for c = 1, 2 do | |
| junctionTypeMap[junctions[c]] = 1 | |
| end | |
| end | |
| end | |
| local placements = {} | |
| for wallIndex, junctions in pairs(wallJunctionsMap) do | |
| if not wallOrderMap[wallIndex] then | |
| local wallPlacements = {} | |
| if junctionTypeMap[junctions[1]] == 1 and junctionTypeMap[junctions[2]] == 1 then | |
| table.insert(placements, {{wallIndex}, 0, 0}) | |
| elseif junctionTypeMap[junctions[1]] == 1 or junctionTypeMap[junctions[2]] == 1 then | |
| if junctionTypeMap[junctions[1]] ~= 1 then | |
| table.insert(wallPlacements, {{wallIndex}, junctions[1], 0, 0}) | |
| else | |
| table.insert(wallPlacements, {{wallIndex}, junctions[2], 0, 0}) | |
| end | |
| end | |
| while #wallPlacements > 0 do | |
| local newWallPlacements = {} | |
| for _, placement in pairs(wallPlacements) do | |
| local activeWallIndex = placement[1][#placement[1]] | |
| local lineDim = utils.lineDim(walls[activeWallIndex]) | |
| local unnecessaryTurn = 0 | |
| local lineDimCounter = {0, 0} | |
| lineDimCounter[lineDim] = 1 | |
| for _, neighborWallIndex in pairs(junctionWallMap[placement[2]]) do | |
| if neighborWallIndex ~= activeWallIndex then | |
| local neighborLineDim = utils.lineDim(walls[neighborWallIndex]) | |
| if neighborLineDim > 0 then | |
| if lineDimCounter[neighborLineDim] == 1 then | |
| unnecessaryTurn = 1 | |
| break | |
| end | |
| lineDimCounter[neighborLineDim] = lineDimCounter[neighborLineDim] + 1 | |
| end | |
| end | |
| end | |
| for _, neighborWallIndex in pairs(junctionWallMap[placement[2]]) do | |
| if neighborWallIndex ~= activeWallIndex then | |
| local turn = 0 | |
| if utils.lineDim(walls[neighborWallIndex]) ~= lineDim then | |
| turn = 1 | |
| end | |
| local numTurns = placement[3] + turn | |
| local numUnnecessaryTurns = placement[4] + math.min(unnecessaryTurn, turn) | |
| if numTurns <= 2 then | |
| local newWallSequence = {} | |
| for _, placementWallIndex in pairs(placement[1]) do | |
| table.insert(newWallSequence, placementWallIndex) | |
| end | |
| table.insert(newWallSequence, neighborWallIndex) | |
| local neighborWallActiveJunction | |
| local neighborWallJunctions = wallJunctionsMap[neighborWallIndex] | |
| for c = 1, 2 do | |
| local junction = neighborWallJunctions[c] | |
| if junction ~= placement[2] then | |
| neighborWallActiveJunction = junction | |
| end | |
| end | |
| if junctionTypeMap[neighborWallActiveJunction] == 1 then | |
| table.insert(placements, {newWallSequence, numTurns, numUnnecessaryTurns}) | |
| else | |
| table.insert(newWallPlacements, {newWallSequence, neighborWallActiveJunction, numTurns, numUnnecessaryTurns}) | |
| --[[ | |
| print(#newWallSequence) | |
| print(newWallSequence[1]) | |
| print(newWallSequence[2]) | |
| print(numTurns .. ' ' .. neighborWallActiveJunction) | |
| --]] | |
| end | |
| end | |
| end | |
| end | |
| end | |
| wallPlacements = newWallPlacements | |
| --os.exit(1) | |
| end | |
| end | |
| end | |
| if #placements == 0 then | |
| break | |
| end | |
| local minCost | |
| local minCostWallSequence | |
| local minCostNumTurns | |
| --print('num placements: ' .. #placements) | |
| for _, placement in pairs(placements) do | |
| local cost = placement[2] + placement[3] * 2 | |
| local lengths = {0, 0, 0} | |
| for __, wallIndex in pairs(placement[1]) do | |
| --print(_ .. ' ' .. __ .. ' ' .. wallIndex) | |
| --cost = cost + 1 / wallLengthMap[wallIndex] | |
| lengths[wallLengthDimMap[wallIndex][2] + 1] = lengths[wallLengthDimMap[wallIndex][2] + 1] + wallLengthDimMap[wallIndex][1] | |
| end | |
| local length = lengths[1] + lengths[2] + lengths[3] | |
| if length > 0 then | |
| cost = cost + 1 / length | |
| end | |
| if not minCost or cost < minCost then | |
| minCost = cost | |
| minCostWallSequence = placement[1] | |
| minCostNumTurns = placement[2] | |
| end | |
| --[[ | |
| if order == 6 or order == 7 then | |
| print('order: ' .. order) | |
| print('index: ' .. _) | |
| print('cost: ' .. cost) | |
| for __, wallIndex in pairs(placement[1]) do | |
| print(wallIndex) | |
| end | |
| end | |
| ]]-- | |
| end | |
| for _, wallIndex in pairs(minCostWallSequence) do | |
| --print(order .. ' ' .. wallIndex) | |
| wallOrderMap[wallIndex] = order | |
| end | |
| orderPlacementTypeMap[order] = minCostNumTurns + 1 | |
| order = order + 1 | |
| --if order == 3 then | |
| --os.exit(1) | |
| --end | |
| end | |
| local maxOrder = order - 1 | |
| --print(maxOrder) | |
| for order = 0, maxOrder do | |
| local orderWalls = {} | |
| for wallIndex, wall in pairs(walls) do | |
| if wallOrderMap[wallIndex] and wallOrderMap[wallIndex] <= order then | |
| table.insert(orderWalls, wall) | |
| end | |
| end | |
| --image.save('test/walls_' .. order .. '.png', utils.drawLineMask(width, height, orderWalls, lineWidth)) | |
| order = order + 1 | |
| end | |
| return walls, wallOrderMap, orderPlacementTypeMap, points, wallJunctionsMap | |
| end | |
| function utils.getWallPlacement(width, height, representation, lineWidth, representationType, getHeatmap) | |
| local lineWidth = lineWidth or 5 | |
| local walls, wallOrderMap, orderPlacementTypeMap = utils.getWallPlacements(width, height, representation, lineWidth, representationType) | |
| -- if not orderPlacementTypeMap[1] then | |
| -- image.save('test/floorplan.png', floorplan) | |
| -- image.save('test/walls.png', utils.drawSegmentation(utils.drawLineMask(floorplan:size(3), floorplan:size(2), representation.walls, lineWidth, true))) | |
| -- --print(wallOrderMap) | |
| -- --assert(false, 'placement type empty') | |
| -- end | |
| local inputWalls = {} | |
| local outputWalls = {} | |
| local inputOrder = 0 | |
| local outputOrder = 1 | |
| for wallIndex, wall in pairs(walls) do | |
| if wallOrderMap[wallIndex] then | |
| if wallOrderMap[wallIndex] == inputOrder then | |
| table.insert(inputWalls, wall) | |
| elseif wallOrderMap[wallIndex] == outputOrder then | |
| table.insert(outputWalls, wall) | |
| end | |
| end | |
| end | |
| local inputWallMask = utils.drawLineMask(width, height, inputWalls, lineWidth) | |
| local outputWallRepresentation | |
| if getHeatmap then | |
| local outputWallJunctionHeatmap = torch.zeros(height, width) | |
| if orderPlacementTypeMap[outputOrder] then | |
| local kernelSize = kernelSize or 7 | |
| for _, wall in pairs(outputWalls) do | |
| for pointIndex = 1, 2 do | |
| local point = wall[pointIndex] | |
| local x = torch.round(point[1]) | |
| local y = torch.round(point[2]) | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| outputWallJunctionHeatmap[y][x] = 1 | |
| end | |
| end | |
| end | |
| local kernel = image.gaussian(kernelSize) | |
| outputWallJunctionHeatmap = image.convolve(outputWallJunctionHeatmap, kernel, 'same') | |
| end | |
| outputWallJunctionHeatmap = outputWallJunctionHeatmap:repeatTensor(1, 1, 1) | |
| outputWallRepresentation = outputWallJunctionHeatmap | |
| else | |
| local maxNumJunctions = 4 | |
| local outputWallJunctions = torch.zeros(maxNumJunctions * 2) | |
| if orderPlacementTypeMap[outputOrder] then | |
| local junctions = {} | |
| for _, wall in pairs(outputWalls) do | |
| for pointIndex = 1, 2 do | |
| local point = wall[pointIndex] | |
| table.insert(junctions, point) | |
| end | |
| end | |
| local sortedJunctions = {} | |
| for _, junction in pairs(junctions) do | |
| local exists = false | |
| for __, sortedJunction in pairs(sortedJunctions) do | |
| if math.abs(junction[1] - sortedJunction[1]) <= lineWidth and math.abs(junction[2] - sortedJunction[2]) <= lineWidth then | |
| sortedJunction[1] = (junction[1] + sortedJunction[1]) / 2 | |
| sortedJunction[2] = (junction[2] + sortedJunction[2]) / 2 | |
| exists = true | |
| end | |
| end | |
| if not exists then | |
| table.insert(sortedJunctions, junction) | |
| end | |
| end | |
| table.sort(sortedJunctions, function(a, b) return a[2] < b[2] or (a[2] == b[2] and a[1] < b[1]) end) | |
| for index = 1, math.min(maxNumJunctions, #sortedJunctions) do | |
| for c = 1, 2 do | |
| outputWallJunctions[(index - 1) * 2 + c] = sortedJunctions[index][c] | |
| end | |
| end | |
| end | |
| outputWallRepresentation = outputWallJunctions | |
| end | |
| if not orderPlacementTypeMap[outputOrder] then | |
| orderPlacementTypeMap[outputOrder] = 4 | |
| end | |
| return inputWallMask, outputWallRepresentation, orderPlacementTypeMap[outputOrder] | |
| end | |
| function utils.getJunctionPlacement(width, height, representation, lineWidth, representationType, kernelSize) | |
| local lineWidth = lineWidth or 5 | |
| local kernelSize = kernelSize or 7 | |
| local walls, wallOrderMap, orderPlacementTypeMap, points, wallJunctionsMap = utils.getWallPlacements(width, height, representation, lineWidth, representationType) | |
| -- if not orderPlacementTypeMap[1] then | |
| -- image.save('test/floorplan.png', floorplan) | |
| -- image.save('test/walls.png', utils.drawSegmentation(utils.drawLineMask(floorplan:size(3), floorplan:size(2), representation.walls, lineWidth, true))) | |
| -- --print(wallOrderMap) | |
| -- --assert(false, 'placement type empty') | |
| -- end | |
| local orderWallsMap = {} | |
| for wallIndex, order in pairs(wallOrderMap) do | |
| if not orderWallsMap[order] then | |
| orderWallsMap[order] = {} | |
| end | |
| --print(order .. ' ' .. wallIndex .. ' ' .. wallJunctionsMap[wallIndex][1] .. ' ' .. wallJunctionsMap[wallIndex][2]) | |
| table.insert(orderWallsMap[order], wallIndex) | |
| end | |
| --print(orderWallsMap[0]) | |
| local orderedJunctions = {} | |
| local activeOrder = -1 | |
| local activeJunction = -1 | |
| local usedJunctionMask = {} | |
| local firstJunction = 0 | |
| while true do | |
| local nextJunctionFound = false | |
| if activeOrder >= 0 then | |
| local orderWalls = orderWallsMap[activeOrder] | |
| for _, wallIndex in pairs(orderWalls) do | |
| for pointIndex = 1, 2 do | |
| if wallJunctionsMap[wallIndex][pointIndex] == activeJunction and not usedJunctionMask[wallJunctionsMap[wallIndex][3 - pointIndex]] then | |
| nextJunctionFound = true | |
| activeJunction = wallJunctionsMap[wallIndex][3 - pointIndex] | |
| usedJunctionMask[activeJunction] = true | |
| --print(activeOrder .. ' ' .. activeJunction) | |
| table.insert(orderedJunctions, {activeJunction, activeOrder}) | |
| break | |
| end | |
| end | |
| if nextJunctionFound then | |
| break | |
| end | |
| end | |
| end | |
| if not nextJunctionFound then | |
| if firstJunction > 0 then | |
| table.insert(orderedJunctions, {firstJunction, activeOrder}) | |
| nextJunctionFound = true | |
| firstJunction = 0 | |
| end | |
| end | |
| if not nextJunctionFound then | |
| activeOrder = activeOrder + 1 | |
| usedJunctionMask = {} | |
| local orderWalls = orderWallsMap[activeOrder] | |
| if not orderWalls then | |
| break | |
| end | |
| local minValue | |
| local minValueWallIndex | |
| local minValuePointIndex | |
| for _, wallIndex in pairs(orderWalls) do | |
| for pointIndex = 1, 2 do | |
| local junction = wallJunctionsMap[wallIndex][pointIndex] | |
| local value = points[junction][1][1] * height + points[junction][1][2] | |
| if not minValue or value < minValue then | |
| minValue = value | |
| minValueWallIndex = wallIndex | |
| minValuePointIndex = pointIndex | |
| end | |
| end | |
| end | |
| activeJunction = wallJunctionsMap[minValueWallIndex][minValuePointIndex] | |
| usedJunctionMask[activeJunction] = true | |
| --print(activeOrder .. ' ' .. activeJunction) | |
| table.insert(orderedJunctions, {activeJunction, activeOrder}) | |
| if activeOrder == 0 then | |
| firstJunction = activeJunction | |
| end | |
| end | |
| end | |
| local orderedPoints = {} | |
| --image.save('test/walls.png', utils.drawLineMask(width, height, walls, lineWidth)) | |
| for index, junction in pairs(orderedJunctions) do | |
| table.insert(orderedPoints, points[junction[1]]) | |
| if index >= 2 then | |
| local orderedWalls = utils.pointsToLines(width, height, orderedPoints, lineWidth, true, nil, true) | |
| --image.save('test/walls_' .. index .. '.png', utils.drawLineMask(width, height, orderedWalls, lineWidth)) | |
| end | |
| end | |
| assert(activeOrder > 0 and #orderedJunctions > 1) | |
| local index = math.random(#orderedJunctions - 1) | |
| while orderedJunctions[index][2] ~= orderedJunctions[index + 1][2] do | |
| index = math.random(#orderedJunctions - 1) | |
| end | |
| local previousPoints = {} | |
| for i = 1, index do | |
| table.insert(previousPoints, points[orderedJunctions[i][1]]) | |
| end | |
| local wallMask = utils.drawLineMask(width, height, utils.pointsToLines(width, height, previousPoints, lineWidth, true, nil, true), lineWidth) | |
| local activeJunction = orderedJunctions[index][1] | |
| local activeJunctionMask = torch.zeros(height, width) | |
| activeJunctionMask[math.max(math.min(torch.round(points[activeJunction][1][2]), height), 1)][math.max(math.min(torch.round(points[activeJunction][1][1]), width), 1)] = 1 | |
| local nextJunction = orderedJunctions[index + 1][1] | |
| local nextJunctionMask = torch.zeros(height, width) | |
| nextJunctionMask[math.max(math.min(torch.round(points[nextJunction][1][2]), height), 1)][math.max(math.min(torch.round(points[nextJunction][1][1]), width), 1)] = 1 | |
| if kernelSize > 0 then | |
| local kernel = image.gaussian(kernelSize) | |
| activeJunctionMask = image.convolve(activeJunctionMask, kernel, 'same') | |
| nextJunctionMask = image.convolve(nextJunctionMask, kernel, 'same') | |
| end | |
| return wallMask, activeJunctionMask, nextJunctionMask | |
| end | |
| function utils.getIconTensor(width, height, representation) | |
| local nameIconsMap = {} | |
| for _, icon in pairs(representation.icons) do | |
| local name = icon[3][1] | |
| if not nameIconsMap[name] then | |
| nameIconsMap[name] = {} | |
| end | |
| table.insert(nameIconsMap[name], icon) | |
| end | |
| local iconTensor = torch.zeros(4, 5) | |
| local iconNums = {} | |
| iconNums['bathtub'] = 1 | |
| iconNums['cooking_counter'] = 1 | |
| iconNums['toilet'] = 1 | |
| iconNums['entrance'] = 1 | |
| --numberMap.icons['washing_basin'] = {6, 7, 8, 9} | |
| --numberMap.icons['special'] = {10, 11, 12} | |
| --numberMap.icons['stairs'] = {13} | |
| local index = 1 | |
| for name, num in pairs(iconNums) do | |
| for i = 1, num do | |
| if nameIconsMap[name] and nameIconsMap[name][i] then | |
| local icon = nameIconsMap[name][i] | |
| iconTensor[index][1] = (icon[1][1] + icon[2][1]) / 2 / width | |
| iconTensor[index][2] = (icon[1][2] + icon[2][2]) / 2 / height | |
| iconTensor[index][3] = math.abs(icon[1][1] - icon[2][1]) / width | |
| iconTensor[index][4] = math.abs(icon[1][2] - icon[2][2]) / height | |
| iconTensor[index][5] = 1 | |
| index = index + 1 | |
| end | |
| end | |
| end | |
| return iconTensor | |
| end | |
| function utils.getSegmentationWithConflict(width, height, representation, lineWidth) | |
| local lineWidth = lineWidth or 5 | |
| local lineMask = utils.drawLineMask(width, height, representation.walls, lineWidth) | |
| local rooms, numRooms = utils.findConnectedComponents(1 - lineMask) | |
| local backgroundRoomIndex | |
| local imageCorners = {{1, 1}, {width, 1}, {width, height}, {1, height}} | |
| for _, imageCorner in pairs(imageCorners) do | |
| local roomIndex = rooms[imageCorner[2]][imageCorner[1]] | |
| if roomIndex > 0 then | |
| if not backgroundRoomIndex then | |
| backgroundRoomIndex = roomIndex | |
| elseif roomIndex ~= backgroundRoomIndex then | |
| rooms[rooms:eq(roomIndex)] = backgroundRoomIndex | |
| end | |
| end | |
| end | |
| local floorplanSegmentation = torch.zeros(#rooms) | |
| for _, label in pairs(representation.labels) do | |
| local center = {torch.round((label[1][1] + label[2][1]) / 2), torch.round((label[1][2] + label[2][2]) / 2)} | |
| local number = utils.getNumber('labels', label[3]) | |
| local roomIndex = rooms[{center[2], center[1]}] | |
| if roomIndex > 0 and roomIndex ~= backgroundRoomIndex then | |
| local segmentIndex = floorplanSegmentation[{center[2], center[1]}] | |
| if segmentIndex == 0 or (number < segmentIndex and (number ~= 2 or segmentIndex ~= 3)) or (number == 3 and segmentIndex == 2) then | |
| floorplanSegmentation[rooms:eq(roomIndex)] = number | |
| end | |
| end | |
| end | |
| floorplanSegmentation[rooms:eq(backgroundRoomIndex)] = 0 | |
| for _, icon in pairs(representation.icons) do | |
| local number = utils.getNumber('icons', icon[3]) | |
| floorplanSegmentation[{{math.max(math.min(icon[1][2], icon[2][2]), 1), math.min(math.max(icon[1][2], icon[2][2]), height)}, {math.max(math.min(icon[1][1], icon[2][1]), 1), math.min(math.max(icon[1][1], icon[2][1]), width)}}] = number + 10 | |
| --[[ | |
| local success, _ = pcall(function() | |
| floorplanSegmentation[{{math.max(icon[1][2], 1), math.min(icon[2][2], height)}, {math.max(icon[1][1], 1), math.min(icon[2][1], width)}}] = number + 10 | |
| end | |
| ) | |
| if not success then | |
| print(width) | |
| print(height) | |
| print(icon[1][1]) | |
| print(icon[1][2]) | |
| print(icon[2][1]) | |
| print(icon[2][2]) | |
| end | |
| ]]-- | |
| end | |
| floorplanSegmentation[rooms:eq(0)] = -1 | |
| --image.save('test/floorplan_segmentation.png', utils.drawSegmentation(floorplanSegmentation)) | |
| --os.exit(1) | |
| local doorWidth = doorWidth or 3 | |
| local doorMask = utils.drawLineMask(width, height, representation.doors, doorWidth, nil, 0) | |
| floorplanSegmentation[doorMask:byte()] = -2 | |
| return floorplanSegmentation | |
| end | |
| function utils.getRoomSegmentationQP(width, height, representation, lineWidth, floorplan) | |
| local lineWidth = lineWidth or 5 | |
| representation.walls = utils.mergeLines(representation.walls, lineWidth) | |
| for _, wall in pairs(representation.walls) do | |
| local lineDim = utils.lineDim(wall) | |
| if lineDim > 0 then | |
| if wall[1][lineDim] > wall[2][lineDim] then | |
| print('opposite wall exists') | |
| os.exit(1) | |
| end | |
| end | |
| end | |
| local points = utils.linesToPoints(width, height, representation.walls, lineWidth) | |
| local wallMask = utils.drawLineMask(width, height, representation.walls, lineWidth) | |
| --image.save('test/wall.png', wallMask) | |
| local rooms, numRooms = utils.findConnectedComponents(1 - wallMask) | |
| local backgroundRoomIndex | |
| local imageCorners = {{width, 1}, {width, height}, {1, height}} | |
| for _, imageCorner in pairs(imageCorners) do | |
| local roomIndex = rooms[imageCorner[2]][imageCorner[1]] | |
| if roomIndex > 0 and not backgroundRoomIndex then | |
| backgroundRoomIndex = roomIndex | |
| end | |
| if roomIndex > 0 and roomIndex ~= backgroundRoomIndex then | |
| rooms[rooms:eq(roomIndex)] = backgroundRoomIndex | |
| end | |
| end | |
| if not backgroundRoomIndex then | |
| backgroundRoomIndex = numRooms | |
| end | |
| local roomLabels = {} | |
| for _, label in pairs(representation.labels) do | |
| local center = {torch.round((label[1][1] + label[2][1]) / 2), torch.round((label[1][2] + label[2][2]) / 2)} | |
| --local number = utils.getNumber('labels', label[3]) | |
| local roomIndex = rooms[{center[2], center[1]}] | |
| if roomIndex > 0 and roomIndex ~= backgroundRoomIndex then | |
| if not roomLabels[roomIndex] then | |
| roomLabels[roomIndex] = {} | |
| end | |
| table.insert(roomLabels[roomIndex], label) | |
| end | |
| end | |
| utils.saveRepresentation('test/points.txt', {points = points}) | |
| local segmentations = torch.zeros(10, height, width) | |
| local radius = 5 | |
| for _, label in pairs(representation.labels) do | |
| local number = utils.getNumber('labels', label[3]) | |
| local center = {torch.round((label[1][1] + label[2][1]) / 2), torch.round((label[1][2] + label[2][2]) / 2)} | |
| local minX = math.max(center[1] - radius, 1) | |
| local minY = math.max(center[2] - radius, 1) | |
| local maxX = math.min(center[1] + radius, width) | |
| local maxY = math.min(center[2] + radius, height) | |
| segmentations[{number, {minY, maxY}, {minX, maxX}}] = 1 | |
| end | |
| for roomIndex, labels in pairs(roomLabels) do | |
| if #labels == 1 then | |
| local number = utils.getNumber('labels', labels[1][3]) | |
| segmentations[number][rooms:eq(roomIndex)] = 1 | |
| elseif #labels == 0 then | |
| segmentations[10][rooms:eq(roomIndex)] = 1 | |
| end | |
| end | |
| for segmentIndex = 1, 10 do | |
| image.save('test/segment_' .. segmentIndex .. '.png', segmentations[segmentIndex]) | |
| end | |
| image.save('test/segment_0.png', torch.zeros(height, width)) | |
| py.execute('import os') | |
| local ret = py.execute('os.system("python ../InverseCAD/PostProcessing/QP_segmentation.py")') | |
| local floorplanSegmentation = torch.zeros(height, width) | |
| for segmentIndex = 1, 10 do | |
| local success, segmentImage = pcall(function() | |
| return image.load('test/segment_result_' .. segmentIndex .. '.png', 1) | |
| end) | |
| if not success then | |
| image.save('test/floorplan.png', floorplan) | |
| os.exit(1) | |
| return utils.getRoomSegmentationHeuristic(width, height, representation, lineWidth, floorplan) | |
| end | |
| floorplanSegmentation[segmentImage:gt(0.5)] = segmentIndex | |
| end | |
| floorplanSegmentation[rooms:eq(backgroundRoomIndex)] = 0 | |
| floorplanSegmentation[wallMask:gt(0)] = -1 | |
| for roomIndex, labels in pairs(roomLabels) do | |
| if #labels == 1 then | |
| local number = utils.getNumber('labels', labels[1][3]) | |
| floorplanSegmentation[rooms:eq(roomIndex)] = number | |
| elseif #labels == 0 then | |
| floorplanSegmentation[rooms:eq(roomIndex)] = 10 | |
| end | |
| end | |
| local hasConflict = false | |
| for roomIndex, labels in pairs(roomLabels) do | |
| if #labels > 1 then | |
| hasConflict = true | |
| break | |
| end | |
| end | |
| if hasConflict then | |
| image.save('test/floorplan.png', floorplan) | |
| representation.points = points | |
| --image.save('test/representation.png', utils.drawRepresentationImage(floorplan:size(3), floorplan:size(2), nil, nil, floorplan, representation, 'P', 'L')) | |
| image.save('test/segmentation.png', utils.drawSegmentation(floorplanSegmentation)) | |
| os.exit(1) | |
| end | |
| return floorplanSegmentation:repeatTensor(1, 1, 1) | |
| end | |
| function utils.getRoomSegments(width, height, representation, lineWidth, floorplan, getHeatmaps) | |
| local lineWidth = lineWidth or 5 | |
| lineWidth = 3 | |
| --representation.walls = utils.mergeLines(representation.walls, lineWidth) | |
| --representation.walls = utils.pointsToLines(width, height, representation.points, lineWidth, true) | |
| --local lineMask = utils.drawLineMask(width, height, lines, lineWidth) | |
| local boundaryMask = utils.drawLineMask(width, height, representation.walls, lineWidth) | |
| local rooms, numRooms = utils.findConnectedComponents(1 - boundaryMask) | |
| image.save('test/walls.png', boundaryMask) | |
| image.save('test/rooms.png', utils.drawSegmentation(rooms)) | |
| local backgroundRoomIndex | |
| local imageCorners = {{1, 1}, {width, 1}, {width, height}, {1, height}} | |
| for _, imageCorner in pairs(imageCorners) do | |
| local roomIndex = rooms[imageCorner[2]][imageCorner[1]] | |
| if roomIndex > 0 then | |
| if not backgroundRoomIndex then | |
| backgroundRoomIndex = roomIndex | |
| elseif roomIndex ~= backgroundRoomIndex then | |
| rooms[rooms:eq(roomIndex)] = backgroundRoomIndex | |
| end | |
| end | |
| end | |
| if not backgroundRoomIndex then | |
| backgroundRoomIndex = numRooms | |
| end | |
| local roomLabels = {} | |
| for _, label in pairs(representation.labels) do | |
| local center = {torch.round((label[1][1] + label[2][1]) / 2), torch.round((label[1][2] + label[2][2]) / 2)} | |
| --local number = utils.getNumber('labels', label[3]) | |
| local roomIndex = rooms[{center[2], center[1]}] | |
| if roomIndex > 0 and roomIndex ~= backgroundRoomIndex then | |
| if not roomLabels[roomIndex] then | |
| roomLabels[roomIndex] = {} | |
| end | |
| table.insert(roomLabels[roomIndex], label) | |
| end | |
| end | |
| local roomSegments = {} | |
| for roomIndex, labels in pairs(roomLabels) do | |
| local roomMask = rooms:eq(roomIndex) | |
| if roomMask:nonzero():size(1) > lineWidth * lineWidth * 2 then | |
| for _, label in pairs(labels) do | |
| local roomLabel = utils.getNumber('labels', label[3]) | |
| table.insert(roomSegments, {roomMask, roomLabel}) | |
| end | |
| end | |
| --[[ | |
| local roomLabel | |
| if #labels == 1 then | |
| roomLabel = utils.getNumber('labels', labels[1][3]) | |
| elseif #labels > 1 then | |
| for _, label in pairs(labels) do | |
| local number = utils.getNumber('labels', label[3]) | |
| if not roomLabel or (number < roomLabel and roomLabel ~= 3) or (number == 3 and roomLabel == 2) then | |
| roomLabel = number | |
| end | |
| end | |
| end | |
| local roomMask = rooms:eq(roomIndex) | |
| if roomLabel and roomMask:nonzero():size(1) > lineWidth * lineWidth * 2 then | |
| table.insert(roomSegments, {roomMask, roomLabel}) | |
| end | |
| ]]-- | |
| end | |
| return roomSegments | |
| end | |
| function utils.getSegmentation(width, height, representation, lineWidth, floorplan, getHeatmaps) | |
| local lineWidth = lineWidth or 5 | |
| --representation.walls = utils.mergeLines(representation.walls, lineWidth) | |
| --local boundaryMask = utils.drawLineMask(width, height, representation.walls, lineWidth) | |
| local points | |
| if representation.points then | |
| points = representation.points | |
| else | |
| points = utils.linesToPoints(width, height, representation.walls, lineWidth) | |
| end | |
| --local shortWalls, wallJunctionsMap = utils.pointsToLines(width, height, points, lineWidth, true, lineWidth) | |
| local shortWalls = representation.walls | |
| local wallJunctionsMap = {} | |
| for _, wall in pairs(shortWalls) do | |
| local junctions = {} | |
| for index = 1, 2 do | |
| local minDistance = math.max(width, height) | |
| local minDistancePointIndex | |
| for pointIndex, point in pairs(points) do | |
| point = point[1] | |
| local distance = utils.calcDistance(point, wall[index]) | |
| if distance < minDistance then | |
| minDistancePointIndex = pointIndex | |
| minDistance = distance | |
| end | |
| end | |
| table.insert(junctions, minDistancePointIndex) | |
| end | |
| table.insert(wallJunctionsMap, junctions) | |
| end | |
| for _, wall in pairs(shortWalls) do | |
| break | |
| local lineDim = utils.lineDim(wall, lineWidth) | |
| local wallFixedValue = (wall[1][3 - lineDim] + wall[2][3 - lineDim]) / 2 | |
| local wallMinValue = math.min(wall[1][lineDim], wall[2][lineDim]) | |
| local wallMaxValue = math.max(wall[1][lineDim], wall[2][lineDim]) | |
| for _, longWall in pairs(representation.walls) do | |
| if utils.lineDim(longWall, lineWidth) == lineDim then | |
| local longWallFixedValue = (longWall[1][3 - lineDim] + longWall[2][3 - lineDim]) / 2 | |
| local longWallMinValue = math.min(longWall[1][lineDim], longWall[2][lineDim]) | |
| local longWallMaxValue = math.max(longWall[1][lineDim], longWall[2][lineDim]) | |
| if math.abs(longWallFixedValue - wallFixedValue) <= lineWidth and longWallMinValue - lineWidth <= wallMinValue and longWallMaxValue + lineWidth >= wallMaxValue then | |
| if #longWall == 4 then | |
| table.insert(shortWalls[_], longWall[4]) | |
| break | |
| end | |
| end | |
| end | |
| end | |
| end | |
| local wallMask = utils.drawLineMask(width, height, shortWalls, lineWidth, true, 0) | |
| local boundaryMask = utils.drawLineMask(width, height, shortWalls, lineWidth) | |
| --image.save('test/boundary.png', boundaryMask) | |
| --image.save('test/walls.png', wallMask) | |
| local rooms, numRooms = utils.findConnectedComponents(1 - boundaryMask) | |
| local backgroundRoomIndex | |
| local imageCorners = {{1, 1}, {width, 1}, {width, height}, {1, height}} | |
| for _, imageCorner in pairs(imageCorners) do | |
| local roomIndex = rooms[imageCorner[2]][imageCorner[1]] | |
| if roomIndex > 0 then | |
| if not backgroundRoomIndex then | |
| backgroundRoomIndex = roomIndex | |
| elseif roomIndex ~= backgroundRoomIndex then | |
| rooms[rooms:eq(roomIndex)] = backgroundRoomIndex | |
| end | |
| end | |
| end | |
| if not backgroundRoomIndex then | |
| backgroundRoomIndex = numRooms | |
| end | |
| local floorplanSegmentation = torch.zeros(#rooms) | |
| floorplanSegmentation[rooms:eq(backgroundRoomIndex)] = 11 | |
| --floorplanSegmentation[rooms:eq(backgroundRoomIndex)] = 0 | |
| local roomLabels = {} | |
| for _, label in pairs(representation.labels) do | |
| local center = {torch.round((label[1][1] + label[2][1]) / 2), torch.round((label[1][2] + label[2][2]) / 2)} | |
| --local number = utils.getNumber('labels', label[3]) | |
| local roomIndex = rooms[{center[2], center[1]}] | |
| if roomIndex > 0 and roomIndex ~= backgroundRoomIndex then | |
| if not roomLabels[roomIndex] then | |
| roomLabels[roomIndex] = {} | |
| end | |
| table.insert(roomLabels[roomIndex], label) | |
| end | |
| end | |
| local hasConflict = false | |
| if true then | |
| for roomIndex, labels in pairs(roomLabels) do | |
| if #labels == 1 then | |
| floorplanSegmentation[rooms:eq(roomIndex)] = utils.getNumber('labels', labels[1][3]) | |
| elseif #labels > 1 then | |
| hasConflict = true | |
| local majorNumber | |
| for _, label in pairs(labels) do | |
| local number = utils.getNumber('labels', label[3]) | |
| if not majorNumber or (number < majorNumber and majorNumber ~= 3) or (majorNumber == 2 and number == 3) then | |
| majorNumber = number | |
| end | |
| end | |
| local inferiorLabels = {} | |
| for _, label in pairs(labels) do | |
| local number = utils.getNumber('labels', label[3]) | |
| if number ~= majorNumber then | |
| table.insert(inferiorLabels, label) | |
| end | |
| end | |
| local roomMask = rooms:eq(roomIndex) | |
| roomMask = image.dilate(roomMask) | |
| local roomWallsTensor = wallMask[roomMask] | |
| local roomWallsIndices = roomWallsTensor:nonzero() | |
| assert(##roomWallsIndices > 0, 'conflict room without a wall') | |
| roomWallsTensor = roomWallsTensor:index(1, roomWallsIndices[{{}, 1}]) | |
| local roomWalls = {} | |
| for i = 1, roomWallsTensor:size(1) do | |
| roomWalls[roomWallsTensor[i]] = true | |
| end | |
| local junctionWallsMap = {} | |
| for wallIndex, _ in pairs(roomWalls) do | |
| for index = 1, 2 do | |
| local junctionIndex = wallJunctionsMap[wallIndex][index] | |
| if not junctionWallsMap[junctionIndex] then | |
| junctionWallsMap[junctionIndex] = {} | |
| end | |
| table.insert(junctionWallsMap[junctionIndex], wallIndex) | |
| end | |
| end | |
| for _, label in pairs(inferiorLabels) do | |
| local labelPoint = {(label[1][1] + label[2][1]) / 2, (label[1][2] + label[2][2]) / 2} | |
| local junctionWallPairMap = {} | |
| local wallNeighbors = {} | |
| for junction, walls in pairs(junctionWallsMap) do | |
| if #walls == 2 and utils.lineDim(shortWalls[walls[1]]) == utils.lineDim(shortWalls[walls[2]]) then | |
| for index = 1, 2 do | |
| if not wallNeighbors[walls[index]] then | |
| wallNeighbors[walls[index]] = {} | |
| end | |
| table.insert(wallNeighbors[walls[index]], walls[3 - index]) | |
| end | |
| end | |
| if #walls >= 2 then | |
| local junctionPoint = points[junction][1] | |
| for _, wallIndex in pairs(walls) do | |
| local lineDim = utils.lineDim(shortWalls[wallIndex]) | |
| local wallPoint | |
| for _, pointIndex in pairs(wallJunctionsMap[wallIndex]) do | |
| if pointIndex ~= junction then | |
| wallPoint = points[pointIndex][1] | |
| end | |
| end | |
| for _, neighborWallIndex in pairs(walls) do | |
| if neighborWallIndex > wallIndex and utils.lineDim(shortWalls[neighborWallIndex]) + lineDim == 3 then | |
| local neighborWallPoint | |
| for _, pointIndex in pairs(wallJunctionsMap[neighborWallIndex]) do | |
| if pointIndex ~= junction then | |
| neighborWallPoint = points[pointIndex][1] | |
| end | |
| end | |
| if (labelPoint[3 - lineDim] - junctionPoint[3 - lineDim]) * (neighborWallPoint[3 - lineDim] - junctionPoint[3 - lineDim]) > 0 and (labelPoint[lineDim] - junctionPoint[lineDim]) * (wallPoint[lineDim] - junctionPoint[lineDim]) > 0 then | |
| junctionWallPairMap[junction] = {wallIndex, neighborWallIndex} | |
| elseif #walls == 2 and math.abs(labelPoint[3 - lineDim] - junctionPoint[3 - lineDim]) >= math.abs(labelPoint[3 - lineDim] - neighborWallPoint[3 - lineDim]) - lineWidth then | |
| for index = 1, 2 do | |
| if not wallNeighbors[walls[index]] then | |
| wallNeighbors[walls[index]] = {} | |
| end | |
| table.insert(wallNeighbors[walls[index]], walls[3 - index]) | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| local wallGroups = {} | |
| local wallGroupMap = {} | |
| local visitedWalls = {} | |
| for wallIndex, _ in pairs(wallNeighbors) do | |
| if not visitedWalls[wallIndex] then | |
| local wallGroup = {} | |
| wallGroup[wallIndex] = true | |
| visitedWalls[wallIndex] = true | |
| while true do | |
| local hasChange = false | |
| for groupWall, _ in pairs(wallGroup) do | |
| for _, neighbor in pairs(wallNeighbors[groupWall]) do | |
| if not wallGroup[neighbor] then | |
| wallGroup[neighbor] = true | |
| visitedWalls[neighbor] = true | |
| hasChange = true | |
| end | |
| end | |
| end | |
| if not hasChange then | |
| break | |
| end | |
| end | |
| local group = {} | |
| for groupWall, _ in pairs(wallGroup) do | |
| table.insert(group, groupWall) | |
| wallGroupMap[groupWall] = #wallGroups + 1 | |
| end | |
| table.insert(wallGroups, group) | |
| end | |
| end | |
| for wallIndex, _ in pairs(roomWalls) do | |
| if not wallGroupMap[wallIndex] then | |
| wallGroupMap[wallIndex] = #wallGroups + 1 | |
| table.insert(wallGroups, {wallIndex}) | |
| end | |
| end | |
| --[[ | |
| if label[3][1] == 'corridor' then | |
| for i = 1, #wallGroups do | |
| local testingWalls = {} | |
| for _, wallIndex in pairs(wallGroups[i]) do | |
| table.insert(testingWalls, shortWalls[wallIndex]) | |
| end | |
| image.save('test/walls_' .. i .. '.png', utils.drawLineMask(width, height, testingWalls, lineWidth)) | |
| end | |
| os.exit(1) | |
| end | |
| ]]-- | |
| local candidateRegions = {} | |
| for junction, wallPair in pairs(junctionWallPairMap) do | |
| for neighborJunction, neighborWallPair in pairs(junctionWallPairMap) do | |
| if neighborJunction < junction then | |
| for index = 1, 2 do | |
| for neighborIndex = 1, 2 do | |
| if wallGroupMap[wallPair[index]] == wallGroupMap[neighborWallPair[neighborIndex]] then | |
| local lineDim = utils.lineDim(shortWalls[wallPair[3 - index]]) | |
| local wallGroup = wallGroups[wallGroupMap[wallPair[3 - index]]] | |
| local minValue | |
| local maxValue | |
| for _, wallIndex in pairs(wallGroup) do | |
| local wall = shortWalls[wallIndex] | |
| if not minValue or wall[1][lineDim] < minValue then | |
| minValue = wall[1][lineDim] | |
| end | |
| if not maxValue or wall[2][lineDim] > maxValue then | |
| maxValue = wall[2][lineDim] | |
| end | |
| end | |
| --local minValue = math.min(shortWalls[wallPair[3 - index]][1][lineDim], shortWalls[wallPair[3 - index]][2][lineDim]) | |
| --local maxValue = math.max(shortWalls[wallPair[3 - index]][1][lineDim], shortWalls[wallPair[3 - index]][2][lineDim]) | |
| local neighborWallGroup = wallGroups[wallGroupMap[neighborWallPair[3 - neighborIndex]]] | |
| local neighborMinValue | |
| local neighborMaxValue | |
| for _, wallIndex in pairs(neighborWallGroup) do | |
| local wall = shortWalls[wallIndex] | |
| if not neighborMinValue or wall[1][lineDim] < neighborMinValue then | |
| neighborMinValue = wall[1][lineDim] | |
| end | |
| if not neighborMaxValue or wall[2][lineDim] > neighborMaxValue then | |
| neighborMaxValue = wall[2][lineDim] | |
| end | |
| end | |
| local minValue = math.max(minValue, neighborMinValue) | |
| local maxValue = math.min(maxValue, neighborMaxValue) | |
| local fixedValue_1 = shortWalls[wallPair[3 - index]][1][3 - lineDim] | |
| local fixedValue_2 = shortWalls[neighborWallPair[3 - neighborIndex]][1][3 - lineDim] | |
| fixedValue_1, fixedValue_2 = math.min(fixedValue_1, fixedValue_2), math.max(fixedValue_1, fixedValue_2) | |
| if minValue + lineWidth < maxValue and fixedValue_1 + lineWidth < fixedValue_2 then | |
| if lineDim == 1 then | |
| table.insert(candidateRegions, {{minValue, fixedValue_1}, {maxValue, fixedValue_2}}) | |
| else | |
| table.insert(candidateRegions, {{fixedValue_1, minValue}, {fixedValue_2, maxValue}}) | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| if #candidateRegions > 0 then | |
| --print(labelPoint[1] .. ' ' .. labelPoint[2]) | |
| local minDistance | |
| local selectedRegion | |
| for _, region in pairs(candidateRegions) do | |
| local center = {(region[1][1] + region[2][1]) / 2, (region[1][2] + region[2][2]) / 2} | |
| local distance = utils.calcDistance(center, labelPoint) | |
| --print(distance) | |
| --print(region[1][1] .. ' ' .. region[1][2] .. ' ' .. region[2][1] .. ' ' .. region[2][2]) | |
| if not minDistance or distance < minDistance then | |
| selectedRegion = region | |
| minDistance = distance | |
| end | |
| end | |
| for pointIndex = 1, 2 do | |
| selectedRegion[pointIndex][1] = math.min(math.max(selectedRegion[pointIndex][1], 1), width) | |
| selectedRegion[pointIndex][2] = math.min(math.max(selectedRegion[pointIndex][2], 1), height) | |
| end | |
| floorplanSegmentation[{{selectedRegion[1][2], selectedRegion[2][2]}, {selectedRegion[1][1], selectedRegion[2][1]}}] = utils.getNumber('labels', label[3]) | |
| roomMask[{{selectedRegion[1][2], selectedRegion[2][2]}, {selectedRegion[1][1], selectedRegion[2][1]}}] = 0 | |
| end | |
| end | |
| floorplanSegmentation[roomMask] = majorNumber | |
| end | |
| end | |
| for roomIndex = 1, numRooms - 1 do | |
| break | |
| if roomIndex ~= backgroundRoomIndex and not roomLabels[roomIndex] then | |
| floorplanSegmentation[rooms:eq(roomIndex)] = 10 | |
| end | |
| end | |
| else | |
| for roomIndex = 1, numRooms - 1 do | |
| local roomMask = rooms:eq(roomIndex) | |
| local roomSegmentation = torch.ones(height, width) * 11 | |
| local majorNumber | |
| roomMask = image.dilate(roomMask) | |
| local roomWallsTensor = wallMask[roomMask] | |
| local roomWallsIndices = roomWallsTensor:nonzero() | |
| assert(##roomWallsIndices > 0, 'conflict room without a wall') | |
| roomWallsTensor = roomWallsTensor:index(1, roomWallsIndices[{{}, 1}]) | |
| local roomWalls = {} | |
| for i = 1, roomWallsTensor:size(1) do | |
| roomWalls[roomWallsTensor[i]] = true | |
| end | |
| for wallIndex_1, _ in pairs(roomWalls) do | |
| for wallIndex_2, __ in pairs(roomWalls) do | |
| if wallIndex_1 < wallIndex_2 then | |
| local wall_1 = shortWalls[wallIndex_1] | |
| local wall_2 = shortWalls[wallIndex_2] | |
| local lineDim_1 = utils.lineDim(wall_1) | |
| local lineDim_2 = utils.lineDim(wall_2) | |
| if lineDim_1 == lineDim_2 and lineDim_1 > 0 and #wall_1 == 4 and #wall_2 == 4 then | |
| local lineDim = lineDim_1 | |
| local fixedValue_1, fixedValue_2, roomLabel | |
| local minValue = math.max(wall_1[1][lineDim], wall_2[1][lineDim]) | |
| local maxValue = math.min(wall_1[2][lineDim], wall_2[2][lineDim]) | |
| if minValue < maxValue then | |
| --print({wall_1, wall_2}) | |
| if wall_1[1][3 - lineDim] < wall_2[1][3 - lineDim] and wall_1[4][2] == wall_2[4][1] then | |
| roomLabel = wall_1[4][2] | |
| fixedValue_1 = wall_1[1][3 - lineDim] | |
| fixedValue_2 = wall_2[1][3 - lineDim] | |
| end | |
| if wall_1[1][3 - lineDim] > wall_2[1][3 - lineDim] and wall_1[4][1] == wall_2[4][2] then | |
| roomLabel = wall_1[4][1] | |
| fixedValue_1 = wall_2[1][3 - lineDim] | |
| fixedValue_2 = wall_1[1][3 - lineDim] | |
| end | |
| if roomLabel then | |
| if not majorNumber or (roomLabel < majorNumber and majorNumber ~= 3) or (majorNumber == 2 and roomLabel == 3) then | |
| majorNumber = roomLabel | |
| end | |
| local invalidMask = false | |
| for wallIndex, _ in pairs(roomWalls) do | |
| if wallIndex ~= wallIndex_1 and wallIndex ~= wallIndex_2 then | |
| local wall = shortWalls[wallIndex] | |
| if utils.lineDim(wall) == lineDim and wall[1][3 - lineDim] > fixedValue_1 and wall[1][3 - lineDim] < fixedValue_2 and math.max(wall[1][lineDim], minValue) < math.min(wall[2][lineDim], maxValue) then | |
| invalidMask = true | |
| break | |
| end | |
| end | |
| end | |
| if not invalidMask then | |
| local mask = torch.zeros(height, width) | |
| mask:narrow(3 - lineDim, minValue, maxValue - minValue + 1):narrow(lineDim, fixedValue_1, fixedValue_2 - fixedValue_1 + 1):fill(1) | |
| mask = mask:byte() | |
| mask:cmul(roomSegmentation:gt(roomLabel)) | |
| if roomLabel == 2 then | |
| mask:csub(roomSegmentation:eq(3)) | |
| mask = mask:gt(0) | |
| elseif roomLabel == 3 then | |
| mask:add(roomSegmentation:eq(2)) | |
| mask = mask:gt(0) | |
| end | |
| --image.save('test/mask_' .. roomLabel .. '_' .. wallIndex_1 .. '_' .. wallIndex_2 .. '.png', mask:double()) | |
| roomSegmentation[mask] = roomLabel | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| local mask = torch.cmul(roomMask, roomSegmentation:eq(11)) | |
| if roomIndex ~= backgroundRoomIndex then | |
| if majorNumber then | |
| roomSegmentation[mask] = majorNumber | |
| else | |
| roomSegmentation[mask] = 10 | |
| end | |
| end | |
| --image.save('test/segmentation_' .. roomIndex .. '.png', utils.drawSegmentation(roomSegmentation)) | |
| floorplanSegmentation[roomMask] = roomSegmentation[roomMask] | |
| end | |
| end | |
| if getHeatmaps then | |
| --image.save('test/floorplan.png', floorplan) | |
| --image.save('test/segmentation.png', utils.drawSegmentation(floorplanSegmentation)) | |
| local heatmaps = torch.zeros(30, height, width) | |
| for segmentIndex = 1, 11 do | |
| heatmaps[segmentIndex]:copy(floorplanSegmentation:eq(segmentIndex):double()) | |
| end | |
| local indexOffset = 11 | |
| for _, icon in pairs(representation.icons) do | |
| local number = utils.getNumber('icons', icon[3]) | |
| heatmaps[indexOffset + number][{{math.max(math.min(icon[1][2], icon[2][2]), 1), math.min(math.max(icon[1][2], icon[2][2]), height)}, {math.max(math.min(icon[1][1], icon[2][1]), 1), math.min(math.max(icon[1][1], icon[2][1]), width)}}] = 1 | |
| end | |
| indexOffset = indexOffset + 10 | |
| local groupWalls = {} | |
| for _, wall in pairs(representation.walls) do | |
| local group = wall[3][2] | |
| if not groupWalls[group] then | |
| groupWalls[group] = {} | |
| end | |
| table.insert(groupWalls[group], wall) | |
| end | |
| for group, walls in pairs(groupWalls) do | |
| local lineMask = utils.drawLineMask(width, height, walls, lineWidth):byte() | |
| heatmaps[indexOffset + group][lineMask] = 1 | |
| end | |
| indexOffset = indexOffset + 2 | |
| local doorWidth = doorWidth or lineWidth - 2 | |
| local groupDoors = {} | |
| local doorNumberGroupMap = {} | |
| doorNumberGroupMap[1] = 1 | |
| doorNumberGroupMap[2] = 1 | |
| doorNumberGroupMap[3] = 2 | |
| doorNumberGroupMap[4] = 3 | |
| doorNumberGroupMap[5] = 3 | |
| doorNumberGroupMap[6] = 4 | |
| doorNumberGroupMap[7] = 5 | |
| doorNumberGroupMap[8] = 5 | |
| doorNumberGroupMap[9] = 6 | |
| doorNumberGroupMap[10] = 2 | |
| doorNumberGroupMap[11] = 2 | |
| for _, door in pairs(representation.doors) do | |
| local number = door[3][2] | |
| local group = doorNumberGroupMap[number] | |
| if not groupDoors[group] then | |
| groupDoors[group] = {} | |
| end | |
| table.insert(groupDoors[group], door) | |
| end | |
| for group, doors in pairs(groupDoors) do | |
| local lineMask = utils.drawLineMask(width, height, doors, doorWidth):byte() | |
| heatmaps[indexOffset + group][lineMask] = 1 | |
| end | |
| return heatmaps | |
| end | |
| if true then | |
| local indexOffset = 11 | |
| local groupWalls = {} | |
| for _, wall in pairs(representation.walls) do | |
| local group = wall[3][2] | |
| if not groupWalls[group] then | |
| groupWalls[group] = {} | |
| end | |
| table.insert(groupWalls[group], wall) | |
| end | |
| for group = 2, 1, -1 do | |
| local walls = groupWalls[group] | |
| if walls then | |
| local lineMask = utils.drawLineMask(width, height, walls, lineWidth):byte() | |
| floorplanSegmentation[lineMask] = indexOffset + group | |
| end | |
| end | |
| --[[ | |
| for group, walls in pairs(groupWalls) do | |
| local lineMask = utils.drawLineMask(width, height, walls, lineWidth):byte() | |
| floorplanSegmentation[lineMask] = indexOffset + group | |
| end | |
| ]]-- | |
| --floorplanSegmentation[floorplanSegmentation:eq(0)] = 11 | |
| local floorplanSegmentation_2 = torch.zeros(#floorplanSegmentation) | |
| local indexOffset = 0 | |
| for _, icon in pairs(representation.icons) do | |
| local number = utils.getNumber('icons', icon[3]) | |
| floorplanSegmentation_2[{{math.max(math.min(icon[1][2], icon[2][2], height), 1), math.min(math.max(icon[1][2], icon[2][2], 1), height)}, {math.max(math.min(icon[1][1], icon[2][1], width), 1), math.min(math.max(icon[1][1], icon[2][1], 1), width)}}] = indexOffset + number | |
| end | |
| indexOffset = indexOffset + 10 | |
| local doorWidth = doorWidth or lineWidth - 2 | |
| local groupDoors = {} | |
| local doorNumberGroupMap = {} | |
| doorNumberGroupMap[1] = 1 | |
| doorNumberGroupMap[2] = 1 | |
| doorNumberGroupMap[3] = 2 | |
| doorNumberGroupMap[4] = 3 | |
| doorNumberGroupMap[5] = 3 | |
| doorNumberGroupMap[6] = 4 | |
| doorNumberGroupMap[7] = 5 | |
| doorNumberGroupMap[8] = 5 | |
| doorNumberGroupMap[9] = 6 | |
| doorNumberGroupMap[10] = 2 | |
| doorNumberGroupMap[11] = 2 | |
| for _, door in pairs(representation.doors) do | |
| local number = door[3][2] | |
| local group = doorNumberGroupMap[number] | |
| if not groupDoors[group] then | |
| groupDoors[group] = {} | |
| end | |
| table.insert(groupDoors[group], door) | |
| end | |
| for group, doors in pairs(groupDoors) do | |
| local lineMask = utils.drawLineMask(width, height, doors, doorWidth):byte() | |
| floorplanSegmentation_2[lineMask] = indexOffset + group | |
| end | |
| indexOffset = indexOffset + 6 | |
| --floorplanSegmentation_2[floorplanSegmentation_2:eq(0)] = indexOffset + 1 | |
| return torch.cat(floorplanSegmentation:repeatTensor(1, 1, 1), floorplanSegmentation_2:repeatTensor(1, 1, 1), 1), shortWalls | |
| end | |
| local indexOffset = 11 | |
| for _, icon in pairs(representation.icons) do | |
| local number = utils.getNumber('icons', icon[3]) | |
| floorplanSegmentation[{{math.max(math.min(icon[1][2], icon[2][2]), 1), math.min(math.max(icon[1][2], icon[2][2]), height)}, {math.max(math.min(icon[1][1], icon[2][1]), 1), math.min(math.max(icon[1][1], icon[2][1]), width)}}] = indexOffset + number | |
| end | |
| indexOffset = indexOffset + 10 | |
| local groupWalls = {} | |
| for _, wall in pairs(representation.walls) do | |
| local group = wall[3][2] | |
| if not groupWalls[group] then | |
| groupWalls[group] = {} | |
| end | |
| table.insert(groupWalls[group], wall) | |
| end | |
| for group, walls in pairs(groupWalls) do | |
| local lineMask = utils.drawLineMask(width, height, walls, lineWidth):byte() | |
| floorplanSegmentation[lineMask] = indexOffset + group | |
| end | |
| indexOffset = indexOffset + 2 | |
| local doorWidth = doorWidth or lineWidth - 2 | |
| local groupDoors = {} | |
| local doorNumberGroupMap = {} | |
| doorNumberGroupMap[1] = 1 | |
| doorNumberGroupMap[2] = 1 | |
| doorNumberGroupMap[3] = 2 | |
| doorNumberGroupMap[4] = 3 | |
| doorNumberGroupMap[5] = 3 | |
| doorNumberGroupMap[6] = 4 | |
| doorNumberGroupMap[7] = 5 | |
| doorNumberGroupMap[8] = 5 | |
| doorNumberGroupMap[9] = 6 | |
| doorNumberGroupMap[10] = 2 | |
| doorNumberGroupMap[11] = 2 | |
| for _, door in pairs(representation.doors) do | |
| local number = door[3][2] | |
| local group = doorNumberGroupMap[number] | |
| if not groupDoors[group] then | |
| groupDoors[group] = {} | |
| end | |
| table.insert(groupDoors[group], door) | |
| end | |
| for group, doors in pairs(groupDoors) do | |
| local lineMask = utils.drawLineMask(width, height, doors, doorWidth):byte() | |
| floorplanSegmentation[lineMask] = indexOffset + group | |
| end | |
| if hasConflict and false then | |
| image.save('test/floorplan.png', floorplan) | |
| representation.points = points | |
| --image.save('test/representation.png', utils.drawRepresentationImage(floorplan:size(3), floorplan:size(2), nil, nil, floorplan, representation, 'P', 'L')) | |
| image.save('test/segmentation.png', utils.drawSegmentation(floorplanSegmentation)) | |
| os.exit(1) | |
| end | |
| return floorplanSegmentation:repeatTensor(1, 1, 1), shortWalls | |
| end | |
| function utils.getPatch(floorplan, representation, opt, mode, head, regression, numLabels) | |
| local lineWidth = opt.lineWidth or 5 | |
| local shift = opt.shift or lineWidth | |
| local width, height = floorplan:size(3), floorplan:size(2) | |
| local numberItems = {} | |
| local numbers = {} | |
| if mode == 'points' then | |
| representation.points = utils.linesToPoints(width, height, representation.walls, lineWidth) | |
| end | |
| for _, item in pairs(representation[mode]) do | |
| local number = utils.getNumber(mode, item[3]) | |
| if not numberItems[number] then | |
| numberItems[number] = {} | |
| table.insert(numbers, number) | |
| end | |
| table.insert(numberItems[number], {{math.min(item[1][1], item[2][1]), math.min(item[1][2], item[2][2])}, {math.max(item[1][1], item[2][1]), math.max(item[1][2], item[2][2])}}) | |
| end | |
| if head == 1 and #numbers == 0 then | |
| return nil, nil | |
| end | |
| local negativeProb = 0.5 | |
| if numLabels then | |
| negativeProb = 1.0 / (numLabels + 1) | |
| end | |
| if head == 2 and (torch.uniform() < negativeProb or #numbers == 0) then | |
| local cx, cy, w, h | |
| for i = 1, 100 do | |
| cx = torch.random(width) | |
| cy = torch.random(height) | |
| local scale = 2 ^ torch.uniform(-1, 2) | |
| w = torch.round(opt.patchDim * scale) | |
| h = torch.round(opt.patchDim * scale) | |
| local isNegativeSample = true | |
| for number, items in pairs(numberItems) do | |
| for _, item in pairs(items) do | |
| if math.abs((item[1][1] + item[2][1]) / 2 - cx) < opt.shift and math.abs((item[1][2] + item[2][2]) / 2 - cy) < opt.shift then | |
| isNegativeSample = false | |
| break | |
| elseif item[1][1] > cx - w / 2 and item[1][2] > cy - h / 2 and item[2][1] < cx + w / 2 and item[2][2] < cy + h / 2 then | |
| isNegativeSample = false | |
| break | |
| end | |
| end | |
| if not isNegativeSample then | |
| break | |
| end | |
| end | |
| if isNegativeSample then | |
| break | |
| end | |
| end | |
| --[[ | |
| local success, _ = pcall(function() | |
| local patch = image.crop(floorplan, torch.round(math.max(cx - w / 2, 0)), torch.round(math.max(cy - h / 2, 0)), torch.round(math.min(cx + w / 2, width)), torch.round(math.min(cy + h / 2, height))) | |
| end) | |
| if not success then | |
| print(#floorplan) | |
| print(cx .. ' ' .. cy .. ' ' .. w .. ' ' .. h) | |
| os.exit(1) | |
| end | |
| ]]-- | |
| local patch = image.crop(floorplan, torch.round(math.max(cx - w / 2, 0)), torch.round(math.max(cy - h / 2, 0)), torch.round(math.min(cx + w / 2, width)), torch.round(math.min(cy + h / 2, height))) | |
| patch = image.scale(patch, opt.patchDim, opt.patchDim) | |
| patch:mul(2):add(-1) | |
| local label = -torch.ones(1) | |
| if numLabels then | |
| label:fill(numLabels + 1) | |
| end | |
| return patch, label | |
| end | |
| local numberSelected = numbers[torch.random(#numbers)] | |
| local itemsSelected = numberItems[numberSelected] | |
| local itemIndexSelected = torch.random(#itemsSelected) | |
| local itemSelected = itemsSelected[itemIndexSelected] | |
| local cx = itemSelected[1][1] + (itemSelected[2][1] - itemSelected[1][1]) * torch.uniform() | |
| local cy = itemSelected[1][2] + (itemSelected[2][2] - itemSelected[1][2]) * torch.uniform() | |
| local centerFound = false | |
| for attempt = 1, 10 do | |
| local closest = true | |
| local distanceSelected = (((itemSelected[2][1] + itemSelected[1][1]) / 2 - cx)^2 + ((itemSelected[2][2] + itemSelected[1][2]) / 2 - cy)^2)^0.5 | |
| for number, items in pairs(numberItems) do | |
| for itemIndex, item in pairs(items) do | |
| if number ~= numberSelected or itemIndex ~= itemIndexSelected then | |
| local distance = (((item[2][1] + item[1][1]) / 2 - cx)^2 + ((item[2][2] + item[1][2]) / 2 - cy)^2)^0.5 | |
| if 0.8 * distance < distanceSelected then | |
| closest = false | |
| break | |
| end | |
| end | |
| end | |
| if not closest then | |
| break | |
| end | |
| end | |
| if closest then | |
| centerFound = true | |
| break | |
| end | |
| cx = itemSelected[1][1] + math.max(itemSelected[2][1] - itemSelected[1][1], shift) * torch.uniform() | |
| cy = itemSelected[1][2] + math.max(itemSelected[2][2] - itemSelected[1][2], shift) * torch.uniform() | |
| end | |
| if not centerFound then | |
| return nil, nil | |
| end | |
| local item = itemSelected | |
| local maxDim = math.max(item[2][1] - item[1][1], item[2][2] - item[1][2], lineWidth) | |
| local minScale = opt.minScale or 1 | |
| local maxScale = opt.maxScale or 2.5 | |
| if head == 2 and regression then | |
| maxScale = 2 | |
| end | |
| if model == 'points' then | |
| minScale = 2 | |
| maxScale = 3 | |
| end | |
| local w = math.max(math.min(maxDim * (2 ^ torch.uniform(minScale, maxScale) - 1), cx * 2, cy * 2, (width - cx) * 2, (height - cy) * 2), cx - item[1][1], cy - item[1][2], item[2][1] - cx, item[2][2] - cy, lineWidth) | |
| local h = w | |
| if maxDim == 0 then | |
| return nil, nil | |
| end | |
| --print(maxDim) | |
| --print(cx .. ' ' .. cy .. ' ' .. w .. ' ' .. h) | |
| --print(item[1][1] .. ' ' .. item[1][2] .. ' ' .. item[2][1] .. ' ' .. item[2][2]) | |
| local patch = image.crop(floorplan, torch.round(math.max(cx - w / 2, 0)), torch.round(math.max(cy - h / 2, 0)), torch.round(math.min(cx + w / 2, width)), torch.round(math.min(cy + h / 2, height))) | |
| patch = image.scale(patch, opt.patchDim, opt.patchDim) | |
| patch:mul(2):add(-1) | |
| if head == 1 then | |
| if not regression then | |
| local mask = torch.zeros(height, width) | |
| mask:narrow(1, item[1][2], item[2][2] - item[1][2] + 1):narrow(2, item[1][1], item[2][1] - item[1][1] + 1):fill(1) | |
| mask = image.crop(mask, torch.round(math.max(cx - w / 2, 0)), torch.round(math.max(cy - h / 2, 0)), torch.round(math.min(cx + w / 2, width)), torch.round(math.min(cy + h / 2, height))) | |
| mask = image.scale(mask, opt.outputDim, opt.outputDim, 'simple'):gt(0.5):double() | |
| mask:mul(2):add(-1) | |
| if torch.uniform() < 0.5 then | |
| patch = image.hflip(patch) | |
| mask = image.hflip(mask) | |
| end | |
| return patch, mask | |
| else | |
| local cxBox = (item[1][1] + item[2][1]) / 2 - torch.round(math.max(cx - w / 2, 0)) | |
| local cyBox = (item[1][2] + item[2][2]) / 2 - torch.round(math.max(cy - h / 2, 0)) | |
| local wBox = item[2][1] - item[1][1] | |
| local hBox = item[2][2] - item[1][2] | |
| wBox = math.max(wBox, lineWidth) | |
| hBox = math.max(hBox, lineWidth) | |
| local wPatch = torch.round(math.min(cx + w / 2, width)) - torch.round(math.max(cx - w / 2, 0)) | |
| local hPatch = torch.round(math.min(cy + h / 2, height)) - torch.round(math.max(cy - h / 2, 0)) | |
| local coordinates = torch.Tensor({cxBox / wPatch, cyBox / hPatch, math.log(wBox / wPatch), math.log(hBox / hPatch)}) | |
| --local lineDim = utils.lineDim(item) | |
| --if lineDim > 0 then | |
| --coordinates[5 - lineDim] = | |
| --end | |
| if torch.uniform() < 0.5 then | |
| patch = image.hflip(patch) | |
| coordinates[1] = 1 - coordinates[1] | |
| end | |
| return patch, coordinates | |
| end | |
| else | |
| local label = torch.ones(1) | |
| if numLabels then | |
| label:fill(numberSelected) | |
| end | |
| return patch, label | |
| end | |
| end | |
| function utils.getPatchGrid(floorplan, representation, opt, mode, head, regression, numLabels) | |
| local lineWidth = opt.lineWidth or 5 | |
| local oriWidth, oriHeight = floorplan:size(3), floorplan:size(2) | |
| local sampleDim = opt.sampleDim or 448 | |
| local gridDim = sampleDim / 16 | |
| local patchDim = opt.patchDim or 16 * 7 | |
| local minScale = opt.minScale or 0 | |
| local maxScale = opt.maxScale or 1 | |
| local scale = 2^torch.uniform(minScale, maxScale) | |
| local width = torch.round(oriWidth * scale) | |
| local height = torch.round(oriHeight * scale) | |
| local sample = image.scale(floorplan, width, height) | |
| representation = utils.scaleRepresentation(representation, oriWidth, oriHeight, width, height) | |
| if width < sampleDim and height < sampleDim then | |
| local newSample = torch.zeros(3, sampleDim, sampleDim) | |
| newSample[{{}, {1, height}, {1, width}}]:copy(sample) | |
| sample = newSample | |
| elseif width < sampleDim then | |
| local newSample = torch.zeros(3, height, sampleDim) | |
| newSample[{{}, {1, height}, {1, width}}]:copy(sample) | |
| sample = newSample | |
| elseif height < sampleDim then | |
| local newSample = torch.zeros(3, sampleDim, width) | |
| newSample[{{}, {1, height}, {1, width}}]:copy(sample) | |
| sample = newSample | |
| end | |
| width = sample:size(3) | |
| height = sample:size(2) | |
| local x = 0 | |
| if width > sampleDim then | |
| x = torch.random(width - sampleDim) | |
| end | |
| local y = 0 | |
| if height > sampleDim then | |
| y = torch.random(height - sampleDim) | |
| end | |
| --print(#sample) | |
| --print(x .. ' ' .. y) | |
| --print(sampleDim) | |
| sample = image.crop(sample, x, y, x + sampleDim, y + sampleDim) | |
| representation = utils.cropRepresentation(representation, x, y, x + sampleDim, y + sampleDim) | |
| local items = representation[mode] | |
| local cellWidth = sampleDim / gridDim | |
| local cellHeight = sampleDim / gridDim | |
| local gridTensor = torch.zeros(6, gridDim, gridDim) | |
| for gridY = 1, gridDim do | |
| for gridX = 1, gridDim do | |
| local cx = cellWidth * (gridX - 1 + 0.5) | |
| local cy = cellHeight * (gridY - 1 + 0.5) | |
| local w = patchDim | |
| local h = patchDim | |
| local maxIOU | |
| local maxIOUCoordinates | |
| local maxIOUNumber | |
| local sumIOU = 0 | |
| local atCenter = false | |
| for _, item in pairs(items) do | |
| if math.max(item[1][1], cx - w / 2) <= math.min(item[2][1], cx + w / 2) and math.max(item[1][2], cy - h / 2) <= math.min(item[2][2], cy + h / 2) then | |
| local intersection = (math.min(cx + w / 2, item[2][1] + 1) - math.max(cx - w / 2, item[1][1])) * (math.min(cy + h / 2, item[2][2] + 1) - math.max(cy - h / 2, item[1][2])) | |
| --local union = w * h + (item[2][1] - item[1][1]) * (item[2][2] - item[1][2]) - intersection | |
| local union = (item[2][1] - item[1][1] + 1) * (item[2][2] - item[1][2] + 1) | |
| local IOU = intersection / union | |
| if not maxIOU or IOU > maxIOU then | |
| maxIOU = IOU | |
| local x_1 = math.max(item[1][1], cx - w / 2) | |
| local x_2 = math.min(item[2][1], cx + w / 2) | |
| local y_1 = math.max(item[1][2], cy - h / 2) | |
| local y_2 = math.min(item[2][2], cy + h / 2) | |
| local boxX = ((x_1 + x_2) / 2 - (cx - w / 2)) / w | |
| local boxY = ((y_1 + y_2) / 2 - (cy - h / 2)) / h | |
| local boxW = math.log(math.max(x_2 - x_1 + 1, lineWidth) / w) | |
| local boxH = math.log(math.max(y_2 - y_1 + 1, lineWidth) / h) | |
| maxIOUCoordinates = torch.Tensor({boxX, boxY, boxW, boxH}) | |
| maxIOUNumber = utils.getNumber(mode, item[3]) | |
| if item[1][1] < cx and item[2][1] > cx and item[1][2] < cy and item[2][2] > cy then | |
| atCenter = true | |
| end | |
| end | |
| sumIOU = sumIOU + IOU | |
| end | |
| end | |
| if maxIOU then | |
| --[[ | |
| print(maxIOUNumber) | |
| print(gridX .. ' ' .. gridY) | |
| print(maxIOU .. ' ' .. sumIOU) | |
| print(maxIOUCoordinates) | |
| ]]-- | |
| if not atCenter then | |
| maxIOU = math.min(maxIOU / sumIOU, maxIOU) | |
| end | |
| gridTensor[{{1, 4}, gridY, gridX}] = maxIOUCoordinates | |
| gridTensor[{5, gridY, gridX}] = maxIOU | |
| gridTensor[{6, gridY, gridX}] = maxIOUNumber | |
| if maxIOU > 0.95 and false then | |
| print(gridTensor[{{}, gridY, gridX}]) | |
| local x_1 = math.max(0, cx - w / 2) | |
| local x_2 = math.min(sampleDim, cx + w / 2) | |
| local y_1 = math.max(0, cy - h / 2) | |
| local y_2 = math.min(sampleDim, cy + h / 2) | |
| image.save('test/sample_' .. maxIOUNumber .. '.png', image.crop(sample, x_1, y_1, x_2, y_2)) | |
| end | |
| end | |
| end | |
| end | |
| if false then | |
| representation.points = utils.linesToPoints(sample:size(3), sample:size(2), representation.walls, lineWidth) | |
| image.save('test/floorplan.png', sample) | |
| image.save('test/representation.png', utils.drawRepresentationImage(sample:size(3), sample:size(2), nil, nil, sample, representation, 'P', 'L')) | |
| os.exit(1) | |
| end | |
| sample:mul(2):add(-1) | |
| return sample, gridTensor | |
| end | |
| function utils.drawJunctionMasks(width, height, points, maxNumPoints, pointLineLength, kernelSize) | |
| local junctionMasks = torch.zeros(maxNumPoints, height, width) | |
| local pointLineLength = pointLineLength or 15 | |
| local kernelSize = kernelSize or 7 | |
| local pointOrientations = {} | |
| for pointIndex, point in pairs(points) do | |
| local orientations = {} | |
| local orientation = point[3][3] | |
| if point[3][2] == 1 then | |
| table.insert(orientations, (orientation + 2 - 1) % 4 + 1) | |
| elseif point[3][2] == 2 then | |
| table.insert(orientations, orientation) | |
| table.insert(orientations, (orientation + 3 - 1) % 4 + 1) | |
| elseif point[3][2] == 3 then | |
| for i = 1, 4 do | |
| if i ~= orientation then | |
| table.insert(orientations, i) | |
| end | |
| end | |
| else | |
| for i = 1, 4 do | |
| table.insert(orientations, i) | |
| end | |
| end | |
| pointOrientations[pointIndex] = orientations | |
| end | |
| local deltas = {{0, -1}, {1, 0}, {0, 1}, {-1, 0}} | |
| local maxNumPoints = maxNumPoints or #points | |
| for pointIndex, point in pairs(points) do | |
| if pointIndex > maxNumPoints then | |
| break | |
| end | |
| for _, orientation in pairs(pointOrientations[pointIndex]) do | |
| local x = torch.round(point[1][1]) | |
| local y = torch.round(point[1][2]) | |
| local delta = deltas[orientation] | |
| x = math.max(math.min(x, width), 1) | |
| y = math.max(math.min(y, height), 1) | |
| for i = 1, pointLineLength do | |
| if x < 1 or x > width or y < 1 or y > height then | |
| break | |
| end | |
| junctionMasks[pointIndex][y][x] = 1 | |
| x = x + delta[1] | |
| y = y + delta[2] | |
| end | |
| end | |
| end | |
| if kernelSize > 0 then | |
| local kernel = image.gaussian(kernelSize) | |
| for i = 1, junctionMasks:size(1) do | |
| junctionMasks[i] = image.convolve(junctionMasks[i], kernel, 'same') | |
| end | |
| end | |
| return junctionMasks | |
| end | |
| function utils.getRoomJunctions(width, height, representation, lineWidth, maxNumPoints) | |
| local lineWidth = lineWidth or 5 | |
| local walls = utils.pointsToLines(width, height, representation.points, lineWidth) | |
| local wallMask = utils.drawLineMask(width, height, walls, lineWidth) | |
| local rooms, numRooms = utils.findConnectedComponents(1 - wallMask) | |
| numRooms = numRooms - 1 | |
| local backgroundRoomIndex | |
| local imageCorners = {{1, 1}, {width, 1}, {width, height}, {1, height}} | |
| for _, imageCorner in pairs(imageCorners) do | |
| local roomIndex = rooms[imageCorner[2]][imageCorner[1]] | |
| if roomIndex > 0 then | |
| if not backgroundRoomIndex then | |
| backgroundRoomIndex = roomIndex | |
| elseif roomIndex ~= backgroundRoomIndex then | |
| rooms[rooms:eq(roomIndex)] = backgroundRoomIndex | |
| end | |
| end | |
| end | |
| if not backgroundRoomIndex then | |
| backgroundRoomIndex = numRooms | |
| end | |
| local roomIndexMap = {} | |
| roomIndexMap[rooms[1][1]] = 0 | |
| for _, label in pairs(representation.labels) do | |
| local center = {torch.round((label[1][1] + label[2][1]) / 2), torch.round((label[1][2] + label[2][2]) / 2)} | |
| local number = utils.getNumber('labels', label[3]) | |
| local roomIndex = rooms[{center[2], center[1]}] | |
| if roomIndex > 0 and roomIndex ~= backgroundRoomIndex then | |
| local segmentIndex = roomIndexMap[roomIndex] | |
| if not segmentIndex or (number < segmentIndex and (number ~= 2 or segmentIndex ~= 3)) or (number == 3 and segmentIndex == 2) then | |
| roomIndexMap[roomIndex] = number | |
| end | |
| end | |
| end | |
| local roomJunctionCounter = {} | |
| local maxNumPoints = maxNumPoints or 40 | |
| local deltas = {{-1, -1}, {-1, 1}, {1, -1}, {1, 1}} | |
| for pointIndex, point in pairs(representation.points) do | |
| if pointIndex > maxNumPoints then | |
| break | |
| end | |
| for _, delta in pairs(deltas) do | |
| local x = torch.round(point[1][1]) | |
| local y = torch.round(point[1][2]) | |
| while x >= 1 and x <= width and y >= 1 and y <= height and wallMask[y][x] == 1 do | |
| x = x + delta[1] | |
| y = y + delta[2] | |
| end | |
| if x >= 1 and x <= width and y >= 1 and y <= height then | |
| local roomIndex = rooms[y][x] | |
| if roomJunctionCounter[roomIndex] == nil then | |
| roomJunctionCounter[roomIndex] = {} | |
| end | |
| if roomJunctionCounter[roomIndex][pointIndex] == nil then | |
| roomJunctionCounter[roomIndex][pointIndex] = 0 | |
| end | |
| roomJunctionCounter[roomIndex][pointIndex] = roomJunctionCounter[roomIndex][pointIndex] + 1 | |
| end | |
| end | |
| end | |
| local roomJunctions = torch.zeros(10, maxNumPoints) | |
| for roomIndex = 1, numRooms do | |
| local label = roomIndexMap[roomIndex] | |
| if label == nil then | |
| label = 10 | |
| end | |
| if label >= 1 and label <= 10 and roomIndex ~= backgroundRoomIndex then | |
| local junctionCounter = roomJunctionCounter[roomIndex] | |
| if junctionCounter then | |
| for pointIndex, count in pairs(junctionCounter) do | |
| roomJunctions[label][pointIndex] = 1 | |
| end | |
| end | |
| end | |
| end | |
| local junctionMasks = utils.drawJunctionMasks(width, height, representation.points, maxNumPoints) | |
| return junctionMasks, roomJunctions | |
| end | |
| function utils.parseFloorplan(floorplanInput) | |
| local binaryThreshold = 0.7 | |
| --image.save('test/floorplan.png', floorplanInput) | |
| --image.save('test/floorplan_segmentation.png', ut.drawSegmentation(floorplanSegmentation, numSegments)) | |
| --image.save('test/floorplan_binary.png', floorplanBinary) | |
| if borderSegmentReversed == 0 then | |
| return | |
| end | |
| local characterMasks = torch.cmul(1 - floorplanSegmentationReversed:eq(backgroundSegmentReversed), 1 - floorplanSegmentationReversed:eq(borderSegmentReversed)) | |
| --image.save('test/invalid_mask_1.png', 1 - floorplanSegmentationReversed:eq(foregroundSegment)) | |
| --image.save('test/invalid_mask_2.png', 1 - floorplanSegmentationReversed:eq(maxCountSegment)) | |
| --image.save('test/character_masks.png', characterMasks:double()) | |
| local borderMask = floorplanSegmentationReversed:eq(borderSegmentReversed) | |
| local borderIndices = borderMask:nonzero() | |
| local borderSegment = floorplanSegmentation[{borderIndices[1][1], borderIndices[1][2]}] | |
| --[[ | |
| local borderThickness = 0 | |
| local borderMaskEroded = borderMask:clone() | |
| while true do | |
| if ##borderMaskEroded:eq(1):nonzero() > 0 then | |
| borderMaskEroded = image.erode(borderMaskEroded) | |
| borderThickness = borderThickness + 2 | |
| else | |
| borderThickness = borderThickness + 1 | |
| break | |
| end | |
| end | |
| ]]-- | |
| --print("borderThickness: " .. borderThickness) | |
| borderMask = image.dilate(borderMask) | |
| local floorplanSegmentationErodedOnce = torch.cmul(floorplanSegmentation, (1 - borderMask):int()) | |
| --floorplanSegmentationErodedOnce = floorplanSegmentation:clone() | |
| --[[ | |
| local numDilations = math.floor(borderThickness / 2) | |
| for i = 2, numDilations do | |
| borderMask = image.dilate(borderMask) | |
| end | |
| local floorplanSegmentationEroded = torch.cmul(floorplanSegmentation, (1 - borderMask):int()) | |
| ]]-- | |
| local floorplanSegmentationRefined = floorplanSegmentation:clone() | |
| local backgroundSegment = floorplanSegmentation[1][1] | |
| floorplanSegmentationRefined[floorplanSegmentation:eq(backgroundSegment)] = 1 | |
| floorplanSegmentationRefined[floorplanSegmentation:eq(borderSegment)] = 0 | |
| local segmentInfo = { | |
| mins = {}, | |
| maxs = {}, | |
| means = {}, | |
| nums = {}, | |
| samples = {}, | |
| } | |
| for segment = 1, numSegments do | |
| if segment ~= backgroundSegment and segment ~= borderSegment then | |
| local indices = floorplanSegmentation:eq(segment):nonzero() | |
| segmentInfo.mins[segment] = torch.min(indices, 1)[1] | |
| segmentInfo.maxs[segment] = torch.max(indices, 1)[1] | |
| segmentInfo.means[segment] = torch.mean(indices:double(), 1)[1] | |
| segmentInfo.nums[segment] = (#indices)[1] | |
| segmentInfo.samples[segment] = indices[1] | |
| end | |
| end | |
| local largeSegments = {} | |
| local smallSegments = {} | |
| for segment = 1, numSegments do | |
| if segment ~= backgroundSegment and segment ~= borderSegment then | |
| --if ##floorplanSegmentationEroded:eq(segment):nonzero() == 0 then | |
| --floorplanSegmentationRefined[floorplanSegmentation:eq(segment)] = 0 | |
| if ##floorplanSegmentationErodedOnce:eq(segment):nonzero() > 0 and (#floorplanSegmentationErodedOnce:eq(segment):nonzero())[1] == segmentInfo.nums[segment] then | |
| table.insert(smallSegments, segment) | |
| elseif segmentInfo.maxs[segment][1] - segmentInfo.mins[segment][1] < borderThickness or segmentInfo.maxs[segment][2] - segmentInfo.mins[segment][2] < borderThickness then | |
| floorplanSegmentationRefined[floorplanSegmentation:eq(segment)] = 0 | |
| else | |
| table.insert(largeSegments, segment) | |
| end | |
| end | |
| end | |
| local characterSegmentsReversed = {} | |
| for segmentReversed = 1, numSegmentsReversed do | |
| if segmentReversed ~= backgroundSegmentReversed and segmentReversed ~= borderSegmentReversed then | |
| table.insert(characterSegmentsReversed, segmentReversed) | |
| end | |
| end | |
| --print(smallSegments) | |
| local characterSegmentInfo = { | |
| mins = {}, | |
| maxs = {}, | |
| means = {}, | |
| nums = {}, | |
| } | |
| for _, characterSegmentReversed in pairs(characterSegmentsReversed) do | |
| local indices = floorplanSegmentationReversed:eq(characterSegmentReversed):nonzero() | |
| characterSegmentInfo.mins[characterSegmentReversed] = torch.min(indices, 1)[1] | |
| characterSegmentInfo.maxs[characterSegmentReversed] = torch.max(indices, 1)[1] | |
| characterSegmentInfo.means[characterSegmentReversed] = torch.mean(indices:double(), 1)[1] | |
| characterSegmentInfo.means[characterSegmentReversed] = (#indices)[1] | |
| end | |
| for _, smallSegment in pairs(smallSegments) do | |
| local minDistance | |
| local minDistanceLargeSegment | |
| for _, largeSegment in pairs(largeSegments) do | |
| if segmentInfo.mins[smallSegment][1] >= segmentInfo.mins[largeSegment][1] and segmentInfo.mins[smallSegment][2] >= segmentInfo.mins[largeSegment][2] and segmentInfo.maxs[smallSegment][1] <= segmentInfo.maxs[largeSegment][1] and segmentInfo.maxs[smallSegment][2] <= segmentInfo.maxs[largeSegment][2] then | |
| local distance = torch.norm(segmentInfo.means[smallSegment] - segmentInfo.means[largeSegment]) | |
| if minDistance == nil or distance < minDistance then | |
| minDistanceLargeSegment = largeSegment | |
| minDistance = distance | |
| end | |
| end | |
| end | |
| if minDistanceLargeSegment == nil then | |
| minDistanceLargeSegment = 1 | |
| end | |
| floorplanSegmentationRefined[floorplanSegmentation:eq(smallSegment)] = minDistanceLargeSegment | |
| end | |
| for _, characterSegmentReversed in pairs(characterSegmentsReversed) do | |
| local minDistance | |
| local minDistanceLargeSegment | |
| for _, largeSegment in pairs(largeSegments) do | |
| if characterSegmentInfo.mins[characterSegmentReversed][1] >= segmentInfo.mins[largeSegment][1] and characterSegmentInfo.mins[characterSegmentReversed][2] >= segmentInfo.mins[largeSegment][2] and characterSegmentInfo.maxs[characterSegmentReversed][1] <= segmentInfo.maxs[largeSegment][1] and characterSegmentInfo.maxs[characterSegmentReversed][2] <= segmentInfo.maxs[largeSegment][2] then | |
| local distance = torch.norm(characterSegmentInfo.means[characterSegmentReversed] - segmentInfo.means[largeSegment]) | |
| if minDistance == nil or distance < minDistance then | |
| minDistanceLargeSegment = largeSegment | |
| minDistance = distance | |
| end | |
| end | |
| end | |
| if minDistanceLargeSegment == nil then | |
| minDistanceLargeSegment = 1 | |
| end | |
| floorplanSegmentationRefined[floorplanSegmentationReversed:eq(characterSegmentReversed)] = minDistanceLargeSegment | |
| characterMasks[floorplanSegmentationReversed:eq(characterSegmentReversed)] = minDistanceLargeSegment | |
| end | |
| --[[ | |
| segmentInfo.indices = {} | |
| for _, largeSegment in pairs(largeSegments) do | |
| local indices = floorplanSegmentationRefined:eq(largeSegment):nonzero() | |
| segmentInfo.indices[largeSegment] = indices | |
| segmentInfo.mins[largeSegment] = torch.min(indices, 1)[1] | |
| segmentInfo.maxs[largeSegment] = torch.max(indices, 1)[1] | |
| segmentInfo.means[largeSegment] = torch.mean(indices:double(), 1)[1] | |
| segmentInfo.nums[largeSegment] = (#indices)[1] | |
| end | |
| ]]-- | |
| local borderMask = floorplanSegmentationRefined:eq(0) | |
| borderMask[floorplanSegmentationRefined:eq(1)] = 1 | |
| borderMask = image.erode(borderMask) | |
| local borderMaskByte = (borderMask * 255):byte() | |
| local floorplanSegmentationCoarse = torch.IntTensor(borderMaskByte:size()) | |
| local numSegmentCoarse = cv.connectedComponents{255 - borderMaskByte, floorplanSegmentationCoarse} | |
| floorplanSegmentationCoarse = floorplanSegmentationCoarse + 1 | |
| --image.save('test/floorplan_segmentation_coarse.png', ut.drawSegmentation(floorplanSegmentationCoarse, numSegmentCoarse)) | |
| local coarseToFineSegmentMap = {} | |
| for segment = 1, numSegmentCoarse do | |
| coarseToFineSegmentMap[segment] = {} | |
| end | |
| for _, largeSegment in pairs(largeSegments) do | |
| table.insert(coarseToFineSegmentMap[floorplanSegmentationCoarse[segmentInfo.samples[largeSegment][1]][segmentInfo.samples[largeSegment][2]]], largeSegment) | |
| end | |
| for segmentCoarse = 1, numSegmentCoarse do | |
| local mins | |
| local maxs | |
| local maxNum | |
| local maxNumSegment | |
| local count = 0 | |
| for _, largeSegment in pairs(coarseToFineSegmentMap[segmentCoarse]) do | |
| count = count + 1 | |
| if maxNum == nil or segmentInfo.nums[largeSegment] > maxNum then | |
| maxNumSegment = largeSegment | |
| maxNum = segmentInfo.nums[largeSegment] | |
| end | |
| if mins == nil then | |
| mins = segmentInfo.mins[largeSegment] | |
| maxs = segmentInfo.maxs[largeSegment] | |
| else | |
| for i = 1, 2 do | |
| mins[i] = math.min(mins[i], segmentInfo.mins[largeSegment][i]) | |
| maxs[i] = math.max(maxs[i], segmentInfo.maxs[largeSegment][i]) | |
| end | |
| end | |
| end | |
| if count > 1 then | |
| local newSegmentWidth = maxs[1] - mins[1] + 1 | |
| local newSegmentHeight = maxs[2] - mins[2] + 1 | |
| local cornerMask = torch.zeros((#floorplanSegmentation)[1], (#floorplanSegmentation)[2]) | |
| cornerMask:narrow(1, mins[1], newSegmentWidth):narrow(2, mins[2], newSegmentHeight):fill(1) | |
| cornerMask = torch.cmul(cornerMask:byte(), 1 - floorplanSegmentationCoarse:eq(segmentCoarse)) | |
| cornerMask:narrow(1, segmentInfo.mins[maxNumSegment][1], segmentInfo.maxs[maxNumSegment][1] - segmentInfo.mins[maxNumSegment][1] + 1):narrow(2, segmentInfo.mins[maxNumSegment][2], segmentInfo.maxs[maxNumSegment][2] - segmentInfo.mins[maxNumSegment][2] + 1):fill(0) | |
| --local oriSegmentArea = (segmentInfo.maxs[maxNumSegment][1] - segmentInfo.mins[maxNumSegment][1] + 1) * (segmentInfo.maxs[maxNumSegment][2] - segmentInfo.mins[maxNumSegment][2] + 1) | |
| if ##cornerMask:nonzero() == 0 or (#cornerMask:nonzero())[1] < newSegmentWidth * newSegmentHeight * 0.1 then | |
| for _, largeSegment in pairs(coarseToFineSegmentMap[segmentCoarse]) do | |
| if largeSegment ~= maxNumSegment then | |
| floorplanSegmentationRefined[floorplanSegmentationRefined:eq(largeSegment)] = maxNumSegment | |
| characterMasks[characterMasks:eq(largeSegment)] = maxNumSegment | |
| end | |
| end | |
| end | |
| end | |
| end | |
| local newSegment = 1 | |
| for segment = 1, numSegments do | |
| if ##floorplanSegmentationRefined:eq(segment):nonzero() > 0 then | |
| floorplanSegmentationRefined[floorplanSegmentationRefined:eq(segment)] = newSegment | |
| characterMasks[characterMasks:eq(segment)] = newSegment | |
| newSegment = newSegment + 1 | |
| end | |
| end | |
| local numSegmentsRefined = newSegment - 1 | |
| return floorplanSegmentationRefined, characterMasks, numSegmentsRefined | |
| end | |
| function utils.checkSanity(floorplan, representation, lineWidth) | |
| local lineWidth = lineWidth or 5 | |
| local width = floorplan:size(3) | |
| local height = floorplan:size(2) | |
| for mode, items in pairs(representation) do | |
| if #items == 0 then | |
| print('no ' .. mode) | |
| return false | |
| end | |
| end | |
| for labelIndex_1, label_1 in pairs(representation.labels) do | |
| for labelIndex_2, label_2 in pairs(representation.labels) do | |
| if labelIndex_2 > labelIndex_1 then | |
| if math.abs((label_1[1][1] + label_1[2][1]) / 2 - (label_2[1][1] + label_2[2][1]) / 2) < 5 and math.abs((label_1[1][2] + label_1[2][2]) / 2 - (label_2[1][2] + label_2[2][2]) / 2) < 5 then | |
| print('two labels overlap: ' .. (label_1[1][1] + label_1[2][1]) / 2 .. ' ' .. (label_1[1][2] + label_1[2][2]) / 2 .. ' ' .. label_1[3][1] .. ' ' .. label_2[3][1]) | |
| return false | |
| end | |
| end | |
| end | |
| end | |
| local wallMask = utils.drawLineMask(width, height, representation.walls, lineWidth) | |
| local rooms, numRooms = utils.findConnectedComponents(1 - wallMask) | |
| local backgroundRoomIndex | |
| local imageCorners = {{1, 1}, {width, 1}, {width, height}, {1, height}} | |
| for _, imageCorner in pairs(imageCorners) do | |
| local roomIndex = rooms[imageCorner[2]][imageCorner[1]] | |
| if roomIndex > 0 then | |
| if not backgroundRoomIndex then | |
| backgroundRoomIndex = roomIndex | |
| elseif roomIndex ~= backgroundRoomIndex then | |
| rooms[rooms:eq(roomIndex)] = backgroundRoomIndex | |
| end | |
| end | |
| end | |
| if not backgroundRoomIndex then | |
| backgroundRoomIndex = numRooms | |
| end | |
| local roomLabelsMap = {} | |
| for roomIndex = 1, numRooms - 1 do | |
| roomLabelsMap[roomIndex] = {} | |
| end | |
| for _, label in pairs(representation.labels) do | |
| local center = {torch.round((label[1][1] + label[2][1]) / 2), torch.round((label[1][2] + label[2][2]) / 2)} | |
| local number = utils.getNumber('labels', label[3]) | |
| local roomIndex = rooms[{center[2], center[1]}] | |
| if roomIndex == 0 then | |
| print('label ' .. label[1][1] .. ' ' .. label[1][2] .. ' ' .. label[2][1] .. ' ' .. label[2][2] .. ' ' .. label[3][1] .. ' is on walls') | |
| return false | |
| end | |
| if roomIndex == backgroundRoomIndex then | |
| if label[3][1] ~= 'corridor' then | |
| print('label ' .. label[1][1] .. ' ' .. label[1][2] .. ' ' .. label[2][1] .. ' ' .. label[2][2] .. ' ' .. label[3][1] .. ' is on background') | |
| return false | |
| end | |
| end | |
| table.insert(roomLabelsMap[roomIndex], label[3][1]) | |
| end | |
| local sanityLabelsPairs = {{'living_room', 'kitchen'}, {'kitchen', 'bedroom'}, {'kitchen', 'corridor'}, {'living_room', 'corridor'}, {'bedroom', 'corridor'}, {'closet', 'corridor'}, {'corridor', 'corridor'}, {'corridor', 'washing_room'}, {'bedroom', 'washing_room'}, {'living_room', 'washing_room'}, {'kitchen', 'washing_room'}, {'restroom', 'washing_room'}} | |
| sanityLabelMap = {} | |
| for _, sanityLabelsPair in pairs(sanityLabelsPairs) do | |
| if not sanityLabelMap[sanityLabelsPair[1]] then | |
| sanityLabelMap[sanityLabelsPair[1]] = {} | |
| end | |
| sanityLabelMap[sanityLabelsPair[1]][sanityLabelsPair[2]] = true | |
| if not sanityLabelMap[sanityLabelsPair[2]] then | |
| sanityLabelMap[sanityLabelsPair[2]] = {} | |
| end | |
| sanityLabelMap[sanityLabelsPair[2]][sanityLabelsPair[1]] = true | |
| end | |
| for roomIndex, labels in pairs(roomLabelsMap) do | |
| if #labels > 1 then | |
| for _, label_1 in pairs(labels) do | |
| for __, label_2 in pairs(labels) do | |
| if __ > _ then | |
| if not sanityLabelMap[label_1] or not sanityLabelMap[label_1][label_2] then | |
| print(label_1 .. ' and ' .. label_2 .. ' exist in the same room') | |
| return | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| local doorWidth = lineWidth - 2 | |
| local doorMask = utils.drawLineMask(width, height, representation.doors, doorWidth, true, 0) | |
| local indices = doorMask:nonzero() | |
| if ##indices == 0 and false then | |
| print('no door pixel') | |
| return false | |
| end | |
| for i = 1, indices:size(1) do | |
| local y = indices[i][1] | |
| local x = indices[i][2] | |
| local wallIndex = wallMask[y][x] | |
| local doorIndex = doorMask[y][x] | |
| local doorDim = utils.lineDim(representation.doors[doorIndex]) | |
| if wallIndex == 0 and doorDim > 0 then | |
| --image.save('test/doors.png', utils.drawSegmentation(doorMask)) | |
| --image.save('test/door.png', doorMask:eq(2):double()) | |
| --print(doorIndex) | |
| --print(representation.doors[doorIndex][1][1]) | |
| --print(representation.doors[doorIndex][1][2]) | |
| --nprint(representation.doors[doorIndex][2][1]) | |
| --print(representation.doors[doorIndex][2][2]) | |
| print('door pixel ' .. x .. ' ' .. y .. ' is not on a wall') | |
| --.. representation.doors[doorIndex][1][doorDim] .. ' ' .. representation.doors[doorIndex][2][doorDim] .. ' ' .. doorDim) | |
| return false | |
| end | |
| end | |
| local gap = lineWidth | |
| for lineIndex_1, line_1 in pairs(representation.doors) do | |
| local lineDim_1 = utils.lineDim(line_1) | |
| if lineDim_1 > 0 then | |
| local fixedValue_1 = (line_1[1][3 - lineDim_1] + line_1[2][3 - lineDim_1]) / 2 | |
| for lineIndex_2, line_2 in pairs(representation.doors) do | |
| if lineIndex_2 > lineIndex_1 then | |
| local lineDim_2 = utils.lineDim(line_2) | |
| if lineDim_2 == lineDim_1 then | |
| local fixedValue_2 = (line_2[1][3 - lineDim_2] + line_2[2][3 - lineDim_2]) / 2 | |
| --local nearestPair, minDistance = utils.findNearestJunctionPair(line_1, line_2, gap, false) | |
| --print(minDistance .. ' ' .. lineDim_1 .. ' ' .. lineDim_2) | |
| local lineDim = lineDim_1 | |
| if math.abs(fixedValue_2 - fixedValue_1) <= gap and math.min(math.max(line_2[1][lineDim], line_2[2][lineDim]), math.max(line_1[1][lineDim], line_1[2][lineDim])) > math.max(math.min(line_2[1][lineDim], line_2[2][lineDim]), math.min(line_1[1][lineDim], line_1[2][lineDim])) then | |
| --local pointIndex_1 = nearestPair[1] | |
| --local pointIndex_2 = nearestPair[2] | |
| --print(lineIndex_1 .. ' ' .. lineIndex_2) | |
| print('two door overlap ' .. line_1[1][1] .. ' ' .. line_1[1][2]) | |
| --.. ' ' .. line_1[3 - pointIndex_1][lineDim_1] .. ' ' .. line_2[3 - pointIndex_2][lineDim_2] .. ' ' .. lineDim_1) | |
| --print('line dim ' .. lineDim_1) | |
| return false | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| local gap = lineWidth | |
| for lineIndex_1, line_1 in pairs(representation.walls) do | |
| local lineDim_1 = utils.lineDim(line_1) | |
| if lineDim_1 > 0 then | |
| local fixedValue_1 = (line_1[1][3 - lineDim_1] + line_1[2][3 - lineDim_1]) / 2 | |
| for lineIndex_2, line_2 in pairs(representation.walls) do | |
| if lineIndex_2 > lineIndex_1 and line_2[3][2] == line_1[3][2] then | |
| local lineDim_2 = utils.lineDim(line_2) | |
| if lineDim_2 == lineDim_1 then | |
| local fixedValue_2 = (line_2[1][3 - lineDim_2] + line_2[2][3 - lineDim_2]) / 2 | |
| local nearestPair, minDistance = utils.findNearestJunctionPair(line_1, line_2, gap, false) | |
| --print(minDistance .. ' ' .. lineDim_1 .. ' ' .. lineDim_2) | |
| if minDistance <= gap then | |
| local pointIndex_1 = nearestPair[1] | |
| local pointIndex_2 = nearestPair[2] | |
| --print(lineIndex_1 .. ' ' .. lineIndex_2) | |
| if pointIndex_1 > 0 and pointIndex_2 > 0 then | |
| print('two walls should be merged at pixel ' .. line_1[pointIndex_1][1] .. ' ' .. line_1[pointIndex_1][2]) | |
| --.. ' ' .. line_1[3 - pointIndex_1][lineDim_1] .. ' ' .. line_2[3 - pointIndex_2][lineDim_2] .. ' ' .. lineDim_1) | |
| --print('line dim ' .. lineDim_1) | |
| return false | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| return true | |
| end | |
| function utils.drawDoorMasks(width, height, doors, doorWidth) | |
| local labelDoorsMap = {} | |
| for _, door in pairs(doors) do | |
| local label = utils.getNumber('doors', door[3]) | |
| if not labelDoorsMap[label] then | |
| labelDoorsMap[label] = {} | |
| end | |
| table.insert(labelDoorsMap[label], door) | |
| end | |
| local doorMasks = torch.zeros(13, width, height) | |
| for label, labelDoors in pairs(labelDoorsMap) do | |
| doorMasks[label] = utils.drawLineMask(width, height, labelDoors, doorWidth) | |
| end | |
| return doorMasks | |
| end | |
| function utils.filterJunctions(width, height, points, wallMask) | |
| local pointOrientations = {{{3}, {4}, {1}, {2}}, {{4, 1}, {1, 2}, {2, 3}, {3, 4}}, {{2, 3, 4}, {3, 4, 1}, {4, 1, 2}, {1, 2, 3}}, {{1, 2, 3, 4}}} | |
| local orientationRanges = torch.Tensor({{width, 0, 0, 0}, {width, height, width, 0}, {width, height, 0, height}, {0, height, 0, 0}}) | |
| local candidateLines = {} | |
| local pointOrientationLinesMap = {} | |
| local pointNeighbors = {} | |
| for pointIndex, point in pairs(points) do | |
| table.insert(pointOrientationLinesMap, {}) | |
| local pointType = point[3][2] | |
| local orientations = pointOrientations[pointType][point[3][3]] | |
| for _, orientation in pairs(orientations) do | |
| pointOrientationLinesMap[pointIndex][orientation] = {} | |
| end | |
| table.insert(pointNeighbors, {}) | |
| end | |
| local gap = 10 | |
| local gapWeight = 1 | |
| local wallEvidenceWeight = 10 | |
| local shortLineWeight = 10 | |
| --local pointWeight = 1000 | |
| local pointMissingLineWeight = 100 | |
| local closePointWeight = 1000 | |
| local overlappingLineWeight = 1000 | |
| local hugeWeight = 10000 | |
| for pointIndex, point in pairs(points) do | |
| local pointType = point[3][2] | |
| local orientations = pointOrientations[pointType][point[3][3]] | |
| for _, orientation in pairs(orientations) do | |
| local oppositeOrientation = (orientation + 2 - 1) % 4 + 1 | |
| local ranges = orientationRanges[orientation]:clone() | |
| local lineDim = 0 | |
| if orientation == 1 or orientation == 3 then | |
| lineDim = 2 | |
| else | |
| lineDim = 1 | |
| end | |
| deltas = {0, 0} | |
| deltas[3 - lineDim] = gap | |
| for c = 1, 2 do | |
| ranges[c] = math.min(ranges[c], point[1][c] - deltas[c]) | |
| ranges[c + 2] = math.max(ranges[c + 2], point[1][c] + deltas[c]) | |
| end | |
| for neighborPointIndex, neighborPoint in pairs(points) do | |
| if neighborPointIndex > pointIndex then | |
| for __, neighborOrientation in pairs(pointOrientations[neighborPoint[3][2]][neighborPoint[3][3]]) do | |
| if neighborOrientation == oppositeOrientation then | |
| local inRange = true | |
| for c = 1, 2 do | |
| if neighborPoint[1][c] < ranges[c] or neighborPoint[1][c] > ranges[c + 2] then | |
| inRange = false | |
| end | |
| end | |
| if inRange and math.abs(neighborPoint[1][lineDim] - point[1][lineDim]) > math.max(math.abs(neighborPoint[1][3 - lineDim] - point[1][3 - lineDim]), gap) then | |
| --print(pointIndex .. ' ' .. neighborPointIndex .. ' ' .. lineDim .. ' ' .. orientation) | |
| --print(ranges[1] .. ' ' .. ranges[2] .. ' ' .. ranges[3] .. ' ' .. ranges[4]) | |
| --print(point[1][lineDim] .. ' ' .. neighborPoint[1][lineDim] .. ' ' .. point[1][3 - lineDim] .. ' ' .. neighborPoint[1][3 - lineDim]) | |
| local lineIndex = #candidateLines + 1 | |
| table.insert(pointOrientationLinesMap[pointIndex][orientation], lineIndex) | |
| table.insert(pointOrientationLinesMap[neighborPointIndex][oppositeOrientation], lineIndex) | |
| pointNeighbors[pointIndex][neighborPointIndex] = true | |
| pointNeighbors[neighborPointIndex][pointIndex] = true | |
| local cost = 0 | |
| cost = cost + math.abs(neighborPoint[1][3 - lineDim] - point[1][3 - lineDim]) / gap * gapWeight | |
| --local fixedValue = torch.round((neighborPoint[1][3 - lineDim] + point[1][3 - lineDim]) / 2) | |
| --local numNonWallPoints = 0 | |
| --for delta in xrange(int(abs(neighborPoint[lineDim] - point[lineDim]) + 1)): | |
| --intermediatePoint = [0, 0] | |
| --intermediatePoint[lineDim] = int(min(neighborPoint[lineDim], point[lineDim]) + delta) | |
| --intermediatePoint[3 - lineDim] = fixedValue | |
| --if wall_evidence[intermediatePoint[1]][intermediatePoint[0]] < 0.5: | |
| --numNonWallPoints += 1 | |
| --cost += numNonWallPoints / abs(neighborPoint[lineDim] - point[lineDim]) * wallEvidenceWeight | |
| table.insert(candidateLines, {pointIndex, neighborPointIndex, cost}) | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| end | |
| local conflictLinePairs = {} | |
| local D_p = torch.zeros(#points) | |
| for pointIndex = 1, #points do | |
| D_p[pointIndex] = #pointOrientationLinesMap[pointIndex] | |
| end | |
| D_p = torch.diag(D_p) | |
| local D_l = torch.zeros(#candidateLines) | |
| for lineIndex, line in pairs(candidateLines) do | |
| D_l[lineIndex] = line[3] | |
| end | |
| D_l = torch.diag(D_l) | |
| local PLs = torch.zeros(4, #points, #candidateLines) | |
| for pointIndex, orientationLinesMap in pairs(pointOrientationLinesMap) do | |
| for orientation, lines in pairs(orientationLinesMap) do | |
| for _, line in pairs(lines) do | |
| PLs[orientation][pointIndex][line] = 1 | |
| end | |
| end | |
| end | |
| local PP = torch.zeros(#points, #points) | |
| for pointIndex, point in pairs(points) do | |
| for neighborPointIndex, neighborPoint in pairs(points) do | |
| if neighborPointIndex ~= pointIndex then | |
| local distance = math.pow(math.pow(point[1][1] - neighborPoint[1][1], 2) + math.pow(point[1][2] - neighborPoint[1][2], 2), 0.5) | |
| if distance < gap and not pointNeighbors[pointIndex][neighborPointIndex] then | |
| PP[pointIndex][neighborPointIndex] = 1 | |
| end | |
| end | |
| end | |
| end | |
| --for conflictLinePair in conflictLinePairs: | |
| --model.addConstr(l[conflictLinePair[0]] + l[conflictLinePair[1]] <= 1) | |
| local PL2_sum = torch.zeros(#candidateLines, #candidateLines) | |
| local PL_sum = torch.zeros(#points, #candidateLines) | |
| for orientation = 1, 4 do | |
| --PLs[orientation]:fill(0) | |
| PL2_sum = PL2_sum + PLs[orientation]:transpose(1, 2) * PLs[orientation] | |
| PL_sum = PL_sum + PLs[orientation] | |
| end | |
| --[[ | |
| for i = 1, #points do | |
| for j = 1, #points do | |
| if i ~= j then | |
| local same = true | |
| for k = 1, #candidateLines do | |
| if PL_sum[i][k] ~= PL_sum[j][k] then | |
| same = false | |
| break | |
| end | |
| end | |
| if same then | |
| print(i) | |
| print(j) | |
| os.exit(1) | |
| end | |
| end | |
| end | |
| end | |
| ]]-- | |
| --print(PL2_sum:sum(1)) | |
| --print(PL2_sum:sum(2)) | |
| --print(torch.min(torch.diag(D_l))) | |
| local D_p_weight = 50 | |
| local D_l_weight = 1 | |
| local PL_weight = 100 | |
| local PP_weight = 1000 | |
| local L_coef = PL_weight * PL2_sum + D_l_weight * D_l | |
| --PL2_sum = PL2_sum + torch.randn(#PL2_sum) * 0.0001 | |
| local L_coef_inv = torch.inverse(L_coef) | |
| local L_p = L_coef_inv * (PL_weight * PL_sum:transpose(1, 2)) | |
| --local L_c = -0.5 * L_coef_inv * D_l | |
| local P_l = PL_weight * PL_sum | |
| local P_coef = 4 * PL_weight * torch.diag(torch.ones(#points)) + PP_weight * PP:transpose(1, 2) + D_p_weight * D_p - P_l * L_p | |
| local P_coef_inv = torch.inverse(P_coef) | |
| local P_x = P_coef_inv * D_p_weight * D_p | |
| --local P_c = P_coef_inv * (P_l * L_c) | |
| local p = P_x * torch.ones(#points) | |
| local function f(x, p) | |
| local l = L_p * p | |
| print(l) | |
| local f = D_p_weight * D_p * (x - p) * (x - p) + D_l_weight * D_l * l * l | |
| for orientation = 1, 4 do | |
| f = f + PL_weight * (p - PLs[orientation] * l) * (p - PLs[orientation] * l) | |
| end | |
| f = f + PP_weight * PP * p * p | |
| return f | |
| end | |
| print(torch.cat(torch.range(1, p:size(1)), p, 2)) | |
| print(f(torch.ones(#points), p)) | |
| print(f(torch.ones(#points), torch.ones(#points))) | |
| --print(f(torch.ones(#points), torch.zeros(#points))) | |
| os.exit(1) | |
| end | |
| function utils.attachDoorsOnWalls(width, height, representation, lineWidth, doorWidth, maxNumWalls) | |
| local lineWidth = lineWidth or 5 | |
| local maxNumWalls = maxNumWalls or 100 | |
| --utils.printRepresentation(representation) | |
| --os.exit(1) | |
| --print(representation.points) | |
| local walls, wallJunctionsMap = utils.pointsToLines(width, height, representation.points, lineWidth, true) | |
| local wallMask = utils.drawLineMask(width, height, walls, lineWidth, true) | |
| local doorWidth = doorWidth or lineWidth - 2 | |
| local doorMask = utils.drawLineMask(width, height, representation.doors, doorWidth, true, 0) | |
| local indices = doorMask:nonzero() | |
| local wallDoorPointsMap = {} | |
| local doorCounters = {} | |
| if ##indices > 0 then | |
| for i = 1, indices:size(1) do | |
| local y = indices[i][1] | |
| local x = indices[i][2] | |
| local doorIndex = doorMask[y][x] | |
| local wallIndex = wallMask[y][x] | |
| if wallIndex > 0 then | |
| if not wallDoorPointsMap[wallIndex] then | |
| wallDoorPointsMap[wallIndex] = {} | |
| end | |
| if not wallDoorPointsMap[wallIndex][doorIndex] then | |
| wallDoorPointsMap[wallIndex][doorIndex] = {} | |
| end | |
| table.insert(wallDoorPointsMap[wallIndex][doorIndex], {x, y}) | |
| end | |
| if not doorCounters[doorIndex] then | |
| doorCounters[doorIndex] = 0 | |
| end | |
| doorCounters[doorIndex] = doorCounters[doorIndex] + 1 | |
| end | |
| end | |
| local wallDoorTensor = torch.zeros(maxNumWalls, 3) | |
| for wallIndex, doorPoints in pairs(wallDoorPointsMap) do | |
| if wallIndex > maxNumWalls then | |
| break | |
| end | |
| local maxCount | |
| local maxCountDoor | |
| local lineDim = utils.lineDim(walls[wallIndex]) | |
| for doorIndex, points in pairs(doorPoints) do | |
| local doorDim = utils.lineDim(representation.doors[doorIndex]) | |
| if doorDim == lineDim then | |
| local count = #points | |
| if not maxCount or count > maxCount then | |
| maxCountDoor = doorIndex | |
| maxCount = count | |
| end | |
| end | |
| end | |
| if maxCountDoor and maxCount > doorCounters[maxCountDoor] * 0.2 then | |
| local points = doorPoints[maxCountDoor] | |
| if lineDim == 0 then | |
| if math.abs(line[1][1] - line[2][1]) > math.abs(line[1][2] - line[2][2]) then | |
| lineDim = 1 | |
| else | |
| lineDim = 2 | |
| end | |
| end | |
| local minDoor | |
| local maxDoor | |
| for _, point in pairs(points) do | |
| if not minDoor or point[lineDim] < minDoor then | |
| minDoor = point[lineDim] | |
| end | |
| if not maxDoor or point[lineDim] > maxDoor then | |
| maxDoor = point[lineDim] | |
| end | |
| end | |
| local wallIndices = wallMask:eq(wallIndex):nonzero() | |
| local minWall = torch.min(wallIndices, 1)[1][3 - lineDim] | |
| local maxWall = torch.max(wallIndices, 1)[1][3 - lineDim] | |
| local minRatio = math.max((minDoor - minWall) / (maxWall - minWall + 1), 0) | |
| local maxRatio = math.min((maxDoor - minWall) / (maxWall - minWall + 1), 1) | |
| wallDoorTensor[wallIndex][1] = minRatio | |
| wallDoorTensor[wallIndex][2] = maxRatio | |
| wallDoorTensor[wallIndex][3] = utils.getNumber('doors', representation.doors[maxCountDoor][3]) | |
| end | |
| end | |
| local wallHeatmaps = torch.zeros(maxNumWalls, height, width) | |
| for wallIndex, _ in pairs(walls) do | |
| if wallIndex > maxNumWalls then | |
| break | |
| end | |
| wallHeatmaps[wallIndex] = wallMask:eq(wallIndex):double() | |
| end | |
| return wallHeatmaps, wallDoorTensor | |
| end | |
| function utils.drawMasks( img, masks, maxn, alpha, clrs ) | |
| assert(img:isContiguous() and img:dim()==3) | |
| local n, h, w = masks:size(1), masks:size(2), masks:size(3) | |
| if not maxn then maxn=n end | |
| if not alpha then alpha=.4 end | |
| if not clrs then clrs=torch.rand(n,3)*.6+.4 end | |
| for i=1,math.min(maxn,n) do | |
| local M = masks[i]:contiguous():data() | |
| local B = torch.ByteTensor(h,w):zero():contiguous():data() | |
| -- get boundaries B in masks M quickly | |
| for y=0,h-2 do for x=0,w-2 do | |
| local k=y*w+x | |
| if M[k]~=M[k+1] then B[k],B[k+1]=1,1 end | |
| if M[k]~=M[k+w] then B[k],B[k+w]=1,1 end | |
| if M[k]~=M[k+1+w] then B[k],B[k+1+w]=1,1 end | |
| end end | |
| -- softly embed masks into image and add solid boundaries | |
| for j=1,3 do | |
| local O,c,a = img[j]:data(), clrs[i][j], alpha | |
| for k=0,w*h-1 do if M[k]==1 then O[k]=O[k]*(1-a)+c*a end end | |
| for k=0,w*h-1 do if B[k]==1 then O[k]=c end end | |
| end | |
| end | |
| end | |
| function utils.detectIcons(floorplan) | |
| local width, height = floorplan:size(3), floorplan:size(2) | |
| local inputSize = 160 | |
| local bw = inputSize / 2 | |
| if not utils.modelIcon then | |
| require 'nnx' | |
| local inn = require 'inn' | |
| if not nn.SpatialConstDiagonal then | |
| torch.class('nn.SpatialConstDiagonal', 'inn.ConstAffine') | |
| end | |
| --paths.dofile('/home/chenliu/Projects/Floorplan/floorplan/InverseCAD/models/DeepMask.lua') | |
| utils.modelIcon = torch.load('/home/chenliu/Projects/Floorplan/models/icon-deepmask/model_best.t7') | |
| end | |
| -------------------------------------------------------------------------------- | |
| -- function: linear2convTrunk | |
| local function linear2convTrunk(net,fSz) | |
| return net:replace(function(x) | |
| if torch.typename(x):find('Linear') then | |
| local nInp,nOut = x.weight:size(2)/(fSz*fSz),x.weight:size(1) | |
| local w = torch.reshape(x.weight,nOut,nInp,fSz,fSz) | |
| local y = cudnn.SpatialConvolution(nInp,nOut,fSz,fSz,1,1) | |
| y.weight:copy(w); y.gradWeight:copy(w); y.bias:copy(x.bias) | |
| return y | |
| elseif torch.typename(x):find('Threshold') then | |
| return cudnn.ReLU() | |
| elseif torch.typename(x):find('View') or | |
| torch.typename(x):find('SpatialZeroPadding') then | |
| return nn.Identity() | |
| else | |
| return x | |
| end | |
| end | |
| ) | |
| end | |
| -------------------------------------------------------------------------------- | |
| -- function: linear2convHeads | |
| local function linear2convHead(net) | |
| return net:replace(function(x) | |
| if torch.typename(x):find('Linear') then | |
| local nInp,nOut = x.weight:size(2),x.weight:size(1) | |
| local w = torch.reshape(x.weight,nOut,nInp,1,1) | |
| local y = cudnn.SpatialConvolution(nInp,nOut,1,1,1,1) | |
| y.weight:copy(w); y.gradWeight:copy(w); y.bias:copy(x.bias) | |
| return y | |
| elseif torch.typename(x):find('Threshold') then | |
| return cudnn.ReLU() | |
| elseif not torch.typename(x):find('View') and | |
| not torch.typename(x):find('Copy') then | |
| return x | |
| end | |
| end | |
| ) | |
| end | |
| utils.modelIcon:evaluate() | |
| local trunk = utils.modelIcon.modules[1] | |
| local maskBranch = utils.modelIcon.modules[2].modules[1] | |
| local scoreBranch = utils.modelIcon.modules[3] | |
| linear2convTrunk(trunk, inputSize / 16) | |
| linear2convHead(maskBranch) | |
| linear2convHead(scoreBranch) | |
| trunk:cuda() | |
| maskBranch:cuda() | |
| scoreBranch:cuda() | |
| local scales = {} | |
| for i = -1.5,0.5,1.5 do | |
| table.insert(scales, 2^i) | |
| end | |
| local numProposals = 100 | |
| local meanstd = {mean = { 0.5, 0.5, 0.5 }, std = { 0.5, 0.5, 0.5 }} | |
| local inpPad = torch.CudaTensor() | |
| local input = floorplan | |
| if input:type() == 'torch.CudaTensor' then input = input:float() end | |
| local pyramid = nn.ConcatTable() | |
| for i = 1,#scales do | |
| pyramid:add(nn.SpatialReSamplingEx{rwidth=scales[i], | |
| rheight=scales[i], mode='bilinear'}) | |
| end | |
| local inpPyramid = pyramid:forward(input) | |
| -- forward all scales through network | |
| local outPyramidMask,outPyramidScore = {},{} | |
| for i,_ in pairs(inpPyramid) do | |
| local inp = inpPyramid[i]:cuda() | |
| local h,w = inp:size(2),inp:size(3) | |
| -- padding/normalize | |
| inpPad:resize(1,3,h+2*bw,w+2*bw):fill(.5) | |
| inpPad:narrow(1,1,1):narrow(3,bw+1,h):narrow(4,bw+1,w):copy(inp) | |
| for i=1,3 do | |
| inpPad[1][i]:add(-meanstd.mean[i]):div(meanstd.std[i]) | |
| end | |
| -- forward trunk | |
| local outTrunk = trunk:forward(inpPad):squeeze() | |
| -- forward score branch | |
| local outScore = scoreBranch:forward(outTrunk) | |
| table.insert(outPyramidScore,outScore:clone():squeeze()) | |
| -- forward mask branch | |
| local outMask = maskBranch:forward(outTrunk) | |
| table.insert(outPyramidMask,outMask:float():squeeze()) | |
| outScore = outScore:clone():squeeze() | |
| outMask = outMask:float():squeeze() | |
| --[[ | |
| print(#inpPad) | |
| print(#outTrunk) | |
| print(#outMask) | |
| print(#outScore) | |
| print(#floorplan) | |
| print(scales[i]) | |
| os.exit(1) | |
| ]]-- | |
| if false then | |
| for y = 1, outMask:size(3) - 1 do | |
| for x = 1, outMask:size(2) - 1 do | |
| local maskDim = math.sqrt(outMask:size(1)) | |
| local mask = outMask[{{}, x, y}]:contiguous():view(maskDim, maskDim) | |
| local thr = 0.2 | |
| local s = scales[i] | |
| --local sz = math.floor(inputSize / s) | |
| local sz = inputSize | |
| local t = 16 | |
| local delta = inputSize / 2 | |
| mask = image.scale(mask,sz,sz,'bilinear') | |
| mask = mask:gt(thr) | |
| local imgMask = torch.zeros(h + bw * 2, w + bw * 2) | |
| imgMask:zero() | |
| print(#imgMask) | |
| print(math.floor((x - 1) * t) .. ' ' .. math.floor((y - 1) * t) .. ' ' .. sz) | |
| imgMask:narrow(1, math.floor((x - 1) * t + 1), sz):narrow(2, math.floor((y - 1) * t + 1), sz)[mask] = 1 | |
| --imgMask:narrow(1, math.floor((x - 1) * t + 1), sz):narrow(2, math.floor((y - 1) * t + 1), sz)[1 - mask] = -1 | |
| imgMask:narrow(1, math.floor((x - 1) * t + 1), 1):narrow(2, math.floor((y - 1) * t + 1), sz):fill(-1) | |
| imgMask:narrow(1, math.floor((x - 1) * t + 1) + sz - 1, 1):narrow(2, math.floor((y - 1) * t + 1), sz):fill(-1) | |
| imgMask:narrow(1, math.floor((x - 1) * t + 1), sz):narrow(2, math.floor((y - 1) * t + 1), 1):fill(-1) | |
| imgMask:narrow(1, math.floor((x - 1) * t + 1), sz):narrow(2, math.floor((y - 1) * t + 1) + sz - 1, 1):fill(-1) | |
| --imgMask = imgMask:narrow(1, bw + 1, h):narrow(2, bw + 1, w):byte() | |
| local img = torch.ones(3,h+2*bw,w+2*bw) * 0.5 | |
| img:narrow(2,bw+1,h):narrow(3,bw+1,w):copy(inp) | |
| --print(#img) | |
| --print(#imgMask) | |
| img[1][imgMask:eq(1)] = 1 | |
| img[2][imgMask:eq(1)] = 0 | |
| img[3][imgMask:eq(1)] = 0 | |
| img[1][imgMask:eq(-1)] = 0 | |
| img[2][imgMask:eq(-1)] = 0 | |
| img[3][imgMask:eq(-1)] = 1 | |
| local score = outScore[x][y] | |
| image.save('test/masks/mask_' .. i .. '_' .. x .. '_' .. y .. '_' .. score .. '.png', img) | |
| end | |
| end | |
| end | |
| end | |
| local function unfoldMasksMatrix(masks) | |
| local umasks = {} | |
| local oSz = math.sqrt(masks[1]:size(1)) | |
| for _,mask in pairs(masks) do | |
| local umask = mask:reshape(oSz,oSz,mask:size(2),mask:size(3)) | |
| umask=umask:transpose(1,3):transpose(2,3):transpose(2,4):transpose(3,4) | |
| table.insert(umasks,umask) | |
| end | |
| return umasks | |
| end | |
| local masks = unfoldMasksMatrix(outPyramidMask) | |
| local scores = outPyramidScore | |
| local function getTopScores() | |
| local sortedScores = torch.Tensor() | |
| local sortedIds = torch.Tensor() | |
| local pos = torch.Tensor() | |
| local topScores = torch.Tensor() | |
| -- sort scores/ids for each scale | |
| local nScales=#scales | |
| local rowN = scores[nScales]:size(1) * scores[nScales]:size(2) | |
| sortedScores:resize(rowN,nScales):zero() | |
| sortedIds:resize(rowN,nScales):zero() | |
| for s = 1,nScales do | |
| scores[s]:mul(-1):exp():add(1):pow(-1) -- scores2prob | |
| local sc = scores[s] | |
| local h,w = sc:size(1),sc:size(2) | |
| local sc=sc:view(h*w) | |
| local sS,sIds=torch.sort(sc,true) | |
| local sz = sS:size(1) | |
| sortedScores:narrow(2,s,1):narrow(1,1,sz):copy(sS) | |
| sortedIds:narrow(2,s,1):narrow(1,1,sz):copy(sIds) | |
| end | |
| -- get top scores | |
| local np = numProposals | |
| pos:resize(nScales):fill(1) | |
| topScores:resize(np,4):fill(1) | |
| np=math.min(np,rowN) | |
| for i = 1,np do | |
| local scale,score = 0,0 | |
| for k = 1,nScales do | |
| if sortedScores[pos[k]][k] > score then | |
| score = sortedScores[pos[k]][k] | |
| scale = k | |
| end | |
| end | |
| local temp=sortedIds[pos[scale]][scale] | |
| local x=math.floor(temp/scores[scale]:size(2)) | |
| local y=temp%scores[scale]:size(2)+1 | |
| x,y=math.max(1,x),math.max(1,y) | |
| pos[scale]=pos[scale]+1 | |
| topScores:narrow(1,i,1):copy(torch.Tensor({score,scale,x,y})) | |
| end | |
| return topScores | |
| end | |
| local topScores = getTopScores() | |
| local function getTopMasks(thr, h, w) | |
| local imgMask = torch.ByteTensor() | |
| local topMasks = torch.ByteTensor() | |
| thr = math.log(thr/(1-thr)) -- 1/(1+e^-s) > th => s > log(1-th) | |
| local masks,topScores,np = masks, topScores, numProposals | |
| topMasks:resize(np,h,w):zero() | |
| --imgMask:resize(h,w) | |
| --local imgMaskPtr = imgMask:data() | |
| for i = 1,np do | |
| local scale,x,y=topScores[i][2], topScores[i][3], topScores[i][4] | |
| local s = scales[scale] | |
| local sz = math.floor(inputSize / s) | |
| local mask = masks[scale] | |
| x,y = math.min(x,mask:size(1)),math.min(y,mask:size(2)) | |
| mask = mask[x][y]:float() | |
| local mask = image.scale(mask,sz,sz,'bilinear') | |
| mask = mask:gt(thr) | |
| local t = 16 / s | |
| local delta = inputSize / 2 / s | |
| imgMask:resize(h + delta * 2, w + delta * 2) | |
| imgMask:zero() | |
| --print(#imgMask) | |
| --print(math.floor((x - 1) * t) .. ' ' .. math.floor((y - 1) * t) .. ' ' .. sz) | |
| local sz1 = math.min(imgMask:size(1) - math.floor((x - 1) * t), sz) | |
| local sz2 = math.min(imgMask:size(2) - math.floor((y - 1) * t), sz) | |
| imgMask:narrow(1, math.floor((x - 1) * t + 1), sz1):narrow(2, math.floor((y - 1) * t + 1), sz2):copy(mask:narrow(1, 1, sz1):narrow(2, 1, sz2)) | |
| imgMask = imgMask:narrow(1, delta + 1, h):narrow(2, delta + 1, w) | |
| --[[ | |
| local mask_ptr = mask:data() | |
| local t = 16/s | |
| local delta = inputSize/2/s | |
| for im =0, sz-1 do | |
| local ii = math.floor((x-1)*t-delta+im) | |
| for jm = 0,sz- 1 do | |
| local jj=math.floor((y-1)*t-delta+jm) | |
| if mask_ptr[sz*im + jm] > thr and | |
| ii >= 0 and ii <= h-1 and jj >= 0 and jj <= w-1 then | |
| imgMaskPtr[jj+ w*ii]=1 | |
| end | |
| end | |
| end | |
| --]] | |
| topMasks:narrow(1,i,1):copy(imgMask) | |
| end | |
| return topMasks | |
| end | |
| local topMasks = getTopMasks(0.2, height, width) | |
| local IOUThreshold = 0.2 | |
| while true do | |
| local hasChange = false | |
| local invalidMasks = {} | |
| for i = 1, topMasks:size(1) do | |
| for j = i + 1, topMasks:size(1) do | |
| local intersection = torch.cmul(topMasks[i], topMasks[j]):nonzero() | |
| if ##intersection > 0 then | |
| local union = (topMasks[i] + topMasks[j]):nonzero() | |
| if intersection:size(1) / union:size(1) > IOUThreshold then | |
| invalidMasks[j] = true | |
| hasChange = true | |
| end | |
| end | |
| end | |
| end | |
| if not hasChange then | |
| break | |
| end | |
| local newTopMasks | |
| for i = 1, topMasks:size(1) do | |
| if not invalidMasks[i] then | |
| if not newTopMasks then | |
| newTopMasks = topMasks[{{i}}] | |
| else | |
| newTopMasks = torch.cat(newTopMasks, topMasks[{{i}}], 1) | |
| end | |
| end | |
| end | |
| topMasks = newTopMasks | |
| end | |
| local res = floorplan:clone() | |
| utils.drawMasks(res, topMasks) | |
| image.save('test/icon_result.png', res) | |
| return | |
| end | |
| function utils.detectItems(floorplan, mode) | |
| local mode = mode or 'icons' | |
| --detectionImage = image.drawRect(detectionImage, x_1, y_1, x_2, y_2, {color = colorMap[number], lineWidth = lineWidth}) | |
| local width, height = floorplan:size(3), floorplan:size(2) | |
| local patchDim = 16 * 7 | |
| local nClasses = 13 | |
| local modelName = 'model' .. mode | |
| if not utils[modelName] then | |
| require 'nnx' | |
| local inn = require 'inn' | |
| if not nn.SpatialConstDiagonal then | |
| torch.class('nn.SpatialConstDiagonal', 'inn.ConstAffine') | |
| end | |
| utils[modelName] = torch.load('/home/chenliu/Projects/Floorplan/models/deepmask-grid-' .. mode .. '/model_best.t7') | |
| end | |
| local model = utils[modelName] | |
| --utils.modelIcon:evaluate() | |
| local trunk = model.modules[1] | |
| local coordinatesBranch = model.modules[2] | |
| local labelBranch = model.modules[3] | |
| local scales = {} | |
| --for i = -1.5,0.5,1.5 do | |
| --table.insert(scales, 2^i) | |
| --end | |
| table.insert(scales, 1) | |
| local meanstd = {mean = { 0.5, 0.5, 0.5 }, std = { 0.5, 0.5, 0.5 }} | |
| local input = floorplan:clone() | |
| if input:type() == 'torch.CudaTensor' then input = input:float() end | |
| local pyramid = nn.ConcatTable() | |
| for i = 1,#scales do | |
| pyramid:add(nn.SpatialReSamplingEx{rwidth=scales[i], | |
| rheight=scales[i], mode='bilinear'}) | |
| end | |
| local inpPyramid = pyramid:forward(input) | |
| local inpPad = torch.CudaTensor() | |
| -- forward all scales through network | |
| --local outPyramidMask,outPyramidScore = {},{} | |
| local boxes = {} | |
| local softmax = nn.SoftMax():cuda() | |
| for i,_ in pairs(inpPyramid) do | |
| local inp = inpPyramid[i]:cuda() | |
| local h,w = inp:size(2),inp:size(3) | |
| -- padding/normalize | |
| --inpPad:resize(1,3,h+2*bw,w+2*bw):fill(.5) | |
| --inpPad:narrow(1,1,1):narrow(3,bw+1,h):narrow(4,bw+1,w):copy(inp) | |
| --for i=1,3 do | |
| --inpPad[1][i]:add(-meanstd.mean[i]):div(meanstd.std[i]) | |
| --end | |
| inpPad = inp:repeatTensor(1, 1, 1, 1) | |
| inpPad[1][i]:add(-meanstd.mean[i]):div(meanstd.std[i]) | |
| -- forward trunk | |
| local outTrunk = trunk:forward(inpPad) | |
| -- forward score branch | |
| local outLabel = labelBranch:forward(outTrunk) | |
| --table.insert(outPyramidScore,outScore:clone():squeeze()) | |
| -- forward mask branch | |
| local outCoordinates = coordinatesBranch:forward(outTrunk) | |
| --table.insert(outPyramidMask,outMask:float():squeeze()) | |
| --outScore = outScore:clone():squeeze() | |
| --outMask = outMask:float():squeeze() | |
| outCoordinates = outCoordinates:squeeze() | |
| local gridHeight, gridWidth = outCoordinates:size(2), outCoordinates:size(3) | |
| outLabel = outLabel:view(gridHeight, gridWidth, -1):transpose(2, 3):transpose(1, 2) | |
| local cellWidth = w / gridWidth | |
| local cellHeight = h / gridHeight | |
| for gridX = 1, gridWidth do | |
| for gridY = 1, gridHeight do | |
| local labelProbs = outLabel[{{}, gridY, gridX}] | |
| labelProbs = softmax:forward(labelProbs) | |
| local prob, pred = torch.max(labelProbs, 1) | |
| if pred[1] <= nClasses then | |
| local box = torch.cat(outCoordinates[{{}, gridY, gridX}], labelProbs, 1):double() | |
| local x_1 = cellWidth * (gridX - 1 + 0.5) + (box[1] - 0.5) * patchDim - math.exp(box[3]) * patchDim / 2 | |
| local y_1 = cellHeight * (gridY - 1 + 0.5) + (box[2] - 0.5) * patchDim - math.exp(box[4]) * patchDim / 2 | |
| local x_2 = cellWidth * (gridX - 1 + 0.5) + (box[1] - 0.5) * patchDim + math.exp(box[3]) * patchDim / 2 | |
| local y_2 = cellHeight * (gridY - 1 + 0.5) + (box[2] - 0.5) * patchDim + math.exp(box[4]) * patchDim / 2 | |
| box[1] = x_1 / w * width | |
| box[2] = y_1 / h * height | |
| box[3] = x_2 / w * width | |
| box[4] = y_2 / h * height | |
| table.insert(boxes, box) | |
| end | |
| end | |
| end | |
| end | |
| --os.exit(1) | |
| local items = {} | |
| local numItemsPerClass = 100 | |
| local confidenceThreshold = 0.3 | |
| if true then | |
| local boxesTensor = torch.zeros(#boxes, 5) | |
| for boxIndex, box in pairs(boxes) do | |
| boxesTensor[boxIndex]:narrow(1, 1, 4):copy(box[{{1, 4}}]) | |
| boxesTensor[boxIndex][5] = torch.max(box[{{5, -1}}]) | |
| end | |
| local nms = require 'nms' | |
| require 'cunn' | |
| local threshold = 0 | |
| local validIndices = nms.gpu_nms(boxesTensor:cuda(), threshold) | |
| --boxesTensor = boxesTensor:index(1, validIndices) | |
| local newBoxes = {} | |
| for i = 1, validIndices:size(1) do | |
| table.insert(newBoxes, boxes[validIndices[i]]) | |
| end | |
| boxes = newBoxes | |
| end | |
| for number = 1, nClasses do | |
| table.sort(boxes, function(a, b) return a[4 + number] > b[4 + number] end) | |
| local boxesTensor = torch.zeros(#boxes, 5) | |
| for boxIndex, box in pairs(boxes) do | |
| boxesTensor[boxIndex] = torch.cat(box[{{1, 4}}], box[{{4 + number, 4 + number}}], 1) | |
| end | |
| if false then | |
| local nms = require 'nms' | |
| require 'cunn' | |
| local threshold = 0 | |
| local validIndices = nms.gpu_nms(boxesTensor:cuda(), threshold) | |
| boxesTensor = boxesTensor:index(1, validIndices) | |
| end | |
| for i = 1, math.min(numItemsPerClass, boxesTensor:size(1)) do | |
| local box = boxesTensor[i] | |
| if box[5] < confidenceThreshold then | |
| break | |
| end | |
| local itemInfo = utils.getItemInfo(mode, number) | |
| --table.insert(itemInfo, box[5]) | |
| local item = {{box[1], box[2]}, {box[3], box[4]}, itemInfo} | |
| table.insert(items, item) | |
| end | |
| end | |
| if true then | |
| local detectionImage = floorplan:clone() | |
| local colorMap = utils.getColorMap() | |
| local lineWidth = 2 | |
| for _, item in pairs(items) do | |
| local x_1 = math.min(math.max(torch.round(item[1][1]), 1 + lineWidth), width - lineWidth) | |
| local y_1 = math.min(math.max(torch.round(item[1][2]), 1 + lineWidth), height - lineWidth) | |
| local x_2 = math.max(math.min(torch.round(item[2][1]), width - lineWidth), 1 + lineWidth) | |
| local y_2 = math.max(math.min(torch.round(item[2][2]), height - lineWidth), 1 + lineWidth) | |
| --if item[3][4] > 0.5 then | |
| detectionImage = image.drawRect(detectionImage, x_1, y_1, x_2, y_2, {lineWidth = lineWidth, color = colorMap[number]}) | |
| --end | |
| end | |
| image.save('test/detection_result_' .. mode .. '.png', detectionImage) | |
| os.exit(1) | |
| end | |
| return items | |
| end | |
| function utils.evaluateResult(width, height, representationTarget, representationPrediction, opt) | |
| local lineWidth = opt.lineWidth or 5 | |
| local maxDim = math.max(width, height) | |
| if (not representationTarget.points or #representationTarget.points == 0) and (representationTarget.walls and #representationTarget.walls > 0) then | |
| representationTarget.walls = utils.mergeLines(representationTarget.walls, lineWidth) | |
| representationTarget.points = utils.linesToPoints(width, height, representationTarget.walls, lineWidth) | |
| end | |
| if (not representationPrediction.points or #representationPrediction.points == 0) and (representationPrediction.walls and #representationPrediction.walls > 0) then | |
| representationPrediction.points = utils.linesToPoints(width, height, representationPrediction.walls, lineWidth) | |
| end | |
| --local longWalls = utils.pointsToLines(width, height, points, lineWidth) | |
| --representationPrediction.points = utils.linesToPoints(width, height, longWalls, lineWidth) | |
| local pointMap = {} | |
| local numCorrectPredictions = 0 | |
| for pointIndexTarget, pointTarget in pairs(representationTarget.points) do | |
| local minDistance | |
| local minDistancePointIndexPrediction | |
| for pointIndexPrediction, pointPrediction in pairs(representationPrediction.points) do | |
| --[[ | |
| local distance = utils.calcDistance(pointPrediction[1], pointTarget[1]) | |
| if distance < opt.pointDistanceThreshold then | |
| print({'point', pointPrediction[1], pointPrediction[3], pointTarget[3]}) | |
| end | |
| ]]-- | |
| if utils.getNumber('points', pointPrediction[3]) == utils.getNumber('points', pointTarget[3]) then | |
| local distance = utils.calcDistance(pointPrediction[1], pointTarget[1]) | |
| if not minDistance or distance < minDistance then | |
| minDistancePointIndexPrediction = pointIndexPrediction | |
| minDistance = distance | |
| end | |
| end | |
| end | |
| if minDistance and minDistance < opt.pointDistanceThreshold * maxDim then | |
| pointMap[pointIndexTarget] = minDistancePointIndexPrediction | |
| numCorrectPredictions = numCorrectPredictions + 1 | |
| else | |
| print('missing wall junction') | |
| print(minDistance) | |
| print(pointTarget) | |
| --if minDistancePointIndexPrediction then | |
| --print(representationPrediction.points[minDistancePointIndexPrediction]) | |
| --end | |
| end | |
| end | |
| local result = {} | |
| for _, mode in pairs({'Wall Junction', 'Door', 'Object', 'Room'}) do | |
| result[mode] = {} | |
| for _, value in pairs({'numCorrectPredictions', 'numTargets', 'numPredictions'}) do | |
| result[mode][value] = 0 | |
| end | |
| end | |
| result["Wall Junction"].numCorrectPredictions = result["Wall Junction"].numCorrectPredictions + numCorrectPredictions | |
| result["Wall Junction"].numTargets = result["Wall Junction"].numTargets + #representationTarget.points | |
| result["Wall Junction"].numPredictions = result["Wall Junction"].numPredictions + #representationPrediction.points | |
| local doorMap = {} | |
| local numCorrectPredictions = 0 | |
| local doorNumberGroupMap = {} | |
| doorNumberGroupMap[1] = 1 | |
| doorNumberGroupMap[2] = 1 | |
| doorNumberGroupMap[3] = 2 | |
| doorNumberGroupMap[4] = 3 | |
| doorNumberGroupMap[5] = 3 | |
| doorNumberGroupMap[6] = 4 | |
| doorNumberGroupMap[7] = 5 | |
| doorNumberGroupMap[8] = 5 | |
| doorNumberGroupMap[9] = 6 | |
| doorNumberGroupMap[10] = 2 | |
| doorNumberGroupMap[11] = 2 | |
| for doorIndexTarget, doorTarget in pairs(representationTarget.doors) do | |
| local minDistance | |
| local minDistanceDoorIndexPrediction | |
| local numberTarget = doorNumberGroupMap[utils.getNumber('doors', doorTarget[3])] | |
| for doorIndexPrediction, doorPrediction in pairs(representationPrediction.doors) do | |
| local numberPrediction = doorPrediction[3][2] | |
| --[[ | |
| local distance = math.min(math.max(utils.calcDistance(doorPrediction[1], doorTarget[1]), utils.calcDistance(doorPrediction[2], doorTarget[2])), math.max(utils.calcDistance(doorPrediction[1], doorTarget[2]), utils.calcDistance(doorPrediction[2], doorTarget[1]))) | |
| if distance < opt.doorDistanceThreshold then | |
| print({'door', numberPrediction, numberTarget}) | |
| end | |
| ]]-- | |
| if numberPrediction == numberTarget or true then | |
| local distance = math.max(utils.calcDistance(doorPrediction[1], doorTarget[1]), utils.calcDistance(doorPrediction[2], doorTarget[2])) | |
| if not minDistance or distance < minDistance then | |
| minDistanceDoorIndexPrediction = doorIndexPrediction | |
| minDistance = distance | |
| end | |
| end | |
| end | |
| if minDistance and minDistance < opt.doorDistanceThreshold * maxDim then | |
| doorMap[doorIndexTarget] = minDistanceDoorIndexPrediction | |
| numCorrectPredictions = numCorrectPredictions + 1 | |
| else | |
| print('missing door') | |
| print(doorTarget) | |
| print(minDistance) | |
| if minDistanceDoorIndexPrediction then | |
| print(representationPrediction.doors[minDistanceDoorIndexPrediction]) | |
| end | |
| end | |
| end | |
| result["Door"].numCorrectPredictions = result["Door"].numCorrectPredictions + numCorrectPredictions | |
| result["Door"].numTargets = result["Door"].numTargets + #representationTarget.doors | |
| result["Door"].numPredictions = result["Door"].numPredictions + #representationPrediction.doors | |
| local iconMap = {} | |
| local numCorrectPredictions = 0 | |
| for iconIndexTarget, iconTarget in pairs(representationTarget.icons) do | |
| local maxIOU | |
| local maxIOUIconIndexPrediction | |
| for iconIndexPrediction, iconPrediction in pairs(representationPrediction.icons) do | |
| if utils.getNumber('icons', iconPrediction[3]) == utils.getNumber('icons', iconTarget[3]) or (iconTarget[3][1] == 'special' and iconPrediction[3][1] == 'special') then | |
| local IOU = utils.calcIOU(iconPrediction, iconTarget) | |
| if not maxIOU or IOU > maxIOU then | |
| maxIOUIconIndexPrediction = iconIndexPrediction | |
| maxIOU = IOU | |
| end | |
| end | |
| end | |
| if maxIOU and maxIOU > opt.iconIOUThreshold then | |
| iconMap[iconIndexTarget] = maxIOUIconIndexPrediction | |
| numCorrectPredictions = numCorrectPredictions + 1 | |
| else | |
| print('missing object') | |
| print(iconTarget) | |
| print(maxIOU) | |
| end | |
| end | |
| result["Object"].numCorrectPredictions = result["Object"].numCorrectPredictions + numCorrectPredictions | |
| result["Object"].numTargets = result["Object"].numTargets + #representationTarget.icons | |
| result["Object"].numPredictions = result["Object"].numPredictions + #representationPrediction.icons | |
| --[[ | |
| local segmentationTarget = utils.getSegmentation(width, height, representationTarget)[1] | |
| local segmentationPrediction = utils.getSegmentation(width, height, representationPrediction)[1] | |
| local segmentMasksTarget = {} | |
| for segmentIndex = 1, 11 do | |
| local segmentMask = segmentationTarget:eq(segmentIndex):double() | |
| local segments, numSegments = utils.findConnectedComponents(segmentMask) | |
| for segment = 1, numSegments - 1 do | |
| local mask = segments:eq(segment) | |
| table.insert(segmentMasksTarget, {mask, segmentIndex}) | |
| end | |
| end | |
| local segmentMasksPrediction = {} | |
| for segmentIndex = 1, 11 do | |
| local segmentMask = segmentationPrediction:eq(segmentIndex):double() | |
| local segments, numSegments = utils.findConnectedComponents(segmentMask) | |
| for segment = 1, numSegments - 1 do | |
| local mask = segments:eq(segment) | |
| table.insert(segmentMasksPrediction, {mask, segmentIndex}) | |
| end | |
| end | |
| ]]-- | |
| local segmentMasksTarget = utils.getRoomSegments(width, height, representationTarget) | |
| local segmentMasksPrediction = utils.getRoomSegments(width, height, representationPrediction) | |
| --[[ | |
| for segmentIndexTarget, segmentTarget in pairs(segmentMasksPrediction) do | |
| print(segmentTarget[2]) | |
| print(segmentTarget[1]:nonzero()[1]) | |
| print(#segmentTarget[1]:nonzero()) | |
| end | |
| print(segmentMasksPrediction.labels) | |
| os.exit(1) | |
| ]]-- | |
| local segmentMap = {} | |
| local numCorrectPredictions = 0 | |
| local roomTypeMap = {} | |
| for roomType = 1, 10 do | |
| roomTypeMap[roomType] = {} | |
| roomTypeMap[roomType][roomType] = true | |
| end | |
| roomTypeMap[1][2] = true | |
| roomTypeMap[2][1] = true | |
| roomTypeMap[2][3] = true | |
| roomTypeMap[3][2] = true | |
| roomTypeMap[4][5] = true | |
| roomTypeMap[5][4] = true | |
| roomTypeMap[7][10] = true | |
| roomTypeMap[8][1] = true | |
| roomTypeMap[8][2] = true | |
| roomTypeMap[9][4] = true | |
| roomTypeMap[9][5] = true | |
| roomTypeMap[10][4] = true | |
| roomTypeMap[10][5] = true | |
| roomTypeMap[10][7] = true | |
| roomTypeMap[10][8] = true | |
| roomTypeMap[10][9] = true | |
| --local IOU = utils.calcIOUMask(segmentMasksPrediction[4][1], segmentMasksTarget[4][1]) | |
| --print(IOU) | |
| --os.exit(1) | |
| for segmentIndexTarget, segmentTarget in pairs(segmentMasksTarget) do | |
| local maxIOU | |
| local maxIOUSegmentIndexPrediction | |
| for segmentIndexPrediction, segmentPrediction in pairs(segmentMasksPrediction) do | |
| --if segmentPrediction[2] == segmentTarget[2] then | |
| if roomTypeMap[segmentTarget[2]][segmentPrediction[2]] then | |
| local IOU = utils.calcIOUMask(segmentPrediction[1], segmentTarget[1]) | |
| if not maxIOU or IOU > maxIOU then | |
| maxIOUSegmentIndexPrediction = segmentIndexPrediction | |
| maxIOU = IOU | |
| end | |
| end | |
| end | |
| if maxIOU and maxIOU > opt.segmentIOUThreshold then | |
| segmentMap[segmentIndexTarget] = maxIOUSegmentIndexPrediction | |
| numCorrectPredictions = numCorrectPredictions + 1 | |
| else | |
| print('missing segment') | |
| print(segmentTarget[2]) | |
| print({segmentTarget[1]:nonzero()[1][2], segmentTarget[1]:nonzero()[1][1]}) | |
| print(maxIOU) | |
| end | |
| end | |
| result["Room"].numCorrectPredictions = result["Room"].numCorrectPredictions + numCorrectPredictions | |
| result["Room"].numTargets = result["Room"].numTargets + #segmentMasksTarget | |
| local numSegmentPredictions = 0 | |
| for segmentIndexPrediction, segmentPrediction in pairs(segmentMasksPrediction) do | |
| numSegmentPredictions = numSegmentPredictions + 1 | |
| end | |
| result["Room"].numPredictions = result["Room"].numPredictions + numSegmentPredictions | |
| return result | |
| end | |
| function utils.evaluateDetection(floorplan, representation, mode) | |
| local mode = mode or 'icons' | |
| --detectionImage = image.drawRect(detectionImage, x_1, y_1, x_2, y_2, {color = colorMap[number], lineWidth = lineWidth}) | |
| local width, height = floorplan:size(3), floorplan:size(2) | |
| local patchDim = 16 * 7 | |
| local nClasses = 13 | |
| if not utils.modelIcon then | |
| require 'nnx' | |
| local inn = require 'inn' | |
| if not nn.SpatialConstDiagonal then | |
| torch.class('nn.SpatialConstDiagonal', 'inn.ConstAffine') | |
| end | |
| --paths.dofile('/home/chenliu/Projects/Floorplan/floorplan/InverseCAD/models/SpatialSymmetricPadding.lua') | |
| utils.modelIcon = torch.load('/home/chenliu/Projects/Floorplan/models/deepmask-grid-' .. mode .. '/model_best.t7') | |
| end | |
| --utils.modelIcon:evaluate() | |
| local trunk = utils.modelIcon.modules[1] | |
| local coordinatesBranch = utils.modelIcon.modules[2] | |
| local labelBranch = utils.modelIcon.modules[3] | |
| local scales = {} | |
| --for i = -1.5,0.5,1.5 do | |
| --table.insert(scales, 2^i) | |
| --end | |
| table.insert(scales, 1) | |
| local meanstd = {mean = { 0.5, 0.5, 0.5 }, std = { 0.5, 0.5, 0.5 }} | |
| local input = floorplan | |
| if input:type() == 'torch.CudaTensor' then input = input:float() end | |
| local pyramid = nn.ConcatTable() | |
| for i = 1,#scales do | |
| pyramid:add(nn.SpatialReSamplingEx{rwidth=scales[i], | |
| rheight=scales[i], mode='bilinear'}) | |
| end | |
| local inpPyramid = pyramid:forward(input) | |
| local inpPad = torch.CudaTensor() | |
| -- forward all scales through network | |
| --local outPyramidMask,outPyramidScore = {},{} | |
| local boxes = {} | |
| local softmax = nn.SoftMax():cuda() | |
| for i,_ in pairs(inpPyramid) do | |
| local inp = inpPyramid[i]:cuda() | |
| local h,w = inp:size(2),inp:size(3) | |
| -- padding/normalize | |
| --inpPad:resize(1,3,h+2*bw,w+2*bw):fill(.5) | |
| --inpPad:narrow(1,1,1):narrow(3,bw+1,h):narrow(4,bw+1,w):copy(inp) | |
| --for i=1,3 do | |
| --inpPad[1][i]:add(-meanstd.mean[i]):div(meanstd.std[i]) | |
| --end | |
| inpPad = inp:repeatTensor(1, 1, 1, 1) | |
| inpPad[1][i]:add(-meanstd.mean[i]):div(meanstd.std[i]) | |
| -- forward trunk | |
| local outTrunk = trunk:forward(inpPad) | |
| -- forward score branch | |
| local outLabel = labelBranch:forward(outTrunk) | |
| --table.insert(outPyramidScore,outScore:clone():squeeze()) | |
| -- forward mask branch | |
| local outCoordinates = coordinatesBranch:forward(outTrunk) | |
| --table.insert(outPyramidMask,outMask:float():squeeze()) | |
| --outScore = outScore:clone():squeeze() | |
| --outMask = outMask:float():squeeze() | |
| outCoordinates = outCoordinates:squeeze() | |
| local gridHeight, gridWidth = outCoordinates:size(2), outCoordinates:size(3) | |
| outLabel = outLabel:view(gridHeight, gridWidth, -1):transpose(2, 3):transpose(1, 2) | |
| local cellWidth = w / gridWidth | |
| local cellHeight = h / gridHeight | |
| for gridX = 1, gridWidth do | |
| for gridY = 1, gridHeight do | |
| local labelProbs = outLabel[{{}, gridY, gridX}] | |
| labelProbs = softmax:forward(labelProbs) | |
| local box = torch.cat(outCoordinates[{{}, gridY, gridX}], labelProbs, 1):double() | |
| local x_1 = cellWidth * (gridX - 1 + 0.5) + (box[1] - 0.5) * patchDim - math.exp(box[3]) * patchDim / 2 | |
| local y_1 = cellHeight * (gridY - 1 + 0.5) + (box[2] - 0.5) * patchDim - math.exp(box[4]) * patchDim / 2 | |
| local x_2 = cellWidth * (gridX - 1 + 0.5) + (box[1] - 0.5) * patchDim + math.exp(box[3]) * patchDim / 2 | |
| local y_2 = cellHeight * (gridY - 1 + 0.5) + (box[2] - 0.5) * patchDim + math.exp(box[4]) * patchDim / 2 | |
| box[1] = x_1 / w * width | |
| box[2] = y_1 / h * height | |
| box[3] = x_2 / w * width | |
| box[4] = y_2 / h * height | |
| table.insert(boxes, box) | |
| if false then | |
| print(gridX .. ' ' .. gridY) | |
| print(box) | |
| local sample = image.crop(floorplan, math.max(cellWidth * (gridX - 1 + 0.5) - patchDim / 2, 0), math.max(cellHeight * (gridY - 1 + 0.5) - patchDim / 2, 0), math.min(cellWidth * (gridX - 1 + 0.5) + patchDim / 2, width), math.min(cellHeight * (gridY - 1 + 0.5) + patchDim / 2, height)) | |
| image.save('test/samples/sample_' .. gridX .. '_' .. gridY .. '.png', sample) | |
| end | |
| end | |
| end | |
| end | |
| --os.exit(1) | |
| local APs = {} | |
| for number = 1, nClasses do | |
| table.sort(boxes, function(a, b) return a[4 + number] > b[4 + number] end) | |
| local boxesTensor = torch.zeros(#boxes, 5) | |
| for boxIndex, box in pairs(boxes) do | |
| boxesTensor[boxIndex] = torch.cat(box[{{1, 4}}], box[{{4 + number, 4 + number}}], 1) | |
| end | |
| local nms = require 'nms' | |
| require 'cunn' | |
| local threshold = 0 | |
| local validIndices = nms.gpu_nms(boxesTensor:cuda(), threshold) | |
| boxesTensor = boxesTensor:index(1, validIndices) | |
| local numberItems = {} | |
| for _, item in pairs(representation[mode]) do | |
| if utils.getNumber(mode, item[3]) == number then | |
| table.insert(numberItems, item) | |
| end | |
| end | |
| if #numberItems > 0 then | |
| local AP = utils.calcAP(boxesTensor, numberItems) | |
| APs[number] = AP | |
| print(number .. ' ' .. AP) | |
| end | |
| end | |
| local APSum = 0 | |
| local count = 0 | |
| for _, AP in pairs(APs) do | |
| APSum = APSum + AP | |
| count = count + 1 | |
| end | |
| print(APSum / count) | |
| local detectionImage = floorplan:clone() | |
| local colorMap = utils.getColorMap() | |
| local lineWidth = 2 | |
| local numProposals = 10 | |
| for number = 1, nClasses do | |
| --table.sort(boxes, function(a, b) return torch.max(a[{{5, 4 + nClasses}}]) > torch.max(b[{{5, 4 + nClasses}}]) end) | |
| table.sort(boxes, function(a, b) return a[4 + number] > b[4 + number] end) | |
| local boxesTensor = torch.zeros(#boxes, 5) | |
| for boxIndex, box in pairs(boxes) do | |
| boxesTensor[boxIndex] = torch.cat(box[{{1, 4}}], box[{{4 + number, 4 + number}}], 1) | |
| end | |
| --print(boxesTensor[{{}, 5}]) | |
| local nms = require 'nms' | |
| require 'cunn' | |
| --require 'cutorch' | |
| local threshold = 0 | |
| local validIndices = nms.gpu_nms(boxesTensor:cuda(), threshold) | |
| boxesTensor = boxesTensor:index(1, validIndices) | |
| --print(boxesTensor[{{}, 5}]) | |
| --os.exit(1) | |
| --print(colorMap[number][1]) | |
| --print(colorMap[number][2]) | |
| --print(colorMap[number][3]) | |
| for i = 1, math.min(numProposals, boxesTensor:size(1)) do | |
| local box = boxesTensor[i] | |
| local x_1 = math.min(math.max(torch.round(box[1]), 1 + lineWidth), width - lineWidth) | |
| local y_1 = math.min(math.max(torch.round(box[2]), 1 + lineWidth), height - lineWidth) | |
| local x_2 = math.max(math.min(torch.round(box[3]), width - lineWidth), 1 + lineWidth) | |
| local y_2 = math.max(math.min(torch.round(box[4]), height - lineWidth), 1 + lineWidth) | |
| --local prob, pred = torch.max(box[{{5, -1}}], 1) | |
| if number == 1 then | |
| --print(box) | |
| --local prob = box[5] | |
| print(box[5] .. ' ' .. number) | |
| end | |
| if box[5] > 0.5 then | |
| detectionImage = image.drawRect(detectionImage, x_1, y_1, x_2, y_2, {lineWidth = lineWidth, color = colorMap[number]}) | |
| end | |
| end | |
| end | |
| image.save('test/icon_result.png', detectionImage) | |
| os.exit(1) | |
| end | |
| function utils.calcIOU(box_1, box_2) | |
| local intersectionX = (math.min(box_1[2][1], box_2[2][1]) - math.max(box_1[1][1], box_2[1][1])) | |
| local intersectionY = (math.min(box_1[2][2], box_2[2][2]) - math.max(box_1[1][2], box_2[1][2])) | |
| if intersectionX <= 0 or intersectionY <= 0 then | |
| return 0 | |
| end | |
| local intersection = intersectionX * intersectionY | |
| local union = (box_1[2][1] - box_1[1][1]) * (box_1[2][2] - box_1[1][2]) + (box_2[2][1] - box_2[1][1]) * (box_2[2][2] - box_2[1][2]) - intersection | |
| return intersection / union | |
| end | |
| function utils.calcIOUMask(mask_1, mask_2) | |
| local unionIndices = (mask_1 + mask_2):nonzero() | |
| if ##unionIndices > 0 then | |
| local intersectionIndices = torch.cmul(mask_1, mask_2):nonzero() | |
| if ##intersectionIndices > 0 then | |
| return intersectionIndices:size(1) / unionIndices:size(1) | |
| end | |
| end | |
| return 0 | |
| end | |
| function utils.calcAP(boxesTensor, items, IOUThreshold) | |
| local IOUThreshold = IOUThreshold or 0.5 | |
| local groundTruthMap = {} | |
| local maxBoxIndex = 0 | |
| for itemIndex, item in pairs(items) do | |
| local boxIndex | |
| local prevIOU | |
| local prevScore | |
| for i = 1, boxesTensor:size(1) do | |
| local box = boxesTensor[i] | |
| local IOU = utils.calcIOU({{box[1], box[2]}, {box[3], box[4]}}, item) | |
| local score = box[5] | |
| if prevScore and score < prevScore then | |
| if not boxIndex and IOU > IOUThreshold then | |
| boxIndex = i | |
| end | |
| break | |
| else | |
| if IOU > IOUThreshold and (not prevIOU or IOU > prevIOU) then | |
| boxIndex = i | |
| prevIOU = IOU | |
| prevScore = score | |
| end | |
| end | |
| end | |
| if boxIndex then | |
| groundTruthMap[boxIndex] = itemIndex | |
| maxBoxIndex = math.max(maxBoxIndex, boxIndex) | |
| end | |
| end | |
| local precisions = {} | |
| local numPositiveBoxes = 0 | |
| for boxIndex = 1, maxBoxIndex do | |
| if groundTruthMap[boxIndex] then | |
| numPositiveBoxes = numPositiveBoxes + 1 | |
| end | |
| local recall = numPositiveBoxes / #items | |
| local recallLevel = math.floor(recall * 10) + 1 | |
| local precision = numPositiveBoxes / boxIndex | |
| --print(recallLevel .. ' ' .. precision) | |
| if not precisions[recallLevel] then | |
| precisions[recallLevel] = precision | |
| else | |
| precisions[recallLevel] = math.max(precision, precisions[recallLevel]) | |
| end | |
| end | |
| local AP = 0 | |
| for recallLevel = 10, 1, -1 do | |
| if not precisions[recallLevel + 1] then | |
| break | |
| end | |
| if not precisions[recallLevel] or precisions[recallLevel] < precisions[recallLevel + 1] then | |
| precisions[recallLevel] = precisions[recallLevel + 1] | |
| end | |
| end | |
| for recallLevel = 1, 11 do | |
| --print(precisions[recallLevel]) | |
| if precisions[recallLevel] then | |
| AP = AP + precisions[recallLevel] | |
| end | |
| end | |
| AP = AP / 11 | |
| return AP | |
| end | |
| function utils.extractMaximum(heatmap, numPoints) | |
| local width, height = heatmap:size(2), heatmap:size(1) | |
| local lineWidth = lineWidth or 5 | |
| local mask = heatmap:clone() | |
| local points = {} | |
| for pointIndex = 1, numPoints do | |
| local maxValues, ys = torch.max(mask, 1) | |
| local maxValue, x = torch.max(maxValues, 2) | |
| if maxValue[1][1] <= 0 then | |
| break | |
| end | |
| x = x[1][1] | |
| local y = ys[1][x] | |
| table.insert(points, {x, y}) | |
| mask[{{math.max(y - lineWidth, 1), math.min(y + lineWidth, height)}, {math.max(x - lineWidth, 1), math.min(x + lineWidth, width)}}]:fill(0) | |
| end | |
| return points | |
| end | |
| function utils.extractLinePoints(mask_1, mask_2, lineDim, numPoints, lineWidth, lineMinLength, getHeatmaps) | |
| local width, height = mask_1:size(2), mask_1:size(1) | |
| local lineWidth = lineWidth or 5 | |
| local lineMinLength = lineMinLength or 10 | |
| local maxPoolingLine | |
| if lineDim == 1 then | |
| maxPoolingLine = nn.SpatialMaxPooling(1, lineWidth * 2 + 1, 1, 1, 0, lineWidth) | |
| --maxPoolingLine = nn.SpatialConvolution(1, 1, 1, lineWidth * 2 + 1, 1, 1, 0, lineWidth) | |
| --maxPoolingLine.weight = image.gaussian1D(lineWidth * 2 + 1):view(#maxPoolingLine.weight) | |
| --maxPoolingLine:noBias() | |
| else | |
| maxPoolingLine = nn.SpatialMaxPooling(lineWidth * 2 + 1, 1, 1, 1, lineWidth, 0) | |
| --maxPoolingLine = nn.SpatialConvolution(1, 1, lineWidth * 2 + 1, 1, 1, 1, lineWidth, 0) | |
| --maxPoolingLine.weight = image.gaussian1D(lineWidth * 2 + 1):view(#maxPoolingLine.weight) | |
| --maxPoolingLine:noBias() | |
| end | |
| local lineConfidences = maxPoolingLine:forward(torch.cat(mask_1:repeatTensor(1, 1, 1), mask_2:repeatTensor(1, 1, 1), 1)) | |
| local lineConfidence_1 = lineConfidences[1] | |
| local lineConfidence_2 = lineConfidences[2] | |
| local heatmap_2 = torch.zeros(height, width) | |
| for i = lineMinLength + 1, mask_1:size(3 - lineDim) do | |
| if lineDim == 1 then | |
| heatmap_2[{{}, i}] = torch.cmax(heatmap_2[{{}, i - 1}], lineConfidence_1[{{}, i - lineMinLength}]) | |
| else | |
| heatmap_2[i] = torch.cmax(heatmap_2[i - 1], lineConfidence_1[i - lineMinLength]) | |
| end | |
| end | |
| local heatmap_1 = torch.zeros(height, width) | |
| for i = mask_1:size(3 - lineDim) - lineMinLength, 1, -1 do | |
| if lineDim == 1 then | |
| heatmap_1[{{}, i}] = torch.cmax(heatmap_1[{{}, i + 1}], lineConfidence_2[{{}, i + lineMinLength}]) | |
| else | |
| heatmap_1[i] = torch.cmax(heatmap_1[i + 1], lineConfidence_2[i + lineMinLength]) | |
| end | |
| end | |
| --local maxPoolingSquare = nn.SpatialMaxPooling(lineWidth * 2 + 1, lineWidth * 2 + 1, 1, 1, lineWidth, lineWidth) | |
| local maxPoolingSquare = nn.SpatialConvolution(1, 1, lineWidth * 2 + 1, lineWidth * 2 + 1, 1, 1, lineWidth, lineWidth) | |
| maxPoolingSquare.weight = image.gaussian(lineWidth * 2 + 1):view(#maxPoolingSquare.weight) | |
| maxPoolingSquare:noBias() | |
| --local pointConfidence_1 = maxPoolingSquare:forward(mask_1:repeatTensor(1, 1, 1))[1]:clone() | |
| --local pointConfidence_2 = maxPoolingSquare:forward(mask_2:repeatTensor(1, 1, 1))[1]:clone() | |
| local pointConfidence_1 = mask_1 | |
| local pointConfidence_2 = mask_2 | |
| --heatmap_1:add(pointConfidence_1) | |
| --heatmap_2:add(pointConfidence_2) | |
| --local heatmap_test = heatmap_1:clone() | |
| heatmap_1:cmul(pointConfidence_1) | |
| heatmap_2:cmul(pointConfidence_2) | |
| if getHeatmaps then | |
| return heatmap_1, heatmap_2 | |
| end | |
| local points_1 = utils.extractMaximum(heatmap_1, numPoints) | |
| local points_2 = utils.extractMaximum(heatmap_2, numPoints) | |
| collectgarbage() | |
| if #points_1 > 0 and false then | |
| image.save('test/mask_1.png', mask_1) | |
| image.save('test/mask_2.png', mask_2) | |
| image.save('test/heatmap_test.png', heatmap_test) | |
| image.save('test/heatmap.png', heatmap_1) | |
| image.save('test/point_confidence.png', pointConfidence_1) | |
| image.save('test/line_confidence.png', lineConfidence_2) | |
| os.exit(1) | |
| end | |
| return points_1, points_2 | |
| end | |
| function utils.extractRectanglePoints(mask_1, mask_2, mask_3, mask_4, numPoints, getHeatmap) | |
| local heatmap_1, heatmap_2 = utils.extractLinePoints(mask_1, mask_2, 1, nil, nil, nil, true) | |
| local heatmap_3, heatmap_4 = utils.extractLinePoints(mask_3, mask_4, 1, nil, nil, nil, true) | |
| heatmap_1, heatmap_3 = utils.extractLinePoints(heatmap_1, heatmap_3, 2, nil, 0, nil, true) | |
| heatmap_2, heatmap_4 = utils.extractLinePoints(heatmap_2, heatmap_4, 2, nil, 0, nil, true) | |
| if getHeatmap then | |
| return heatmap_1, heatmap_2, heatmap_3, heatmap_4 | |
| end | |
| local points_1 = utils.extractMaximum(heatmap_1, numPoints) | |
| local points_2 = utils.extractMaximum(heatmap_2, numPoints) | |
| local points_3 = utils.extractMaximum(heatmap_3, numPoints) | |
| local points_4 = utils.extractMaximum(heatmap_4, numPoints) | |
| if #points_1 > 0 and false then | |
| image.save('test/mask_1.png', mask_1) | |
| image.save('test/mask_2.png', mask_2) | |
| image.save('test/mask_3.png', mask_3) | |
| image.save('test/mask_4.png', mask_4) | |
| local heatmap_test_1, heatmap_test_2 = utils.extractLinePoints(mask_1, mask_2, 1, nil, nil, nil, true) | |
| local heatmap_test_3, heatmap_test_4 = utils.extractLinePoints(mask_3, mask_4, 1, nil, nil, nil, true) | |
| image.save('test/heatmap_test_1.png', heatmap_test_1) | |
| image.save('test/heatmap_test_2.png', heatmap_test_2) | |
| image.save('test/heatmap_test_3.png', heatmap_test_3) | |
| image.save('test/heatmap_test_4.png', heatmap_test_4) | |
| image.save('test/heatmap_1.png', heatmap_1) | |
| image.save('test/heatmap_2.png', heatmap_2) | |
| image.save('test/heatmap_3.png', heatmap_3) | |
| image.save('test/heatmap_4.png', heatmap_4) | |
| print(#points_1 .. ' ' .. #points_2 .. ' ' .. #points_3 .. ' ' .. #points_4) | |
| end | |
| return points_1, points_2, points_3, points_4 | |
| end | |
| function utils.extractWallPoints(masks, numPoints, lineWidth, lineMinLength, getHeatmaps) | |
| local width, height = masks[1]:size(2), masks[1]:size(1) | |
| local lineWidth = lineWidth or 5 | |
| local lineMinLength = lineMinLength or 10 | |
| local maxPoolingLine_1 = nn.SpatialMaxPooling(1, lineWidth * 2 + 1, 1, 1, 0, lineWidth) | |
| local maxPoolingLine_2 = nn.SpatialMaxPooling(lineWidth * 2 + 1, 1, 1, 1, lineWidth, 0) | |
| local lineConfidences = torch.zeros(4, height, width) | |
| for i = 1, 4 do | |
| lineConfidences[i]:cmax(masks[i]) | |
| end | |
| for i = 1, 4 do | |
| lineConfidences[i % 4 + 1]:cmax(masks[4 + i]) | |
| lineConfidences[(i + 1) % 4 + 1]:cmax(masks[4 + i]) | |
| end | |
| for i = 1, 4 do | |
| lineConfidences[i]:cmax(masks[8 + i]) | |
| lineConfidences[i % 4 + 1]:cmax(masks[8 + i]) | |
| lineConfidences[(i + 2) % 4 + 1]:cmax(masks[8 + i]) | |
| end | |
| for i = 1, 4 do | |
| lineConfidences[i]:cmax(masks[13]) | |
| end | |
| local output_1 = maxPoolingLine_1:forward(torch.cat(lineConfidences[2]:repeatTensor(1, 1, 1), lineConfidences[4]:repeatTensor(1, 1, 1), 1)) | |
| local output_2 = maxPoolingLine_2:forward(torch.cat(lineConfidences[1]:repeatTensor(1, 1, 1), lineConfidences[3]:repeatTensor(1, 1, 1), 1)) | |
| lineConfidences[2] = output_1[1] | |
| lineConfidences[4] = output_1[2] | |
| lineConfidences[1] = output_2[1] | |
| lineConfidences[3] = output_2[2] | |
| local lineHeatmaps = torch.zeros(4, height, width) | |
| for i = lineMinLength + 1, width do | |
| lineHeatmaps[4][{{}, i}] = torch.cmax(lineHeatmaps[4][{{}, i - 1}], lineConfidences[4][{{}, i - lineMinLength}]) | |
| end | |
| for i = lineMinLength + 1, height do | |
| lineHeatmaps[1][i] = torch.cmax(lineHeatmaps[1][i - 1], lineConfidences[1][i - lineMinLength]) | |
| end | |
| for i = width - lineMinLength, 1, -1 do | |
| lineHeatmaps[2][{{}, i}] = torch.cmax(lineHeatmaps[2][{{}, i + 1}], lineConfidences[2][{{}, i + lineMinLength}]) | |
| end | |
| for i = height - lineMinLength, 1, -1 do | |
| lineHeatmaps[3][i] = torch.cmax(lineHeatmaps[3][i + 1], lineConfidences[3][i + lineMinLength]) | |
| end | |
| --[[ | |
| for i = 1, 4 do | |
| image.save('test/line_' .. i .. '.png', lineHeatmaps[i]) | |
| end | |
| os.exit(1) | |
| ]]-- | |
| --local maxPoolingSquare = nn.SpatialMaxPooling(lineWidth * 2 + 1, lineWidth * 2 + 1, 1, 1, lineWidth, lineWidth) | |
| --local maxPoolingSquare = nn.SpatialConvolution(1, 1, lineWidth * 2 + 1, lineWidth * 2 + 1, 1, 1, lineWidth, lineWidth) | |
| --maxPoolingSquare.weight = image.gaussian(lineWidth * 2 + 1):view(#maxPoolingSquare.weight) | |
| --maxPoolingSquare:noBias() | |
| --local pointConfidence_1 = maxPoolingSquare:forward(mask_1:repeatTensor(1, 1, 1))[1]:clone() | |
| --local pointConfidence_2 = maxPoolingSquare:forward(mask_2:repeatTensor(1, 1, 1))[1]:clone() | |
| local heatmaps = masks | |
| for i = 1, 4 do | |
| heatmaps[i]:cmul(lineHeatmaps[(i + 1) % 4 + 1]) | |
| end | |
| for i = 1, 4 do | |
| heatmaps[i]:cmul(lineHeatmaps[i]) | |
| heatmaps[i]:cmul(lineHeatmaps[(i + 2) % 4 + 1]) | |
| end | |
| for i = 1, 4 do | |
| heatmaps[i]:cmul(lineHeatmaps[i % 4 + 1]) | |
| heatmaps[i]:cmul(lineHeatmaps[(i + 1) % 4 + 1]) | |
| heatmaps[i]:cmul(lineHeatmaps[(i + 2) % 4 + 1]) | |
| end | |
| for i = 1, 4 do | |
| heatmaps[13]:cmul(lineHeatmaps[i]) | |
| end | |
| if getHeatmaps then | |
| return heatmaps | |
| end | |
| local points = {} | |
| for i = 1, 13 do | |
| table.insert(points, utils.extractMaximum(heatmaps[i], numPoints)) | |
| end | |
| return points | |
| end | |
| function utils.estimateHeatmaps(modelHeatmap, floorplan, scaleType) | |
| local scaleType = scaleType or 'single' | |
| local width, height = floorplan:size(3), floorplan:size(2) | |
| local sampleDim = 256 | |
| package.path = 'datasets/?.lua;' .. package.path | |
| package.path = '?.lua;' .. package.path | |
| local dataset = require('floorplan-representation') | |
| dataset.split = 'val' | |
| local output | |
| if scaleType == 'single' then | |
| --local input = dataset:preprocessResize(sampleDim, sampleDim)(floorplan):repeatTensor(1, 1, 1, 1):cuda() | |
| --output = modelHeatmap:forward(input)[1]:double() | |
| local floorplanScaled = dataset:preprocessScale(sampleDim)(floorplan) | |
| --image.save('test/scale.png', floorplanScaled:double()) | |
| local offsetX, offsetY | |
| if floorplanScaled:size(2) < sampleDim then | |
| offsetY = math.floor((sampleDim - floorplanScaled:size(2)) / 2) | |
| offsetX = 0 | |
| else | |
| offsetX = math.floor((sampleDim - floorplanScaled:size(3)) / 2) | |
| offsetY = 0 | |
| end | |
| local temp = torch.zeros(3, sampleDim, sampleDim) | |
| temp:narrow(2, offsetY + 1, floorplanScaled:size(2)):narrow(3, offsetX + 1, floorplanScaled:size(3)):copy(floorplanScaled) | |
| local input = temp:repeatTensor(1, 1, 1, 1):cuda() | |
| output = modelHeatmap:forward(input)[1]:double() | |
| output = image.crop(output, offsetX, offsetY, offsetX + floorplanScaled:size(3), offsetY + floorplanScaled:size(2)) | |
| --[[ | |
| image.save('test/heatmaps/floorplan.png', dataset:postprocess()(input[1])) | |
| for i = 1, 13 do | |
| image.save('test/heatmaps/junction_heatmap_' .. i .. '.png', output[i]) | |
| end | |
| os.exit(1) | |
| ]]-- | |
| elseif scaleType == 'full' then | |
| local floorplanNormalized = dataset:preprocessNormalization()(floorplan) | |
| else | |
| local floorplanNormalized = dataset:preprocessNormalization()(floorplan) | |
| local inputParamMap = {} | |
| local scale = 0 | |
| local inputs | |
| while math.max(floorplanNormalized:size(3), floorplanNormalized:size(2)) > sampleDim / 2 or not inputs do | |
| for offsetX = 0, floorplanNormalized:size(3) - 1, sampleDim do | |
| for offsetY = 0, floorplanNormalized:size(2) - 1, sampleDim do | |
| local input = torch.zeros(3, sampleDim, sampleDim) | |
| local inputWidth = math.min(sampleDim, floorplanNormalized:size(3) - offsetX) | |
| local inputHeight = math.min(sampleDim, floorplanNormalized:size(2) - offsetY) | |
| input:narrow(2, 1, inputHeight):narrow(3, 1, inputWidth):copy(image.crop(floorplanNormalized, offsetX, offsetY, offsetX + inputWidth, offsetY + inputHeight)) | |
| input = input:repeatTensor(1, 1, 1, 1) | |
| if not inputs then | |
| inputs = input | |
| else | |
| inputs = torch.cat(inputs, input, 1) | |
| end | |
| inputParamMap[inputs:size(1)] = {scale, offsetX, offsetY} | |
| end | |
| end | |
| scale = scale + 1 | |
| floorplanNormalized = image.scale(floorplanNormalized, floorplanNormalized:size(3) / 2, floorplanNormalized:size(2) / 2) | |
| end | |
| inputs = inputs:cuda() | |
| local outputs = torch.zeros(inputs:size(1), 51, sampleDim, sampleDim) | |
| for batchOffset = 0, inputs:size(1) - 1, 4 do | |
| local numInputs = math.min(inputs:size(1) - batchOffset, 4) | |
| outputs[{{batchOffset + 1, batchOffset + numInputs}}]:copy(modelHeatmap:forward(inputs[{{batchOffset + 1, batchOffset + numInputs}}]):double()) | |
| end | |
| local scaleWeights = {} | |
| for i = 0, 10 do | |
| scaleWeights[i] = 1 | |
| end | |
| local prediction = torch.zeros(51, height, width) | |
| for i = 1, inputs:size(1) do | |
| local inputParam = inputParamMap[i] | |
| local scaleFactor = 2^inputParam[1] | |
| local output = image.scale(outputs[i], outputs[i]:size(3) * scaleFactor, outputs[i]:size(2) * scaleFactor) | |
| local offsetX = inputParam[2] * scaleFactor | |
| local offsetY = inputParam[3] * scaleFactor | |
| local outputWidth = math.min(width - offsetX, output:size(3)) | |
| local outputHeight = math.min(height - offsetY, output:size(2)) | |
| prediction:narrow(2, offsetY + 1, outputHeight):narrow(3, offsetX + 1, outputWidth):add(output:narrow(2, 1, outputHeight):narrow(3, 1, outputWidth) * scaleWeights[inputParam[1]]) | |
| image.save('test/output_' .. i .. '.png', prediction:narrow(1, 1, 13):sum(1)[1]) | |
| end | |
| local scaleSum = 0 | |
| for i = 0, scale - 1 do | |
| scaleSum = scaleSum + scaleWeights[i] | |
| end | |
| prediction:div(scaleSum) | |
| output = prediction | |
| end | |
| if true then | |
| output = image.scale(output:double(), width, height, 'bicubic') | |
| local doorOffset = 13 | |
| local doorHeatmaps = torch.cat(torch.cat(torch.cat(output:narrow(1, doorOffset + 3, 1), output:narrow(1, doorOffset + 2, 1), 1), output:narrow(1, doorOffset + 4, 1), 1), output:narrow(1, doorOffset + 1, 1), 1) | |
| local iconOffset = 17 | |
| local iconHeatmaps = torch.cat(torch.cat(torch.cat(output:narrow(1, iconOffset + 4, 1), output:narrow(1, iconOffset + 3, 1), 1), output:narrow(1, iconOffset + 1, 1), 1), output:narrow(1, iconOffset + 2, 1), 1) | |
| return output:narrow(1, 1, 13):double(), doorHeatmaps, iconHeatmaps, output:narrow(1, 22, 30):double() | |
| end | |
| --local junctionHeatmaps = torch.zeros(nClasses, height, width) | |
| --local junctionHeatmaps = output:narrow(1, 1, 13):double() | |
| local confidenceMasks = output:narrow(1, 1, 13):double() | |
| local junctionHeatmaps = fp_ut.extractWallPoints(confidenceMasks, nil, nil, nil, true) | |
| --[[ | |
| for i = 1, 13 do | |
| image.save('test/mask_' .. i .. '.png', confidenceMasks[i]) | |
| end | |
| for i = 1, 13 do | |
| image.save('test/heatmap_' .. i .. '.png', junctionHeatmaps[i]) | |
| end | |
| ]]-- | |
| local confidenceMasks = output:narrow(1, 13 + 1, 4):double() | |
| local heatmap_1, heatmap_2 = fp_ut.extractLinePoints(confidenceMasks[1], confidenceMasks[2], 1, nil, nil, nil, true) | |
| local heatmap_3, heatmap_4 = fp_ut.extractLinePoints(confidenceMasks[3], confidenceMasks[4], 2, nil, nil, nil, true) | |
| local doorHeatmaps = torch.cat(torch.cat(torch.cat(heatmap_3:repeatTensor(1, 1, 1), heatmap_2:repeatTensor(1, 1, 1), 1), heatmap_4:repeatTensor(1, 1, 1), 1), heatmap_1:repeatTensor(1, 1, 1), 1) | |
| local confidenceMasks = output:narrow(1, 13 + 4 + 1, 4):double() | |
| --[[ | |
| image.save('test/mask_1.png', confidenceMasks[4 * (number - 1) + 1]) | |
| image.save('test/mask_2.png', confidenceMasks[4 * (number - 1) + 2]) | |
| image.save('test/mask_3.png', confidenceMasks[4 * (number - 1) + 3]) | |
| image.save('test/mask_4.png', confidenceMasks[4 * (number - 1) + 4]) | |
| ]]-- | |
| local heatmap_1, heatmap_2, heatmap_3, heatmap_4 = fp_ut.extractRectanglePoints(confidenceMasks[1], confidenceMasks[2], confidenceMasks[3], confidenceMasks[4], numPoints, true) | |
| local iconHeatmaps = torch.cat(torch.cat(torch.cat(heatmap_4:repeatTensor(1, 1, 1), heatmap_3:repeatTensor(1, 1, 1), 1), heatmap_1:repeatTensor(1, 1, 1), 1), heatmap_2:repeatTensor(1, 1, 1), 1) | |
| local segmentations = output:narrow(1, 22, 30):double() | |
| junctionHeatmaps = image.scale(junctionHeatmaps, width, height) | |
| doorHeatmaps = image.scale(doorHeatmaps, width, height) | |
| iconHeatmaps = image.scale(iconHeatmaps, width, height) | |
| segmentations = image.scale(segmentations, width, height) | |
| return junctionHeatmaps, doorHeatmaps, iconHeatmaps, segmentations | |
| end | |
| function utils.proceduralGeneration(floorplan) | |
| local width, height = floorplan:size(3), floorplan:size(2) | |
| local floorplanByte = (floorplan:transpose(1, 2):transpose(2, 3) * 255):byte() | |
| local floorplanGray = cv.cvtColor({floorplanByte, nil, cv.COLOR_BGR2GRAY}) | |
| local edges = cv.Canny({floorplanByte, 100, 200}) | |
| image.save('test/edges.png', edges:double() / 255) | |
| local lines = cv.HoughLinesP({edges, 1, cv.CV_PI / 180, 100}) | |
| local lineImage = floorplanByte:clone() | |
| local lineMask = torch.ByteTensor(#floorplanGray):zero() | |
| lineImage:zero() | |
| for i = 1, lines:size(1) do | |
| local line = lines[i][1] | |
| cv.line({lineImage, {line[1], line[2]}, {line[3], line[4]}, {255, 255, 255}}) | |
| cv.line({lineMask, {line[1], line[2]}, {line[3], line[4]}, {255}}) | |
| end | |
| image.save('test/Hough.png', lineImage:transpose(2, 3):transpose(1, 2):double() / 255) | |
| local lineMask = lineMask:gt(128) | |
| local lineSegments, numLineSegments = utils.findConnectedComponents(lineMask) | |
| image.save('test/line_mask.png', utils.drawSegmentation(lineSegments)) | |
| lineSegments[1]:fill(numLineSegments) | |
| lineSegments[lineSegments:size(1)]:fill(numLineSegments) | |
| lineSegments:narrow(2, 1, 1):fill(numLineSegments) | |
| lineSegments:narrow(2, lineSegments:size(2), 1):fill(numLineSegments) | |
| cv.watershed({floorplanByte, lineSegments}) | |
| local corners = cv.cornerHarris({lineSegments:byte(), nil, 5, 3, 0.04}) | |
| local segmentationImage = utils.drawSegmentation(lineSegments) | |
| --corners = cv.dilate(corners) | |
| --corners = image.dilate(corners:double()) | |
| local cornerMask = corners:gt(corners:max() * 0.01) | |
| local smallSegmentMask = torch.zeros(height, width) | |
| local smallSegmentAreaThreshold = width * height * 0.002 | |
| local smallSegmentWidthThreshold = 3 | |
| for segmentIndex = 1, numLineSegments - 1 do | |
| local segmentMask = lineSegments:eq(segmentIndex) | |
| local indices = segmentMask:nonzero() | |
| if ##indices > 0 then | |
| local mins = torch.min(indices, 1)[1] | |
| local maxs = torch.max(indices, 1)[1] | |
| if (#indices)[1] < smallSegmentAreaThreshold or maxs[1] - mins[1] < smallSegmentWidthThreshold and maxs[2] - mins[2] < smallSegmentWidthThreshold then | |
| smallSegmentMask[segmentMask] = 1 | |
| end | |
| end | |
| end | |
| smallSegmentMask = image.dilate(smallSegmentMask) | |
| --cornerMask[smallSegmentMask:byte()] = 0 | |
| --cornerMask[lineSegments:le(0)] = 0 | |
| segmentationImage[1][cornerMask] = 1 | |
| segmentationImage[2][cornerMask] = 0 | |
| segmentationImage[3][cornerMask] = 0 | |
| image.save('test/segmentation.png', segmentationImage) | |
| print(lines:size(1)) | |
| os.exit(1) | |
| return | |
| end | |
| function utils.writePopupData(width, height, representation, filename) | |
| local floorplanSegmentation, shortWalls = utils.getSegmentation(width, height, representation) | |
| floorplanSegmentation = floorplanSegmentation[1] | |
| local representationFile = io.open(filename .. '.txt', 'w') | |
| representationFile:write(width .. '\t' .. height .. '\n') | |
| representationFile:write(#shortWalls .. '\n') | |
| for _, wall in pairs(shortWalls) do | |
| for pointIndex = 1, 2 do | |
| for c = 1, 2 do | |
| representationFile:write(torch.round(wall[pointIndex][c]) .. '\t') | |
| end | |
| end | |
| local lineDim = utils.lineDim(wall) | |
| local center = {(wall[1][1] + wall[2][1]) / 2, (wall[1][2] + wall[2][2]) / 2} | |
| center[1] = math.max(math.min(center[1], width), 1) | |
| center[2] = math.max(math.min(center[2], height), 1) | |
| local label_1, label_2 | |
| if lineDim == 1 then | |
| for delta = 1, height do | |
| if center[2] - delta < 1 then | |
| break | |
| end | |
| if floorplanSegmentation[center[2] - delta][center[1]] <= 11 then | |
| label_1 = floorplanSegmentation[center[2] - delta][center[1]] | |
| break | |
| end | |
| end | |
| if not label_1 then | |
| label_1 = 11 | |
| end | |
| for delta = 1, height do | |
| if center[2] + delta > height then | |
| break | |
| end | |
| if floorplanSegmentation[center[2] + delta][center[1]] <= 11 then | |
| label_2 = floorplanSegmentation[center[2] + delta][center[1]] | |
| break | |
| end | |
| end | |
| if not label_2 then | |
| label_2 = 11 | |
| end | |
| else | |
| for delta = 1, width do | |
| if center[1] + delta > width then | |
| break | |
| end | |
| if floorplanSegmentation[center[2]][center[1] + delta] <= 11 then | |
| label_1 = floorplanSegmentation[center[2]][center[1] + delta] | |
| break | |
| end | |
| end | |
| if not label_1 then | |
| label_1 = 11 | |
| end | |
| for delta = 1, height do | |
| if center[1] - delta < 1 then | |
| break | |
| end | |
| if floorplanSegmentation[center[2]][center[1] - delta] <= 11 then | |
| label_2 = floorplanSegmentation[center[2]][center[1] - delta] | |
| break | |
| end | |
| end | |
| if not label_2 then | |
| label_2 = 11 | |
| end | |
| end | |
| representationFile:write(label_1 .. '\t' .. label_2 .. '\n') | |
| end | |
| for itemMode, items in pairs(representation) do | |
| if itemMode == 'doors' or itemMode == 'icons' then | |
| for _, item in pairs(items) do | |
| for __, field in pairs(item) do | |
| if __ <= 3 then | |
| for ___, value in pairs(field) do | |
| if __ <= 2 then | |
| value = torch.round(value) | |
| end | |
| representationFile:write(value .. '\t') | |
| end | |
| end | |
| end | |
| representationFile:write('\n') | |
| end | |
| end | |
| end | |
| representationFile:close() | |
| --if floorplan then | |
| --image.save(filename .. '.png', floorplan) | |
| --end | |
| end | |
| return utils | |