Update app.py

app.py

@@ -21,6 +21,17 @@ model_dir = snapshot_download(
     repo_id="srijaydeshpande/spadesegresnet"
 )
 
+color_map = {
+'outside_roi' : (255, 255, 255) # white
+'tumor' : (255, 0, 0) # red
+'stroma' : (0, 0, 255) # blue
+'inflammatory' : (0, 255, 0) # green
+'necrosis' : (255, 255, 0) # yello
+'others' : (8, 133, 161) # cyan
+}
+class_labels = ['outside_roi', 'tumor', 'stroma', 'inflammatory', 'necrosis', 'others']
+colors = ['white', 'red', 'blue', 'green', 'yello', 'cyan']
+
 class SPADE(nn.Module):
     def __init__(self, norm_nc, label_nc, norm):
         super().__init__()
@@ -194,12 +205,12 @@ def generate_colors(n):
   return colors
 
 def generate_colored_image(labels):
-    colors = generate_colors(6)
+    # colors = generate_colors(6)
     w, h = labels.shape
     new_mk = np.empty([w, h, 3])
     for i in range(0,w):
         for j in range(0,h):
-            new_mk[i][j] = colors[labels[i][j]]
+            new_mk[i][j] = color_map[class_labels[labels[i][j]]]
     new_mk = new_mk.astype(np.uint8)
     return Image.fromarray(new_mk)
 
@@ -239,11 +250,10 @@ def segment_image(image):
     pred_labels = np.argmax(pred_labels_probs, axis=1)
     pred_labels = pred_labels[0]
     image = generate_colored_image(pred_labels)
-    class_labels = ['tumor', 'stroma', 'inflammatory', 'necrosis', 'others']
     pixels_counts = []
     total=0
     print(np.unique(pred_labels))
-    for i in range(1,len(class_labels)+1):
+    for i in range(1,len(class_labels)):
         current_count=np.sum(pred_labels == i)
         pixels_counts.append(current_count)
         total+=current_count
@@ -251,7 +261,7 @@ def segment_image(image):
     print(pixels_counts)
     plt.figure(figsize=(10, 6))
     bar_width = 0.15
-    plt.bar(class_labels, pixels_counts, color='blue', width=bar_width)
+    plt.bar(class_labels[1:], pixels_counts, color=colors[1:], width=bar_width)
     plt.xticks(rotation=45, ha='right')
     plt.xlabel('Tissue types', fontsize=17)
     plt.ylabel('Class Percentage', fontsize=17)
@@ -264,11 +274,11 @@ def segment_image(image):
         temp_filename = tmpfile.name
     stats = Image.open(temp_filename)
 
-    legend = Image.open('legend.png')
+    # legend = Image.open('legend.png')
 
     superimposed_image = superimpose_images(img, image)
 
-    return image, legend, stats, superimposed_image
+    return image, stats, superimposed_image
 
 def superimpose_images(image1, image2):
 
@@ -307,7 +317,6 @@ with gr.Row():
         
         # Second column: Remaining three outputs
         with gr.Column():
-            output2 = gr.Image(label="Legend")  # Second output
             output3 = gr.Image(label="Statistics")  # Third output
             output4 = gr.Image(label="Superimposed Map")  # Fourth output
 
@@ -316,7 +325,7 @@ demo = gr.Interface(
     segment_image,
     inputs=input_image,
     examples=examples,
-    outputs=[output1, output2, output3, output4],
+    outputs=[output1, output3, output4],
     title="Breast Cancer Semantic Segmentation"
 )