git.sesse.net Git - voxel-flow/blob - utils/vis_utils.py

   1 """Implements various visualization utils."""
   2
   3 import numpy as np
   4
   5
   6 def viz_flow(u,v,logscale=True,scaledown=6,output=False):
   7   """Flow visualization implentations.
   8   The code obtained from PCA flow github repo.
   9   https://github.com/jswulff/pcaflow/blob/master/pcaflow/utils/viz_flow.py
  10
  11   topleft is zero, u is horiz, v is vertical
  12   red is 3 o'clock, yellow is 6, light blue is 9, blue/purple is 12
  13   """
  14   colorwheel = makecolorwheel()
  15   ncols = colorwheel.shape[0]
  16
  17   radius = np.sqrt(u**2 + v**2)
  18   if output:
  19       print("Maximum flow magnitude: %04f" % np.max(radius))
  20   if logscale:
  21       radius = np.log(radius + 1)
  22       if output:
  23           print("Maximum flow magnitude (after log): %0.4f" % np.max(radius))
  24   radius = radius / scaledown
  25   if output:
  26       print("Maximum flow magnitude (after scaledown): %0.4f" % np.max(radius))
  27   rot = np.arctan2(-v, -u) / np.pi
  28
  29   fk = (rot+1)/2 * (ncols-1)  # -1~1 maped to 0~ncols
  30   k0 = fk.astype(np.uint8)       # 0, 1, 2, ..., ncols
  31
  32   k1 = k0+1
  33   k1[k1 == ncols] = 0
  34
  35   f = fk - k0
  36
  37   ncolors = colorwheel.shape[1]
  38   img = np.zeros(u.shape+(ncolors,))
  39   for i in range(ncolors):
  40       tmp = colorwheel[:,i]
  41       col0 = tmp[k0]
  42       col1 = tmp[k1]
  43       col = (1-f)*col0 + f*col1
  44
  45       idx = radius <= 1
  46       # increase saturation with radius
  47       col[idx] = 1 - radius[idx]*(1-col[idx])
  48       # out of range
  49       col[~idx] *= 0.75
  50       img[:,:,i] = np.floor(255*col).astype(np.uint8)
  51   return img.astype(np.uint8)
  52
  53 def makecolorwheel():
  54         """Flow visualization implentations.
  55   The code obtained from PCA flow github repo.
  56   https://github.com/jswulff/pcaflow/blob/master/pcaflow/utils/viz_flow.py
  57   """
  58   # Create a colorwheel for visualization
  59   RY = 15
  60   YG = 6
  61   GC = 4
  62   CB = 11
  63   BM = 13
  64   MR = 6
  65
  66   ncols = RY + YG + GC + CB + BM + MR
  67
  68   colorwheel = np.zeros((ncols,3))
  69
  70   col = 0
  71   # RY
  72   colorwheel[0:RY,0] = 1
  73   colorwheel[0:RY,1] = np.arange(0,1,1./RY)
  74   col += RY
  75
  76   # YG
  77   colorwheel[col:col+YG,0] = np.arange(1,0,-1./YG)
  78   colorwheel[col:col+YG,1] = 1
  79   col += YG
  80
  81   # GC
  82   colorwheel[col:col+GC,1] = 1
  83   colorwheel[col:col+GC,2] = np.arange(0,1,1./GC)
  84   col += GC
  85
  86   # CB
  87   colorwheel[col:col+CB,1] = np.arange(1,0,-1./CB)
  88   colorwheel[col:col+CB,2] = 1
  89   col += CB
  90
  91   # BM
  92   colorwheel[col:col+BM,2] = 1
  93   colorwheel[col:col+BM,0] = np.arange(0,1,1./BM)
  94   col += BM
  95
  96   # MR
  97   colorwheel[col:col+MR,2] = np.arange(1,0,-1./MR)
  98   colorwheel[col:col+MR,0] = 1
  99
 100   return colorwheel