Merge pull request #1 from Stochastic13/code-cleanup

Code cleanup

InverseTransformSampling.py

@@ -1,36 +1,43 @@
 import numpy as np
 
-def pdftocdf(x,y,p,shape):
-    xs,ys = (shape[0]/np.max(x.flatten()),shape[1]/np.max(y.flatten()))
-    px = np.sum(p,axis=1)
+
+def pdftocdf(x, y, p, shape):
+    xs, ys = (shape[0] / np.max(x), shape[1] / np.max(y))  # correction for arange not always reaching the end value
+    px = np.sum(p, axis=1)  # px is marginal CDF for x alone
     px = np.cumsum(px)
-    p = np.apply_along_axis(np.cumsum,0,p)
-    p = np.apply_along_axis(np.cumsum,1,p)
-    p = p/np.max(p.flatten())
-    return (p,x*xs,y*ys,px/np.max(px.flatten()))
-
-def samplecdf1(ux,uy,p,x,y,px): # no interpolation
-    a1 = np.argmin(np.abs(px-ux))
-    sel_x = x[0,:][a1]
-    sel_y = y[:,0][np.argmin(np.abs(p[a1,:]/np.max(p[a1,:].flatten())-uy))]
-    return sel_x,sel_y
-
-def transformp(cn,p,x,y,shape):
-    p,x,y,px = pdftocdf(x,y,p,shape)
+    px = px / np.max(px)
+    p = np.apply_along_axis(np.cumsum, 0, p)  # p is the joint CDF for x,y
+    p = np.apply_along_axis(np.cumsum, 1, p)
+    p = p / np.max(p)
+    return (p, x * xs, y * ys, px)
+
+
+def samplecdf1(ux, uy, p, x, y, px):  # no interpolation - density of the p grid determines the smoothness of the output
+    a1 = np.argmin(np.abs(px - ux))  # ux,uy are the uniform random variables for Inverse Transformation Sampling
+    sel_x = x[0, :][a1]  # sample x based on ux first
+    sel_y = y[:, 0][np.argmin(np.abs(p[a1, :] / np.max(p[a1, :]) - uy))]  # p[a1,:] give the marginal CDF for y at sel_x
+    return sel_x, sel_y
+
+
+def transformp(cn, p, x, y, shape):
+    p, x, y, px = pdftocdf(x, y, p, shape)
     ans = []
     for i in range(cn):
-        ans.append(samplecdf1(np.random.rand(),np.random.rand(),p,x,y,px))
+        ans.append(samplecdf1(np.random.rand(), np.random.rand(), p, x, y, px))
     return ans
 
-def gaussian(mx,my,sigmax,sigmay,corr=0,spacing=None,shape=None):
-    mx = mx*shape[0]
-    my = my*shape[1]
+
+def gaussian(mx, my, sigmax, sigmay, corr=0, spacing=None, shape=None):
+    mx = mx * shape[0]  # Find the centres (mx,my are relative distances)
+    my = my * shape[1]
     if spacing is None:
-        x = np.linspace(0,shape[0],1500)
-        y = np.linspace(0,shape[1],1500)
+        x = np.linspace(0, shape[0], 1500)
+        y = np.linspace(0, shape[1], 1500)
     else:
-        x = np.arange(0,shape[0],spacing)
-        y = np.arange(0,shape[1],spacing)
-    x,y = np.meshgrid(x,y)
-    p = 1/(2*np.pi*sigmax*sigmay*np.sqrt(1-corr**2))*np.exp(-1/(2-2*corr**2)*((x-mx)**2/sigmax**2 + (y-my)**2/sigmay**2 - 2*corr*(x-mx)*(y-my)/sigmax/sigmay))
-    return (p,x,y)
+        x = np.arange(0, shape[0], spacing)
+        y = np.arange(0, shape[1], spacing)
+    x, y = np.meshgrid(x, y)
+    p = 1 / (2 * np.pi * sigmax * sigmay * np.sqrt(1 - corr ** 2)) * np.exp(-1 / (2 - 2 * corr ** 2) * (
+                (x - mx) ** 2 / sigmax ** 2 + (y - my) ** 2 / sigmay ** 2 - 2 * corr * (x - mx) * (
+                    y - my) / sigmax / sigmay))  # The multivariate Gaussian distribution - returns a probmap
+    return (p, x, y)

README.md

@@ -1,5 +1,5 @@
 # VoronoiTessellations
-Python3 script to create [Voronoi](https://en.wikipedia.org/wiki/Voronoi_diagram) Tessellations (Mosaic pattern) on images. The script basically does 2D k-means-like clustering of the pixels based on fixed pre-defined cluster centres, and averages the RGB values for each cluster group and assigns all the pixels in the group the average value. Further options allow selectively applying average to only one of the RGB channels or exchanging values of the channels randomly (see [docs](https://github.com/Stochastic13/VoronoiTessellations/blob/master/VorTes%20docs.pdf)).
+Python3 script to create [Voronoi](https://en.wikipedia.org/wiki/Voronoi_diagram) Tessellations (Mosaic pattern) on images. The script basically does 2D k-means-like clustering of the pixels based on fixed pre-defined cluster centres (which can be set to be random), and averages the RGB values for each cluster group and assigns all the pixels in the group the average value. Further options allow selectively applying average to only one of the RGB channels or exchanging values of the channels randomly (see [docs](https://github.com/Stochastic13/VoronoiTessellations/blob/master/VorTes%20docs.pdf)).
 
 
 This is my first ever open-source repository. :)
@@ -29,7 +29,7 @@ You can run the following command to see help:
     positional arguments:
       input                 Input Image file
       output                Output Image file
-      cn                    Number of clusters (default = 0.1*size)
+      cn                    Number of clusters (default = 0.1*smaller_dimension)
 
     optional arguments:
       -h, --help            show this help message and exit
@@ -62,11 +62,15 @@ Source:
 <div align=çenter>
   <img src='demo\default_options_2000.jpg' height=250px>
   </div>
+
  gaussian probmap using the `--gaussian` option with `--gaussianvars 0.3 0.8 90 150`
+
 <div align=çenter>
   <img src='demo\gaussian_3000.jpg' height=250px>
   </div>
+
  `--channel rand` and `--channel randdual`
+
 <div align=çenter>
   <img src='demo\channel_1000.jpg' height=250px>
   <img src='demo\channel2_1000.jpg' height=250px>

VoronoiMain.py

@@ -2,118 +2,126 @@
 import numpy as np
 from tessellate_fast import tessel_fast
 from tessellate_lowmem import tessel_low_mem
-from InverseTransformSampling import transformp,gaussian
+from InverseTransformSampling import transformp, gaussian
 import argparse
 
 # Parse command line arguments
 parser = argparse.ArgumentParser()
-parser.add_argument('input',help='Input Image file')
-parser.add_argument('output',help='Output Image file')
-parser.add_argument('cn',default=0,help='Number of clusters (default = 0.1*size)',type=int)
-parser.add_argument('--rescale',default = 1,help='Rescaling factor for large images',type=float)
-parser.add_argument('--border',default=0,help='Make border [1/0]?',type=int)
-parser.add_argument('--method',default='low_mem',help='fast vs low_mem methods. Default is low_mem.')
-parser.add_argument('--threshold',default=200,help='Only for borders. Threshold distance.',type=float)
-parser.add_argument('--clusmap',default=0,help='Load a specific cluster map as tab-separated text file')
-parser.add_argument('--probmap',default=0,help='Load a 2D probability map for cluster generation')
-parser.add_argument('--channel',default=0,help='Whether to tessellate along only R,G,B or combinations?',choices=['r','g','b','rand','rb','rg','bg','randdual'])
-parser.add_argument('--verbose',default=1,help='Print progress?[1/0]',type=int)
-parser.add_argument('--seed',default='None',help='Seed for PRNG')
-parser.add_argument('--gaussianvars',nargs='*',help='Only for gaussian probmap (mx,my,sigmax,sigmay,corr(opt),spacing(opt))')
+parser.add_argument('input', help='Input Image file')
+parser.add_argument('output', help='Output Image file')
+parser.add_argument('cn', default=0, help='Number of clusters (default = 0.1*smaller_dimension)', type=int)
+parser.add_argument('--rescale', default=1, help='Rescaling factor for large images', type=float)
+parser.add_argument('--border', default=0, help='Make border [1/0]?', type=int)
+parser.add_argument('--method', default='low_mem', help='fast vs low_mem methods. Default is low_mem.')
+parser.add_argument('--threshold', default=200, help='Only for borders. Threshold distance.', type=float)
+parser.add_argument('--clusmap', default=0, help='Load a specific cluster map as tab-separated text file')
+parser.add_argument('--probmap', default=0, help='Load a 2D probability map for cluster generation')
+parser.add_argument('--channel', default=0, help='Whether to tessellate along only R,G,B or combinations?',
+                    choices=['r', 'g', 'b', 'rand', 'rb', 'rg', 'bg', 'randdual'])
+parser.add_argument('--verbose', default=1, help='Print progress?[1/0]', type=int)
+parser.add_argument('--seed', default='None', help='Seed for PRNG')
+parser.add_argument('--gaussianvars', nargs='*',
+                    help='Only for gaussian probmap (mx,my,sigmax,sigmay,corr(opt),spacing(opt))')
 args = parser.parse_args()
 
 # seed
-if not args.seed=='None':
+if not args.seed == 'None':
     np.random.seed(int(args.seed))
 
 # load and rescale input image
 img = Image.open(args.input)
-img = img.resize((int(img.size[0]/args.rescale),int(img.size[1]/args.rescale)))
+img = img.resize((int(img.size[0] / args.rescale), int(img.size[1] / args.rescale)))
 img = np.array(img)
 
 # verbose mode?
-verb = [False,True][args.verbose]
-
+verb = [False, True][args.verbose]
 
 if verb:
-    print('Making Clusters.')
+    print('Making cluster centers.')
+
+# default cn
+if args.cn == 0:
+    args.cn = int(0.1 * np.min(img.shape))
+
 # Cluster generation
-if not (args.clusmap == 0): # Pre-formed cluster-map
+if not (args.clusmap == 0):  # Pre-formed cluster-map is desired
     clusters = []
     with open(args.clusmap) as f:
         for row in f:
-            row = row.split('\n')[0]
-            clusters.append(list(map(float,row.split('\t'))))
+            row = row.split('\n')[0]  # Just for extra protection against bad lines?
+            clusters.append(list(map(float, row.split('\t'))))
     clusters = np.array(clusters)
     args.cn = clusters.shape[0]
-elif not (args.probmap == 0): #Probability distribution for Inverse Transform Sampling
+elif not (args.probmap == 0):  # Probability distribution for Inverse Transform Sampling given
     if args.probmap == 'gaussian':
-        if len(args.gaussianvars)<6:
-            defs = [0.5,0.5,100,100,0,None]
-            args.gaussianvars = list(map(float,args.gaussianvars)) + defs[len(args.gaussianvars):len(defs)]
-        arguments = args.gaussianvars+[img.shape]
+        if len(args.gaussianvars) < 6:
+            defs = [0.5, 0.5, 100, 100, 0, None]  # defaults
+            args.gaussianvars = list(map(float, args.gaussianvars)) + defs[len(args.gaussianvars):len(defs)]
+        arguments = args.gaussianvars + [img.shape]
         g = gaussian(*arguments)
-        clusters = transformp(args.cn,g[0],g[1],g[2],img.shape)
+        clusters = transformp(args.cn, g[0], g[1], g[2], img.shape)
         clusters = np.array(tuple(clusters))
-else:
-    clusters = np.array(tuple(zip(np.random.rand(args.cn)*img.shape[0],np.random.rand(args.cn)*img.shape[1])))
+else:  # random cluster map
+    clusters = np.array(tuple(zip(np.random.rand(args.cn) * img.shape[0], np.random.rand(args.cn) * img.shape[1])))
 
 if verb:
-    print('Done.')
+    print('Cluster centers are ready.')
     print('Making Voronoi Tessellations....')
 
 # Tessellating
-if args.method=='fast':
-    dist = tessel_fast(clusters,img.shape,[False,True][args.verbose],[False,True][args.border],args.threshold)
-elif args.method=='low_mem':
-    dist = tessel_low_mem(clusters,img.shape,[False,True][args.verbose],[False,True][args.border],args.threshold)
+if args.method == 'fast':  # dist is the cluster membership array
+    dist = tessel_fast(clusters, img.shape, [False, True][args.verbose], [False, True][args.border], args.threshold)
+elif args.method == 'low_mem':
+    dist = tessel_low_mem(clusters, img.shape, [False, True][args.verbose], [False, True][args.border], args.threshold)
 else:
     print("ERROR: Invalid Method")
-    quit(1)
+    quit(5)
 
 if verb:
     print('\t\t\t\t\t\nDone.')
-# Averaging over Voronoi clusters
+# Averaging pixels over the clusters
 s = set(dist.flatten())
 sl = len(s)
-x=0 # counter for Verbose mode
+x = 0  # counter for verbose mode
 if verb:
     print('\nAveraging over Voronoi clusters.')
-for i in (set(list(range(args.cn))) & s):
+for i in (set(list(range(args.cn))) & s):  # To exclude centers without any membership
     if verb:
-        print(str(int(x/sl*100))+'% done \t\t\r',end='')
+        print(str(int(x / sl * 100)) + '% done \t\t\r', end='')
         x += 1
-    if args.channel in ['r','g','b']:
-        chn = {'r':0,'g':1,'b':2}[args.channel]
+    if args.channel in ['r', 'g', 'b']:  # If averaging only along 1 channel
+        chn = {'r': 0, 'g': 1, 'b': 2}[args.channel]
         img[dist == i, chn] = int(np.mean(img[dist == i, chn].flatten()))
         continue
-    elif args.channel=='rand':
-        chn = np.random.randint(0,3)
+    elif args.channel == 'rand':  # Randomly select a channel and average
+        chn = np.random.randint(0, 3)
         img[dist == i, chn] = int(np.mean(img[dist == i, chn].flatten()))
         continue
-    elif args.channel=='randdual':
-        chn1 = np.random.randint(0,3)
-        chn2 = np.random.randint(0,3)
-        img[dist == i, chn1],img[dist == i, chn2] = (np.array(img[dist == i, chn2],copy=True),np.array(img[dist == i, chn1],copy=True))
+    elif args.channel == 'randdual':  # Randomly exchange the channels
+        chn1 = np.random.randint(0, 3)
+        chn2 = np.random.randint(0, 3)
+        img[dist == i, chn1], img[dist == i, chn2] = (
+            np.array(img[dist == i, chn2], copy=True), np.array(img[dist == i, chn1], copy=True))
         continue
-    elif args.channel in ['rb','rg','gb']:
-        perms = ['rb','rg','gb']
-        chn1,chn2 = [(0,2),(0,1),(1,2)][perms.index(args.channel)]
-        chn1,chn2 = [(chn1,chn2),(0,0)][np.random.randint(0,2)]
-        if chn1==chn2:
+    elif args.channel in ['rb', 'rg', 'gb']:  # Exchange two channels (specific)
+        perms = ['rb', 'rg', 'gb']
+        chn1, chn2 = [(0, 2), (0, 1), (1, 2)][perms.index(args.channel)]
+        chn1, chn2 = [(chn1, chn2), (0, 0)][np.random.randint(0, 2)]
+        if chn1 == chn2:
             continue
-        img[dist == i, chn1],img[dist == i, chn2] = (np.array(img[dist == i, chn2],copy=True),np.array(img[dist == i, chn1],copy=True))
+        img[dist == i, chn1], img[dist == i, chn2] = (
+            np.array(img[dist == i, chn2], copy=True), np.array(img[dist == i, chn1], copy=True))
         continue
-    img[dist==i,0]=int(np.mean(img[dist==i,0].flatten()))
-    img[dist==i,1]=int(np.mean(img[dist==i,1].flatten()))
-    img[dist==i,2]=int(np.mean(img[dist==i,2].flatten()))
-if [False,True][args.border]:
-    img[dist==args.cn+1,0]=0
-    img[dist==args.cn+1,1]=0
-    img[dist==args.cn+1,2]=0
+    img[dist == i, 0] = int(np.mean(img[dist == i, 0]))  # The vanilla averaging
+    img[dist == i, 1] = int(np.mean(img[dist == i, 1]))
+    img[dist == i, 2] = int(np.mean(img[dist == i, 2]))
+if [False, True][args.border]:
+    img[dist == args.cn + 1, 0] = 0  # If drawing borders is set to True
+    img[dist == args.cn + 1, 1] = 0
+    img[dist == args.cn + 1, 2] = 0
 
 if verb:
     print('\t\t\t\t\t\nDone.')
-    print('Saving Output file.')
+    print('Saving output file.')
 img = Image.fromarray(img)
 img.save(args.output)