viz_app.py

import streamlit as st
from detectron2.data import Metadata, MetadataCatalog, DatasetCatalog
from detectron2.utils.visualizer import Visualizer


# Define the plot_single_image function
def plot_single_image(im, predictor, num_classes):
    
    outputs = predictor(im)

    if num_classes == 2:
        modified_metadata = Metadata()
        modified_metadata.thing_classes = ["Nest", "Inorganic_material"]
    elif num_classes == 4:
        modified_metadata = Metadata()
        modified_metadata.thing_classes = ["Plastic", "Fishing_Net", "Rope/Cloth", "Nest"]
        
    st.write('Using Visualizer to draw the predictions on Image')
    v = Visualizer(im[:, :, ::-1], modified_metadata, scale=1.2)
    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    st.image(out.get_image()[:, :, ::-1])

    return outputs