Update app.py

app.py

@@ -106,15 +106,13 @@ if __name__ == "__main__":
     label2color = {"unlabeled":[255, 255, 255], "human": [255, 0, 0], "background": [29, 162, 216], "fish": [255, 255, 0], "sand": [194, 178, 128], "rubble": [161, 153, 128], "unknown hard substrate": [125, 125, 125], "algae covered substrate": [125, 163, 125], "dark": [31, 31, 31], "branching bleached": [252, 231, 240], "branching dead": [123, 50, 86], "branching alive": [226, 91, 157], "stylophora alive": [255, 111, 194], "pocillopora alive": [255, 146, 150], "acropora alive": [236, 128, 255], "table acropora alive": [189, 119, 255], "table acropora dead": [85, 53, 116], "millepora": [244, 150, 115], "turbinaria": [228, 255, 119], "other coral bleached": [250, 224, 225], "other coral dead": [114, 60, 61], "other coral alive": [224, 118, 119], "massive/meandering alive": [236, 150, 21], "massive/meandering dead": [134, 86, 18], "massive/meandering bleached": [255, 248, 228], "meandering alive": [230, 193, 0], "meandering dead": [119, 100, 14], "meandering bleached": [251, 243, 216], "transect line": [0, 255, 0], "transect tools": [8, 205, 12], "sea urchin": [0, 142, 255], "sea cucumber": [0, 231, 255], "anemone": [0, 255, 189], "sponge": [240, 80, 80], "clam": [189, 255, 234], "other animal": [0, 255, 255], "trash": [255, 0, 134], "seagrass": [125, 222, 125], "crown of thorn": [179, 245, 234], "dead clam": [89, 155, 134]}
     label2colorhex = {k:'#%02x%02x%02x' % tuple(v) for k,v in label2color.items()}
     print(label2colorhex)
-    id2color = {int(id): label2color[label] for id, label in id2label.items()}
-
 
     with gr.Blocks(title="Coral Segmentation with SegFormer") as demo:
         gr.Markdown("""<h1><center>Coral Segmentation with SegFormer</center></h1>""") 
         with gr.Row():
             img_input = gr.Image(type="pil", label="Input image")
             # img_output = gr.Image(type="pil", label="Predictions")
-            img_output = gr.AnnotatedImage(label="Predictions", color_map=label2color)
+            img_output = gr.AnnotatedImage(label="Predictions", color_map=label2colorhex)
 
         section_btn = gr.Button("Segment Image")
         section_btn.click(partial(segment_image, preprocessor=preprocessor, model=model), img_input, img_output)