basically working

Author: kwinkunks

docs/index.rst

@@ -21,6 +21,7 @@ User guide
 
     readme
     userguide/Unmap_data_from_an_image.ipynb
+    userguide/Guess_the_colourmap_from_an_image.ipynb
 
 
 API reference

docs/notebooks/Guess_the_colourmap_from_an_image.ipynb

@@ -16,7 +16,8 @@
     "- **You don't have a lot of pixels.** Small images are hard to rip data from. We'd always like more pixels.\n",
     "- **There are a lot of annotations.** Usually these get in the way of the data.\n",
     "- **The image is lossily compressed.** Most PDFs contain JPEGs and JPEG is [lossy](https://en.wikipedia.org/wiki/Lossy_compression). This means the colours are a bit garbled, especially around abrupt edges (like annotations!).\n",
-    "- **The image has hillshading.** The cute shadow effect adds another layer of complexity... but we can still have a go. \n",
+    "- **The image has dithering.** If the number of colours in the image has been reduced at some point, there's a good chance it has been [dithered](https://en.wikipedia.org/wiki/Dither).\n",
+    "- **The image has hillshading or specularity.** Cute 3D effects add another layer of complexity... but we can still have a go. \n",
     "- **It's a 3D perspective plot.** You might be able to recover the data from what you can _see_, but 3D perspective views add another type of distortion that I daresay could be undone, but not by me. \n",
     "- **The image is poor.** If it's a photo or scan of a paper document... well, now you 've got the unknown transform from the data to the image, then the unknown transform of the image to the physical plot, then the unknown transform of the plot to the image you have. I mean, come on.\n",
     "\n",

docs/notebooks/Unmap_data_from_an_image.ipynb

@@ -14,7 +14,7 @@
 from scipy.spatial import cKDTree
 from scipy.cluster.vq import kmeans
 from matplotlib import cm
-from matplotlib.colors import hsv_to_rgb, rgb_to_hsv
+from matplotlib.colors import ListedColormap, hsv_to_rgb, rgb_to_hsv
 from matplotlib.colors import to_rgb, LinearSegmentedColormap
 
 
@@ -147,6 +147,9 @@ def get_cmap(cmap, arr=None, levels=256, quantize=False):
         except ValueError:
             raise ValueError("cmap name not recognized by matplotlib")
 
+    elif isinstance(cmap, LinearSegmentedColormap) or isinstance(cmap, ListedColormap):
+        return cmap
+
     else:
         try:
             cmap = np.array(cmap)

unmap/unmap.py

@@ -65,7 +65,7 @@ def construct_graph(imarray, colors=256, normed=True):
 
     # Normalize.
     if normed:
-        glcm /= (1e-9 + np.sum(glcm, axis=-1))
+        glcm = glcm / (1e-9 + np.sum(glcm, axis=-1))
 
     # Construct and remove self-loops.
     G = nx.from_numpy_array(glcm)
@@ -164,7 +164,7 @@ def longest_shortest_path(G):
         path (list): Longest shortest path.
     """
 
-    dist = lambda *_, d: d['dist']**2
+    dist = lambda u, v, d: d['dist']**2
 
     # Find the longest shortest path.
     paths = nx.shortest_path_length(G, weight=dist)
@@ -181,7 +181,7 @@ def longest_shortest_path(G):
                             target=t)
 
 
-def path_to_cmap(path, unique_colors, colors=256, reverse='auto'):
+def path_to_cmap(path, unique_colors, colors=256, reverse='auto', equilibrate=False):
     """
     Convert a path through the graph to a colormap.
 
@@ -193,11 +193,14 @@ def path_to_cmap(path, unique_colors, colors=256, reverse='auto'):
         reverse (bool): Whether to reverse the colormap. If 'auto', the
             colormap will start with the end closest to dark blue. If False,
             the direction is essentially random.
+        equilibrate (bool): Whether to equilibrate the colormap. This will
+            try to ensure that the colormap's colors are as evenly spaced as
+            possible.
 
     Returns:
         matplotlib.colors.LinearSegmentedColormap: Colormap.
     """
-    cpath = [unique_colors[n] for n in path]
+    cpath = np.asarray(unique_colors)[path]
     if reverse == 'auto':
         cool_dark = np.array([0, 0, 0.5])
         if np.linalg.norm(cpath[0] - cool_dark) > np.linalg.norm(cpath[-1] - cool_dark):
@@ -208,10 +211,23 @@ def path_to_cmap(path, unique_colors, colors=256, reverse='auto'):
         cpath = cpath[::-1]
     if colors is None:
         colors = 2 * len(cpath)  # Not sure what the default should be.
-    return LinearSegmentedColormap.from_list("recovered", cpath, N=colors)
-
-
-def guess_cmap(fname, source_colors=256, target_colors=256, min_weight=0.025, max_dist=0.25, max_neighbours=20, reverse='auto', ):
+    cmap = LinearSegmentedColormap.from_list("recovered", cpath, N=colors)
+    if equilibrate:
+        dists = np.linalg.norm(cpath[:-1] - cpath[1:], axis=-1)
+        invdist = np.cumsum(1 / dists) / (1 / dists).sum()
+        cmap = LinearSegmentedColormap.from_list('recovered', cmap(invdist), N=colors)
+    return cmap
+
+
+def guess_cmap(fname,
+               source_colors=256,
+               target_colors=256,
+               min_weight=0.025,
+               max_dist=0.25,
+               max_neighbours=20,
+               reverse='auto',
+               equilibrate=False
+               ):
     """
     Guess the colormap of an image.
 
@@ -231,4 +247,4 @@ def guess_cmap(fname, source_colors=256, target_colors=256, min_weight=0.025, ma
     G = construct_graph(imarray, colors=source_colors)
     G0 = prune_graph(G, uniq, min_weight=min_weight, max_dist=max_dist, max_neighbours=max_neighbours)
     path = longest_shortest_path(G0)
-    return path_to_cmap(path, uniq, colors=target_colors, reverse=reverse)
+    return path_to_cmap(path, uniq, colors=target_colors, reverse=reverse, equilibrate=equilibrate)