Spaces:
Sleeping
Sleeping
File size: 843 Bytes
239f6f6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
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
|