142 lines
3.6 KiB
Python
142 lines
3.6 KiB
Python
from pathlib import Path
|
|
import sys
|
|
import io
|
|
import streamlit as st
|
|
from PIL import Image, ImageDraw
|
|
import numpy as np
|
|
import os
|
|
import deepstack.core as ds
|
|
import utils
|
|
import const
|
|
|
|
## Depstack setup
|
|
DEEPSTACK_IP = os.getenv("DEEPSTACK_IP", "localhost")
|
|
DEEPSTACK_PORT = os.getenv("DEEPSTACK_PORT", 80)
|
|
DEEPSTACK_API_KEY = os.getenv("DEEPSTACK_API_KEY", "")
|
|
DEEPSTACK_TIMEOUT = int(os.getenv("DEEPSTACK_TIMEOUT", 20))
|
|
|
|
DEFAULT_CONFIDENCE_THRESHOLD = 45 # percent
|
|
TEST_IMAGE = "street.jpg"
|
|
|
|
DEFAULT_ROI_Y_MIN = 0.0
|
|
DEFAULT_ROI_Y_MAX = 1.0
|
|
DEFAULT_ROI_X_MIN = 0.0
|
|
DEFAULT_ROI_X_MAX = 1.0
|
|
DEFAULT_ROI = (
|
|
DEFAULT_ROI_Y_MIN,
|
|
DEFAULT_ROI_X_MIN,
|
|
DEFAULT_ROI_Y_MAX,
|
|
DEFAULT_ROI_X_MAX,
|
|
)
|
|
|
|
predictions = None
|
|
|
|
|
|
@st.cache
|
|
def process_image(pil_image, dsobject):
|
|
image_bytes = utils.pil_image_to_byte_array(pil_image)
|
|
predictions = dsobject.detect(image_bytes)
|
|
return predictions
|
|
|
|
|
|
## Setup sidebar
|
|
st.title("Deepstack Object detection")
|
|
img_file_buffer = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
|
|
|
|
st.sidebar.title("Parameters")
|
|
CONFIDENCE_THRESHOLD = st.sidebar.slider(
|
|
"Confidence threshold", 0, 100, DEFAULT_CONFIDENCE_THRESHOLD, 1
|
|
)
|
|
|
|
CLASSES_TO_INCLUDE = st.sidebar.multiselect(
|
|
"Select object classes to include",
|
|
options=const.CLASSES,
|
|
default=const.DEFAULT_CLASSES,
|
|
)
|
|
|
|
# Get ROI info
|
|
st.sidebar.title("ROI")
|
|
ROI_X_MIN = st.sidebar.slider("x_min", 0.0, 1.0, DEFAULT_ROI_X_MIN)
|
|
ROI_Y_MIN = st.sidebar.slider("y_min", 0.0, 1.0, DEFAULT_ROI_Y_MIN)
|
|
ROI_X_MAX = st.sidebar.slider("x_max", 0.0, 1.0, DEFAULT_ROI_X_MAX)
|
|
ROI_Y_MAX = st.sidebar.slider("y_max", 0.0, 1.0, DEFAULT_ROI_Y_MAX)
|
|
ROI_TUPLE = (
|
|
ROI_Y_MIN,
|
|
ROI_X_MIN,
|
|
ROI_Y_MAX,
|
|
ROI_X_MAX,
|
|
)
|
|
ROI_DICT = {
|
|
"x_min": ROI_X_MIN,
|
|
"y_min": ROI_Y_MIN,
|
|
"x_max": ROI_X_MAX,
|
|
"y_max": ROI_Y_MAX,
|
|
}
|
|
|
|
## Process image
|
|
if img_file_buffer is not None:
|
|
pil_image = Image.open(img_file_buffer)
|
|
|
|
else:
|
|
pil_image = Image.open(TEST_IMAGE)
|
|
|
|
dsobject = ds.DeepstackObject(
|
|
DEEPSTACK_IP, DEEPSTACK_PORT, DEEPSTACK_API_KEY, DEEPSTACK_TIMEOUT
|
|
)
|
|
|
|
predictions = process_image(pil_image, dsobject)
|
|
objects = utils.get_objects(predictions, pil_image.width, pil_image.height)
|
|
all_objects_names = set([obj["name"] for obj in objects])
|
|
|
|
# Filter objects for display
|
|
objects = [obj for obj in objects if obj["confidence"] > CONFIDENCE_THRESHOLD]
|
|
objects = [obj for obj in objects if obj["name"] in CLASSES_TO_INCLUDE]
|
|
objects = [obj for obj in objects if utils.object_in_roi(ROI_DICT, obj["centroid"])]
|
|
|
|
# Draw object boxes
|
|
draw = ImageDraw.Draw(pil_image)
|
|
for obj in objects:
|
|
name = obj["name"]
|
|
confidence = obj["confidence"]
|
|
box = obj["bounding_box"]
|
|
box_label = f"{name}: {confidence:.1f}%"
|
|
|
|
utils.draw_box(
|
|
draw,
|
|
(box["y_min"], box["x_min"], box["y_max"], box["x_max"]),
|
|
pil_image.width,
|
|
pil_image.height,
|
|
text=box_label,
|
|
color=const.YELLOW,
|
|
)
|
|
|
|
# Draw ROI box
|
|
if ROI_TUPLE != DEFAULT_ROI:
|
|
utils.draw_box(
|
|
draw,
|
|
ROI_TUPLE,
|
|
pil_image.width,
|
|
pil_image.height,
|
|
text="ROI",
|
|
color=const.GREEN,
|
|
)
|
|
|
|
# Display image and results
|
|
st.image(
|
|
np.array(pil_image), caption=f"Processed image", use_column_width=True,
|
|
)
|
|
st.subheader("All discovered objects")
|
|
st.write(all_objects_names)
|
|
|
|
st.subheader("Filtered object count")
|
|
obj_types = list(set([obj["name"] for obj in objects]))
|
|
for obj_type in obj_types:
|
|
obj_type_count = len([obj for obj in objects if obj["name"] == obj_type])
|
|
st.write(f"{obj_type} : {obj_type_count}")
|
|
|
|
st.subheader("All filtered objects")
|
|
st.write(objects)
|
|
|
|
st.subheader("Deepstack raw response")
|
|
st.write(predictions)
|