Spaces:
Sleeping
Sleeping
File size: 4,418 Bytes
239f6f6 5c861ff 239f6f6 1f10fa6 10db898 239f6f6 5c861ff 239f6f6 5c861ff 239f6f6 5c861ff 239f6f6 5c861ff b1587c2 5c861ff b1587c2 239f6f6 22d6c0a 239f6f6 18999c7 239f6f6 4d576d9 |
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 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
import streamlit as st
import torch,torchvision
import detectron2
import pickle
from detectron2.utils.logger import setup_logger
logger = setup_logger()
from detectron2.engine import DefaultPredictor
from detectron2.projects import point_rend
from viz_app import plot_single_image # Correct import
from PIL import Image
import numpy as np
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.utils.visualizer import Visualizer
from collections import Counter
# Define configuration parameters
output_directories = {
"4 Classes": "output/pointrend_4_cls_6_nest_8k/",
}
# Function to load the configuration
def load_configuration(output_dir):
with open(output_dir + "cfg.pickle", "rb") as f:
return pickle.load(f)
# Function to run instance segmentation
def run_instance_segmentation(im, predictor, num_classes):
outputs = plot_single_image(im, predictor,num_classes)
masks = outputs['instances'].pred_masks.cpu().numpy()
classes = outputs['instances'].pred_classes.cpu().numpy()
if num_classes == 2:
# Create masks for Nest and Inorganic Material
nest_masks = masks[classes == 0]
inorganic_material_masks = masks[classes == 1]
# Calculate the total pixel area of Nest masks
total_nest_area = np.sum(nest_masks)
# Calculate the total area of Inorganic Material masks
total_inorganic_area = np.sum(inorganic_material_masks)
# Contar la cantidad de máscaras por categoría
mask_counts = Counter(classes)
elif num_classes == 4:
# Create masks for Nest and Inorganic Material
plastic_masks = masks[classes == 0]
fishing_net_masks = masks[classes == 1]
rope_cloth_masks = masks[classes == 2]
nest_masks = masks[classes == 3]
# Calculate the total pixel area of Nest masks
total_nest_area = np.sum(nest_masks)
# Calculate the total area of Inorganic Material masks
plastic_area = np.sum(plastic_masks)
fishing_net_area = np.sum(fishing_net_masks)
rope_cloth_area = np.sum(rope_cloth_masks)
total_inorganic_area = np.sum([plastic_area, fishing_net_area, rope_cloth_area])
# Contar la cantidad de máscaras por categoría
mask_counts = Counter(classes)
return total_nest_area, total_inorganic_area, mask_counts
def click_instance_segmentation(image, model_selection, predictor):
im = np.array(image)[:, :, ::-1]
if st.button("Run Instance Segmentation"):
num_classes = 2 if model_selection == "2 Classes" else 4
total_nest_area, total_inorganic_area, mask_counts = run_instance_segmentation(im, predictor, num_classes)
# Calculate the percentage of inorganic material within the Nest
percentage_inorganic_in_nest = (total_inorganic_area / total_nest_area) * 100
st.write(f"Percentage of Inorganic Material in Nest: {percentage_inorganic_in_nest:.2f}%")
category_list = ['Plastic', 'Fishing_Net', 'Rope/Cloth', "Nest"]
# Mostrar la cantidad de máscaras por categoría
st.write("Mask Counts by Category:")
for category, count in mask_counts.items():
st.write(f"Category {category_list[category]}: {count} masks")
def app():
# Create a sidebar to select the model
model_selection = st.sidebar.radio("Select Model", list(output_directories.keys()))
output_dir = output_directories[model_selection]
cfg = load_configuration(output_dir)
cfg.MODEL.WEIGHTS = output_dir + "model_best.pth"
cfg.MODEL.DEVICE = "cpu"
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.65
predictor = DefaultPredictor(cfg)
default_image_path = 'image/DJI_0228_frame_0022_object_0.jpg'
st.header('Please upload an image')
file = st.file_uploader('', type=['png', 'jpg', 'jpeg'])
if file:
image = Image.open(file)
st.image(image, caption="Uploaded Image", use_column_width=True)
click_instance_segmentation(image, model_selection, predictor)
else:
st.write("No image uploaded. Using default example image.")
image = Image.open(default_image_path)
st.image(image, caption="Default Example Image", use_column_width=True)
click_instance_segmentation(image, model_selection, predictor)
app()
|