Update main.py

main.py

@@ -8,7 +8,30 @@ import huggingface_hub # for loading model
 model1 = huggingface_hub.from_pretrained_keras("cmudrc/microstructure-colorization")
 
 # Get the color map by name:
-cm = matplotlib.pyplot.get_cmap('RdBu')
+# cm = matplotlib.pyplot.get_cmap('RdBu')
+# import matplotlib.colors
+
+cdict = {'red':   [[0.00, 0.10, 0.10],
+                   [0.17, 0.30, 0.30],
+                   [0.33, 0.35, 0.35],
+                   [0.50, 1.00, 1.00],
+                   [0.75, 0.92, 0.92],
+                   [1.00, 0.87, 0.87]],
+         'green': [[0.00, 0.30, 0.30],
+                   [0.17, 0.53, 0.53],
+                   [0.33, 0.72, 0.72],
+                   [0.50, 1.00, 1.00],
+                   [0.75, 0.52, 0.52],
+                   [1.00, 0.32, 0.32]],
+         'blue':  [[0.00, 0.29, 0.29],
+                   [0.17, 0.53, 0.53],
+                   [0.33, 0.73, 0.73],
+                   [0.50, 1.00, 1.00],
+                   [0.75, 0.20, 0.20],
+                   [1.00, 0.15, 0.15]]}
+
+
+hamster_colormap = matplotlib.colors.LinearSegmentedColormap('hamster', segmentdata=cdict, N=256)
 
 #simple image scaling to (nR x nC) size
 def scale(im, nR, nC):