Update app.py

app.py

@@ -18,6 +18,7 @@ import pandas as pd
 import matplotlib.pyplot as plt
 from PIL import Image
 from shiny import App, ui, render, reactive, Session
+from detectron2.utils.visualizer import Visualizer, ColorMode
 
 from python_utils import load_model
 # Load data and compute static values
@@ -105,7 +106,18 @@ def server(input, output, session: Session):
             @render.plot
             def plot_predicitons():
                 fig, ax = plt.subplots()
-                plt.imshow(r["image"])
+                
+                v = Visualizer(r["image"[:, :, ::-1], 
+                               {"thing_classes": ['Seed', 'Viable', 'Non-Viable', 'Empty'],
+                                "thing_colors": [(0, 0, 0), (0, 255, 0), (255, 0, 0), (0, 0, 255)]},
+                               scale=1.2, instance_mode=ColorMode.SEGMENTATION, font_size_scale=1.5)
+                out = v.draw_instance_predictions(results["instances"].to("cpu"), jittering=False)
+                ax.imshow(cv2.cvtColor(out.get_image()[:, :, ::-1], cv2.COLOR_BGR2RGB))
+                
+                ax = plt.Axes(fig, [0., 0., 1., 1.])
+                ax.set_axis_off()
+                fig.add_axes(ax)
+                
                 return fig
                 
             output.append(