add color map

app.py

@@ -38,6 +38,34 @@ import matplotlib.pyplot as plt
 
 from skimage import exposure
 
+cdl_color_map = [{'value': 1, 'label': 'Natural vegetation', 'rgb': (233,255,190)},
+                 {'value': 2, 'label': 'Forest', 'rgb': (149,206,147)},
+                 {'value': 3, 'label': 'Corn', 'rgb': (255,212,0)},
+                 {'value': 4, 'label': 'Soybeans', 'rgb': (38,115,0)},
+                 {'value': 5, 'label': 'Wetlands', 'rgb': (128,179,179)},
+                 {'value': 6, 'label': 'Developed/Barren', 'rgb': (156,156,156)},
+                 {'value': 7, 'label': 'Open Water', 'rgb': (77,112,163)},
+                 {'value': 8, 'label': 'Winter Wheat', 'rgb': (168,112,0)},
+                 {'value': 9, 'label': 'Alfalfa', 'rgb': (255,168,227)},
+                 {'value': 10, 'label': 'Fallow/Idle cropland', 'rgb': (191,191,122)},
+                 {'value': 11, 'label': 'Cotton', 'rgb':(255,38,38)},
+                 {'value': 12, 'label': 'Sorghum', 'rgb':(255,158,15)},
+                 {'value': 13, 'label': 'Other', 'rgb':(0,175,77)}]
+
+
+def apply_color_map(rgb, color_map=cdl_color_map):
+    
+    
+    rgb_mapped = rgb.copy()
+    
+    for map_tmp in cdl_color_map:
+        
+        for i in range(3):
+            rgb_mapped[i] = np.where((rgb[0] == map_tmp['value']) & (rgb[1] == map_tmp['value']) & (rgb[2] == map_tmp['value']), map_tmp['rgb'][i], rgb_mapped[i])
+    
+    return rgb_mapped    
+    
+
 def stretch_rgb(rgb):
     
     ls_pct=0
@@ -170,8 +198,9 @@ def inference_on_file(target_image, model, custom_test_pipeline):
     time_taken = np.round(et - st, 1)
     print(f'Inference completed in {str(time_taken)} seconds')
     
-    output = (result[0][0]*18).astype(np.uint8)
-    output = np.vstack([output[None], output[None], output[None]]).transpose((1,2,0))
+    output=result[0][0] + 1
+    output = np.vstack([output[None], output[None], output[None]]).astype(np.uint8)
+    output=apply_color_map(output).transpose((1,2,0))
         
     return rgb1,rgb2,rgb3,output