require 'torch'
require 'lfs'
cv = require 'cv'
require 'cv.imgproc'

local utils = {}

function utils.to4DTensor(a)
  return utils.toXDTensor(a, 4)
end

function utils.toXDTensor(a, x)
  local sz = torch.LongStorage(x):fill(1)
  for i = 1, a:dim() do
    sz[i] = a:size(i)
  end
  return a:reshape(sz)
end

function utils.toNNSingleTensor(a)
  local opt = opt or {}
  opt.type = opt.type or 'CUDA'

  if opt.type == 'CPU' then 
    return a:double()
  elseif opt.type == 'CUDA' then 
    return a:cuda()
  else
    return a
  end
end

function utils.toNNTensor(a)
  if type(a) == 'userdata' then
    return utils.toNNSingleTensor(a)
  elseif type(a) == 'table' then
    local b = {}
    for key, value in pairs(a) do
      b[key] = utils.toNNTensor(value)
    end
    return b
  else
    return a
  end
end

function utils.getLength(a)
   if type(a) == 'userdata' then
    return a:size(1)
  elseif type(a) == 'table' then
    return #a
  end
end

function utils.getFirstDim(a)
  if type(a) == 'userdata' then
    return a:size(1)
  elseif type(a) == 'table' then
    return utils.getSize(a[1])
  end
end

-- to be replaced
function utils.getSize(a)
  return utils.getFirstDim(a)
end


function utils.getIndices(a, indices)
  if type(a) == 'userdata' then
    return a:index(1, indices)
  elseif type(a) == 'table' then
    local b = {} 
    for key, value in pairs(a) do
      b[key] = utils.getIndices(value, indices)
    end
    return b
  end
end

function utils.loadData(fileName)
  if lfs.attributes(fileName, 'mode') == 'file' then
    return torch.load(fileName)
  else
    return nil
  end
end

-- all inputs are ByteTensors, 0 is false and 1 is true.
function utils.isinf(value)
  return value == math.huge or value == -math.huge
end

function utils.isnan(value)
  return value ~= value
end

function utils.isfinite(value)
  return not utils.isinf(value) and not utils.isnan(value)
end

function utils.bitand(t1, t2)
  return torch.ge(t1 + t2, 2)
end

function utils.bitor(t1, t2)
  return torch.ge(t1 + t2, 1)
end

function utils.bitnot(t)
  return torch.lt(t, 1)
end

function utils.len(t)
  local n
  if type(t) == 'table' then
    n = #t
  elseif t:nDimension() == 1 then
    n = t:size(1)
  else
    error(string.format("t is a high-order tensor! t:nDimension() = %d", t:nDimension()))
  end
  return n
end

function utils.fromHead(t, k)
  local n = utils.len(t)
  local res = {}
  for i = 1, math.min(k, n) do
    table.insert(res, t[i])
  end
  return res
end

function utils.fromTail(t, k)
  local n = utils.len(t)
  local res = {}
  for i = n, math.max(n - k, 0) + 1, -1 do
    table.insert(res, t[i])
  end
  return res
end

-- Find nonzero and return a table.
function utils.nonzero(t)
  local indices = {}
  local n = utils.len(t)

  for i = 1, n do
    if t[i] == 1 then 
      table.insert(indices, i) 
    end
  end
  return indices
end

-- Select rows and return
function utils.selectRows(t, tb)
  assert(tb:size(1) == t:size(1), 'The first dimension of t and tb must be the same')
  local dims = t:size():totable()
  table.remove(dims, 1)

  local res = t.new():resize(tb:sum(), unpack(dims))
  local counter = 0
  for i = 1, t:size(1) do
    if tb[i] == 1 then 
      counter = counter + 1
      res[counter]:copy(t[i]) 
    end
  end
  return res
end

function utils.permCompose(indices1, indices2)
  -- body
  local n1 = utils.len(indices1)
  local n2 = utils.len(indices2)

  local indices = {}
  for i = 1, n2 do
    indices[i] = indices1[indices2[i]]
  end
  return indices
end

function utils.removeKeys(t, exclude)
  local res = {}
  for k, v in pairs(t) do
    if not exclude[k] then res[k] = v end
  end
  return res
end

--[[ 

Find the correlation between two matrices.
t1 : m1 by n
t2 : m2 by n

return matrix of size m1 by m2

--]]
function utils.innerprod(t1, t2, func, reduction)
  local n = t1:size(2)
  assert(n == t2:size(2), string.format('The column of t1 [%d] must be the same as the column of t2 [%d]', t1:size(2), t2:size(2)))
  local m1 = t1:size(1)
  local m2 = t2:size(1)

  local res = torch.Tensor(m1, m2)
  for i = 1, m1 do
    for j = 1, m2 do
      res[i][j] = func(t1[i], t2[j])
    end
  end

  -- Find the one with smallest distance
  local best, bestIndices
  if reduction then
    best, bestIndices = reduction(res, 2)
  end

  return res, best, bestIndices
end

-----------------------------------------------
-- Landmarks related.

function utils.inRect(rect, p, margin)
  margin = margin or 0;

  return rect[1] + margin <= p[1] and p[1] <= rect[3] - margin and
  rect[2] + margin <= p[2] and p[2] <= rect[4] - margin;
end

function utils.fillKernel(m, x, y, r, c)
  assert(m, "Input image should not be null")
  assert(x and y, "Input coordinates should not be null")
  assert(r and c, "Input radius and color should not be null")
  assert(#m:size() == 2, 'fill_kernel: input m is not 2D!')

  -- fill a circle (x, y, r) with number c.
  local w = m:size(2)
  local h = m:size(1)

  local minX = math.min(math.max(x - r, 1), w)
  local maxX = math.min(math.max(x + r, 1), w)
  -- if minX > maxX then minX, maxX = maxX, minX end

  local minY = math.min(math.max(y - r, 1), h)
  local maxY = math.min(math.max(y + r, 1), h)
  -- if minY > maxY then minY, maxY = maxY, minY end

  -- local img_rect = {1, 1, m:size(2), m:size(1)}
  -- assert(util.rect_isin(img_rect, {minX, minY}) == true, string.format("out of bound, minX = %f, minY = %f", minX, minY))
  -- assert(util.rect_isin(img_rect, {maxX, maxY}) == true, string.format("out of bound, maxX = %f, maxY = %f", maxX, maxY))

  m:sub(minY, maxY, minX, maxX):fill(c);
end

-- Input a 2D mask, find its minimal value and associated location.
function utils.imin(m)
  local min1, minI1 = torch.min(m, 1)
  local minVal, minI2 = torch.min(min1, 2)

  local x = minI2[1][1]
  local y = minI1[1][x]
  return x, y, minVal
end

-- Input a 2D mask, find its maximal value and associated location.
function utils.imax(m)
  local max1, maxI1 = torch.max(m, 1)
  local maxVal, maxI2 = torch.max(max1, 2)

  local x = maxI2[1][1]
  local y = maxI1[1][x]
  return x, y, maxVal
end
-------------------------------------Save to json--------------------
function utils.set2Array(t)
  if type(t) ~= 'table' then return end
  t.__array = true
  for i, v in ipairs(t) do
    utils.set2Array(v)
  end
end

function utils.convert2Table(t)
  local res = {}
  if type(t) == 'table' then
    if debug then print("parsing table") end
    for k, v in pairs(t) do
      res[k] = utils.convert2Table(v)
    end
  elseif type(t) == 'number' then
    if debug then print("parsing number") end
    res = t
  elseif torch.typename(t) and torch.typename(t):match('Tensor') then
    -- if t is a tensor
    if debug then print("parsing tensor") end
    res = t:totable()
    utils.set2Array(t)
    -- Layer 
  else
    local typename = type(t)
    typename = typename or torch.typename(t)
    error("Convert_to_table error, unsupported datatype = " .. typename)
  end
  return res
end

function utils.saveJson(t, f)
  -- save a table to json
  if type(t) == 'number' then
    f:write(tostring(t))
    return
  end
  if t.__array then
    -- array must contain all numbers.
    f:write("[\n")
    for i, v in ipairs(t) do
      utils.saveJson(v, f)
      if i ~= #t then f:write(",") end
    end
    f:write("]\n")
  else
    local counter = 0
    for k, v in pairs(t) do counter = counter + 1 end
    f:write("{\n")
    for k, v in pairs(t) do
      f:write(tostring(k) .. " : ")
      utils.saveJson(v, f)
      counter = counter - 1
      if counter >= 1 then f:write(",\n") end
    end
    f:write("}\n")
  end
end

-------------------------------------Save to numpy-----------------
-- run it using PATH=/usr/bin/python

function utils.savePickle(f, t)
  local py = require 'fb.python'
  py.exec([=[
import numpy as np
import cPickle
with open(filename, "wb") as outfile:
    cPickle.dump(variable, outfile, protocol=cPickle.HIGHEST_PROTOCOL)
]=], {variable = t, filename = f})
end

--

local function getTermLength()
  if sys.uname() == 'windows' then return 80 end
  local tputf = io.popen('tput cols', 'r')
  local w = tonumber(tputf:read('*a'))
  local rc = {tputf:close()}
  if rc[3] == 0 then return w
  else return 80 end 
end

local barDone = true
local previous = -1
local tm = ''
local timer
local times
local indices
local termLength = math.min(getTermLength(), 120)

local function formatTime(seconds)
  -- decompose:
  local floor = math.floor
  local days = floor(seconds / 3600/24)
  seconds = seconds - days*3600*24
  local hours = floor(seconds / 3600)
  seconds = seconds - hours*3600
  local minutes = floor(seconds / 60)
  seconds = seconds - minutes*60
  local secondsf = floor(seconds)
  seconds = seconds - secondsf
  local millis = floor(seconds*1000)

  -- string
  local f = ''
  local i = 1
  if days > 0 then f = f .. days .. 'D' i=i+1 end
  if hours > 0 and i <= 2 then f = f .. hours .. 'h' i=i+1 end
  if minutes > 0 and i <= 2 then f = f .. minutes .. 'm' i=i+1 end
  if secondsf > 0 and i <= 2 then f = f .. secondsf .. 's' i=i+1 end
  if millis > 0 and i <= 2 then f = f .. millis .. 'ms' i=i+1 end
  if f == '' then f = '0ms' end

  -- return formatted time
  return f
end

function utils.progress(current, goal, addinfo)
  -- defaults:
  local barLength = termLength - 37 - #addinfo
  local smoothing = 100 
  local maxfps = 10

  -- Compute percentage
  local percent = math.floor(((current) * barLength) / goal)

  -- start new bar
  if (barDone and ((previous == -1) or (percent < previous))) then
    barDone = false
    previous = -1
    tm = ''
    timer = torch.Timer()
    times = {timer:time().real}
    indices = {current}
  else
    io.write('\r')
  end

  --if (percent ~= previous and not barDone) then
  if (not barDone) then
    previous = percent
    -- print bar
    io.write(' [')
    for i=1,barLength do
      if (i < percent) then io.write('=')
      elseif (i == percent) then io.write('>')
      else io.write('.') end
    end
    io.write('] ')
    for i=1,termLength-barLength-4 do io.write(' ') end
    for i=1,termLength-barLength-4 do io.write('\b') end
    -- time stats
    local elapsed = timer:time().real
    local step = (elapsed-times[1]) / (current-indices[1])
    if current==indices[1] then step = 0 end
    local remaining = math.max(0,(goal - current)*step)
    table.insert(indices, current)
    table.insert(times, elapsed)
    if #indices > smoothing then
      indices = table.splice(indices)
      times = table.splice(times)
    end
    -- Print remaining time when running or total time when done.
    if (percent < barLength) then
      tm = ' ETA: ' .. formatTime(remaining)
    else
      tm = ' Tot: ' .. formatTime(elapsed)
    end
    tm = tm .. ' | Step: ' .. formatTime(step) .. ' | ' .. addinfo
    io.write(tm)
    -- go back to center of bar, and print progress
    for i=1,5+#tm+barLength/2 do io.write('\b') end
    io.write(' ', current, '/', goal, ' ')
    -- reset for next bar
    if (percent == barLength) then
      barDone = true
      io.write('\n')
    end
    -- flush
    io.write('\r')
    io.flush()
  end
end

function utils.drawHeatMapImage(image, heatMap, alpha, beta)
   local heatMapImage = image * alpha
   local scale = math.ceil(math.max((#image)[2] / (#heatMap)[1], (#image)[3] / (#heatMap)[2]))
   resizedHeatMap = require('nn').SpatialUpSamplingBilinear(scale):forward(heatMap:repeatTensor(1, 1, 1))[1]
   
   local paddingX = (#image)[2] - (#resizedHeatMap)[1]   
   local paddingLeft = math.floor(paddingX / 2)   
   local paddingRight = paddingX - paddingLeft
   local paddingY = (#image)[3] - (#resizedHeatMap)[2]   
   local paddingTop = math.floor(paddingY / 2)
   local paddingBottom = paddingY - paddingTop
   resizedHeatMap = require('nn').SpatialZeroPadding(paddingLeft, paddingRight, paddingTop, paddingBottom):forward(resizedHeatMap:repeatTensor(1, 1, 1))[1]
   --print(scale)
   --print(padding)
   --print(#resizedHeatMap)
   --os.exit(1)
   
   local heatMapHSL = torch.DoubleTensor(3, (#resizedHeatMap)[1], (#resizedHeatMap)[2])
   heatMapHSL[1] = 2 / 3 * (1 - resizedHeatMap)
   heatMapHSL[2] = 1
   heatMapHSL[3] = 0.5
   local heatMapRGB = require('image').hsl2rgb(heatMapHSL)
   heatMapImage = heatMapImage + heatMapRGB * beta
   heatMapImage:clamp(0, 1)
   return heatMapImage
end



function utils.segmentFloorplan(floorplan, binaryThreshold, numOpenOperations, reverse)

   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 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.drawSegmentation(floorplanComponent, numComponents)
   local colorMap = {}
   for i = 1, numComponents do
      colorMap[i] = torch.rand(3)
   end
   colorMap[0] = torch.zeros(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 utils.parseFloorplan(floorplanInput)
   local binaryThreshold = 0.7
   local floorplanSegmentation, numSegments, floorplanBinary = utils.segmentFloorplan(floorplanInput, binaryThreshold, 0)
   while (#floorplanSegmentation:eq(floorplanSegmentation[1][1]):nonzero())[1] > (#floorplanSegmentation)[1] * (#floorplanSegmentation)[2] * 0.8 and binaryThreshold < 1 do
      binaryThreshold = binaryThreshold + 0.1
      floorplanSegmentation, numSegments, floorplanBinary = utils.segmentFloorplan(floorplanInput, binaryThreshold, 0)
   end
   
   --image.save('test/floorplan.png', floorplanInput)
   --image.save('test/floorplan_segmentation.png', ut.drawSegmentation(floorplanSegmentation, numSegments))
   --image.save('test/floorplan_binary.png', floorplanBinary)
   local floorplanSegmentationReversed, numSegmentsReversed, floorplanBinaryReversed = utils.segmentFloorplan(floorplanInput, binaryThreshold, 0, true)
   --image.save('test/floorplan_binary_reversed.png', floorplanBinaryReversed)
   --image.save('test/floorplan.png', floorplanInput)
   --image.save('test/floorplan_segmentation_reversed.png', utils.drawSegmentation(floorplanSegmentationReversed, numSegmentsReversed))
   local backgroundSegmentReversed = floorplanSegmentationReversed[1][1]
   --print(torch.min(floorplanSegmentationReversed))
   --print(torch.max(floorplanSegmentationReversed))
   --print(numSegmentsReversed)
   local maxCount = 0
   local borderSegmentReversed = 0
   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 == 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)
   local borderThickness = 10
   borderMask = image.dilate(borderMask)
   local floorplanSegmentationErodedOnce = torch.cmul(floorplanSegmentation, (1 - borderMask):int())
   
   --[[
      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

return utils