imagepers.py

   1 #!/usr/bin/python2
   2
   3 """A simple implementation of persistent homology on 2D images."""
   4
   5 __author__ = "Stefan Huber <shuber@sthu.org>"
   6
   7
   8 import union_find
   9
  10
  11 def get(im, p):
  12     return im[p[0]][p[1]]
  13
  14
  15 def iter_neighbors(p, w, h):
  16     y, x = p
  17
  18     # 8-neighborship
  19     neigh = [(y+j, x+i) for i in [-1, 0, 1] for j in [-1, 0, 1]]
  20     # 4-neighborship
  21     # neigh = [(y-1, x), (y+1, x), (y, x-1), (y, x+1)]
  22
  23     for j, i in neigh:
  24         if j < 0 or j >= h:
  25             continue
  26         if i < 0 or i >= w:
  27             continue
  28         if j == y and i == x:
  29             continue
  30         yield j, i
  31
  32
  33 def persistence(im):
  34     h, w = im.shape
  35
  36     # Get indices orderd by value from high to low
  37     indices = [(i, j) for i in range(h) for j in range(w)]
  38     indices.sort(key=lambda p: get(im, p), reverse=True)
  39
  40     # Maintains the growing sets
  41     uf = union_find.UnionFind()
  42
  43     groups0 = {}
  44
  45     def get_comp_birth(p):
  46         return get(im, uf[p])
  47
  48     # Process pixels from high to low
  49     for i, p in enumerate(indices):
  50         v = get(im, p)
  51         ni = [uf[q] for q in iter_neighbors(p, w, h) if q in uf]
  52         nc = sorted([(get_comp_birth(q), q) for q in set(ni)], reverse=True)
  53
  54         if i == 0:
  55             groups0[p] = (v, v, None)
  56
  57         uf.add(p, -i)
  58
  59         if len(nc) > 0:
  60             oldp = nc[0][1]
  61             uf.union(oldp, p)
  62
  63             # Merge all others with oldp
  64             for bl, q in nc[1:]:
  65                 if uf[q] not in groups0:
  66                     #print(i, ": Merge", uf[q], "with", oldp, "via", p)
  67                     groups0[uf[q]] = (bl, bl-v, p)
  68                 uf.union(oldp, q)
  69
  70     groups0 = [(k, groups0[k][0], groups0[k][1], groups0[k][2]) for k in groups0]
  71     groups0.sort(key=lambda g: g[2], reverse=True)
  72
  73     return groups0