util/lua/utils.lua

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 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)
   --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)
   local heatMapImage = image * alpha + 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, denotedColorMap)
   local colorMap = denotedColorMap
   if colorMap == nil then
      colorMap = {}
      for i = 1, numComponents do
         colorMap[i] = torch.rand(3)
      end
      colorMap[0] = torch.zeros(3)
      colorMap[-1] = torch.ones(3)
   end

   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

return utils