Update app.py

app.py

@@ -166,6 +166,27 @@ def load_network_pkl_cpu(f, force_fp16=False):
                 data[key] = new
     return data
 
+color_list = [[255, 255, 255], [204, 0, 0], [76, 153, 0], [204, 204, 0], [51, 51, 255], [204, 0, 204], [0, 255, 255], [255, 204, 204], [102, 51, 0], [255, 0, 0], [102, 204, 0], [255, 255, 0], [0, 0, 153], [0, 0, 204], [255, 51, 153], [0, 204, 204], [0, 51, 0], [255, 153, 51], [0, 204, 0]]
+
+def colormap2labelmap(color_img):
+  im_base = np.zeros((color_img.shape[0], color_img.shape[1]))
+  for idx, color in enumerate(color_list):
+    
+      k1=((color_img == np.asarray(color))[:,:,0])*1
+      k2=((color_img == np.asarray(color))[:,:,1])*1
+      k3=((color_img == np.asarray(color))[:,:,2])*1
+      k=((k1*k2*k3)==1)
+
+      im_base[k] = idx
+  return im_base
+
+
+def checklabelmap(img):
+  labels=np.unique(img)
+  for idx,label in enumerate(labels):
+    img[img==label]=idx
+  return img
+    
 def get_all(cfg,input,truncation_psi,mesh_resolution,random_seed,fps,num_frames):
 
         network=models[cfg]
@@ -194,11 +215,16 @@ def get_all(cfg,input,truncation_psi,mesh_resolution,random_seed,fps,num_frames)
 
         save_dir = Path(outdir)
 
-        input_label = PIL.Image.open(input)
-        input_label = PIL.ImageOps.grayscale(input_label)
+          
+        if isinstance(input,str):
+          input_label =np.asarray( PIL.Image.open(input))
+        else:
+          input_label=np.asarray(input)
+
+        input_label=colormap2labelmap(input_label)
+        input_label=checklabelmap(input_label)
         input_label = np.asarray(input_label).astype(np.uint8)
         input_label = torch.from_numpy(input_label).unsqueeze(0).unsqueeze(0).to(device)
-        print(input_label.shape)
         input_pose = forward_pose.to(device)
 
         # Generate videos
@@ -282,7 +308,7 @@ demo_outputs=[
 
 ]
 examples = [
-    ["seg2cat", "example_input.png", 1, 32, 128, 30, 30],
+    ["seg2cat", "img.png", 1, 32, 128, 30, 30],
 
 ]