Update app.py

app.py

@@ -52,6 +52,11 @@ DATASET_COLORMAPS = {
     "ade20k": colormaps.ADE20K_COLORMAP,
     "voc2012": colormaps.VOC2012_COLORMAP,
 }
+colormap = DATASET_COLORMAPS["ade20k"]
+flattened = np.array(colormap).flatten()
+zeros = np.zeros(768)
+zeros[:flattened.shape[0]] = flattened
+colorMap = list(zeros.astype('uint8'))
 
 model = init_segmentor(cfg)
 load_checkpoint(model, CHECKPOINT_URL, map_location="cpu")
@@ -91,11 +96,11 @@ def create_segmenter(cfg, backbone_model):
 
 
 def render_segmentation(segmentation_logits, dataset):
-    colormap = DATASET_COLORMAPS[dataset]
     colormap_array = np.array(colormap, dtype=np.uint8)
     segmentation_logits += 1
-    segmentation_values = colormap_array[segmentation_logits]
-    segmentation_values = segmentation_values[:, :, ::-1]
+    segmented_image = Image.fromarray(segmentation_logits)
+    segmented_image.putpalette(colorMap)
+    
     unique_labels = np.unique(segmentation_logits)
 
     colormap_array = colormap_array[unique_labels]
@@ -107,7 +112,7 @@ def render_segmentation(segmentation_logits, dataset):
     for idx, color in enumerate(colormap_array):
         color_box = np.zeros((20, 20, 3), dtype=np.uint8)
         color_box[:, :] = color
-        # color_box = cv2.cvtColor(color_box, cv2.COLOR_RGB2BGR)
+        color_box = cv2.cvtColor(color_box, cv2.COLOR_RGB2BGR)
         _, img_data = cv2.imencode(".jpg", color_box)
         img_base64 = base64.b64encode(img_data).decode("utf-8")
         img_data_uri = f"data:image/jpg;base64,{img_base64}"
@@ -115,14 +120,15 @@ def render_segmentation(segmentation_logits, dataset):
 
     html_output += "</div>"
 
-    return Image.fromarray(segmentation_values), html_output
+    return segmented_image, html_output
 
 
 def predict(image_file):
     array = np.array(image_file)[:, :, ::-1]  # BGR
     segmentation_logits = inference_segmentor(model, array)[0]
+    segmentation_logits = segmentation_logits.astype(np.uint8)
     segmented_image, html_output = render_segmentation(segmentation_logits, "ade20k")
-    return np.array(segmented_image), html_output
+    return segmented_image, html_output
 
 description = "Gradio demo for Semantic segmentation. To use it, simply upload your image"
 
@@ -130,10 +136,10 @@ demo = gr.Interface(
     title="Semantic Segmentation - DinoV2",
     fn=predict,
     inputs=gr.inputs.Image(),
-    outputs=[gr.outputs.Image(type="numpy"), gr.outputs.HTML()],
+    outputs=[gr.outputs.Image(type="pil"), gr.outputs.HTML()],
     examples=["example_1.jpg", "example_2.jpg"],
     cache_examples=False,
     description=description,
 )
 
-demo.launch()
+demo.launch()