283 lines
10 KiB
Python
283 lines
10 KiB
Python
# Ultralytics YOLO 🚀, AGPL-3.0 license
|
|
|
|
import sys
|
|
import time
|
|
from threading import Thread
|
|
|
|
from ultralytics import Explorer
|
|
from ultralytics.utils import ROOT, SETTINGS
|
|
from ultralytics.utils.checks import check_requirements
|
|
|
|
check_requirements(("streamlit>=1.29.0", "streamlit-select>=0.3"))
|
|
|
|
import streamlit as st
|
|
from streamlit_select import image_select
|
|
|
|
|
|
def _get_explorer():
|
|
"""Initializes and returns an instance of the Explorer class."""
|
|
exp = Explorer(data=st.session_state.get("dataset"), model=st.session_state.get("model"))
|
|
thread = Thread(
|
|
target=exp.create_embeddings_table,
|
|
kwargs={"force": st.session_state.get("force_recreate_embeddings"), "split": st.session_state.get("split")},
|
|
)
|
|
thread.start()
|
|
progress_bar = st.progress(0, text="Creating embeddings table...")
|
|
while exp.progress < 1:
|
|
time.sleep(0.1)
|
|
progress_bar.progress(exp.progress, text=f"Progress: {exp.progress * 100}%")
|
|
thread.join()
|
|
st.session_state["explorer"] = exp
|
|
progress_bar.empty()
|
|
|
|
|
|
def init_explorer_form(data=None, model=None):
|
|
"""Initializes an Explorer instance and creates embeddings table with progress tracking."""
|
|
if data is None:
|
|
datasets = ROOT / "cfg" / "datasets"
|
|
ds = [d.name for d in datasets.glob("*.yaml")]
|
|
else:
|
|
ds = [data]
|
|
|
|
if model is None:
|
|
models = [
|
|
"yolov8n.pt",
|
|
"yolov8s.pt",
|
|
"yolov8m.pt",
|
|
"yolov8l.pt",
|
|
"yolov8x.pt",
|
|
"yolov8n-seg.pt",
|
|
"yolov8s-seg.pt",
|
|
"yolov8m-seg.pt",
|
|
"yolov8l-seg.pt",
|
|
"yolov8x-seg.pt",
|
|
"yolov8n-pose.pt",
|
|
"yolov8s-pose.pt",
|
|
"yolov8m-pose.pt",
|
|
"yolov8l-pose.pt",
|
|
"yolov8x-pose.pt",
|
|
]
|
|
else:
|
|
models = [model]
|
|
|
|
splits = ["train", "val", "test"]
|
|
|
|
with st.form(key="explorer_init_form"):
|
|
col1, col2, col3 = st.columns(3)
|
|
with col1:
|
|
st.selectbox("Select dataset", ds, key="dataset")
|
|
with col2:
|
|
st.selectbox("Select model", models, key="model")
|
|
with col3:
|
|
st.selectbox("Select split", splits, key="split")
|
|
st.checkbox("Force recreate embeddings", key="force_recreate_embeddings")
|
|
|
|
st.form_submit_button("Explore", on_click=_get_explorer)
|
|
|
|
|
|
def query_form():
|
|
"""Sets up a form in Streamlit to initialize Explorer with dataset and model selection."""
|
|
with st.form("query_form"):
|
|
col1, col2 = st.columns([0.8, 0.2])
|
|
with col1:
|
|
st.text_input(
|
|
"Query",
|
|
"WHERE labels LIKE '%person%' AND labels LIKE '%dog%'",
|
|
label_visibility="collapsed",
|
|
key="query",
|
|
)
|
|
with col2:
|
|
st.form_submit_button("Query", on_click=run_sql_query)
|
|
|
|
|
|
def ai_query_form():
|
|
"""Sets up a Streamlit form for user input to initialize Explorer with dataset and model selection."""
|
|
with st.form("ai_query_form"):
|
|
col1, col2 = st.columns([0.8, 0.2])
|
|
with col1:
|
|
st.text_input("Query", "Show images with 1 person and 1 dog", label_visibility="collapsed", key="ai_query")
|
|
with col2:
|
|
st.form_submit_button("Ask AI", on_click=run_ai_query)
|
|
|
|
|
|
def find_similar_imgs(imgs):
|
|
"""Initializes a Streamlit form for AI-based image querying with custom input."""
|
|
exp = st.session_state["explorer"]
|
|
similar = exp.get_similar(img=imgs, limit=st.session_state.get("limit"), return_type="arrow")
|
|
paths = similar.to_pydict()["im_file"]
|
|
st.session_state["imgs"] = paths
|
|
st.session_state["res"] = similar
|
|
|
|
|
|
def similarity_form(selected_imgs):
|
|
"""Initializes a form for AI-based image querying with custom input in Streamlit."""
|
|
st.write("Similarity Search")
|
|
with st.form("similarity_form"):
|
|
subcol1, subcol2 = st.columns([1, 1])
|
|
with subcol1:
|
|
st.number_input(
|
|
"limit", min_value=None, max_value=None, value=25, label_visibility="collapsed", key="limit"
|
|
)
|
|
|
|
with subcol2:
|
|
disabled = not len(selected_imgs)
|
|
st.write("Selected: ", len(selected_imgs))
|
|
st.form_submit_button(
|
|
"Search",
|
|
disabled=disabled,
|
|
on_click=find_similar_imgs,
|
|
args=(selected_imgs,),
|
|
)
|
|
if disabled:
|
|
st.error("Select at least one image to search.")
|
|
|
|
|
|
# def persist_reset_form():
|
|
# with st.form("persist_reset"):
|
|
# col1, col2 = st.columns([1, 1])
|
|
# with col1:
|
|
# st.form_submit_button("Reset", on_click=reset)
|
|
#
|
|
# with col2:
|
|
# st.form_submit_button("Persist", on_click=update_state, args=("PERSISTING", True))
|
|
|
|
|
|
def run_sql_query():
|
|
"""Executes an SQL query and returns the results."""
|
|
st.session_state["error"] = None
|
|
query = st.session_state.get("query")
|
|
if query.rstrip().lstrip():
|
|
exp = st.session_state["explorer"]
|
|
res = exp.sql_query(query, return_type="arrow")
|
|
st.session_state["imgs"] = res.to_pydict()["im_file"]
|
|
st.session_state["res"] = res
|
|
|
|
|
|
def run_ai_query():
|
|
"""Execute SQL query and update session state with query results."""
|
|
if not SETTINGS["openai_api_key"]:
|
|
st.session_state["error"] = (
|
|
'OpenAI API key not found in settings. Please run yolo settings openai_api_key="..."'
|
|
)
|
|
return
|
|
import pandas # scope for faster 'import ultralytics'
|
|
|
|
st.session_state["error"] = None
|
|
query = st.session_state.get("ai_query")
|
|
if query.rstrip().lstrip():
|
|
exp = st.session_state["explorer"]
|
|
res = exp.ask_ai(query)
|
|
if not isinstance(res, pandas.DataFrame) or res.empty:
|
|
st.session_state["error"] = "No results found using AI generated query. Try another query or rerun it."
|
|
return
|
|
st.session_state["imgs"] = res["im_file"].to_list()
|
|
st.session_state["res"] = res
|
|
|
|
|
|
def reset_explorer():
|
|
"""Resets the explorer to its initial state by clearing session variables."""
|
|
st.session_state["explorer"] = None
|
|
st.session_state["imgs"] = None
|
|
st.session_state["error"] = None
|
|
|
|
|
|
def utralytics_explorer_docs_callback():
|
|
"""Resets the explorer to its initial state by clearing session variables."""
|
|
with st.container(border=True):
|
|
st.image(
|
|
"https://raw.githubusercontent.com/ultralytics/assets/main/logo/Ultralytics_Logotype_Original.svg",
|
|
width=100,
|
|
)
|
|
st.markdown(
|
|
"<p>This demo is built using Ultralytics Explorer API. Visit <a href='https://docs.ultralytics.com/datasets/explorer/'>API docs</a> to try examples & learn more</p>",
|
|
unsafe_allow_html=True,
|
|
help=None,
|
|
)
|
|
st.link_button("Ultrlaytics Explorer API", "https://docs.ultralytics.com/datasets/explorer/")
|
|
|
|
|
|
def layout(data=None, model=None):
|
|
"""Resets explorer session variables and provides documentation with a link to API docs."""
|
|
st.set_page_config(layout="wide", initial_sidebar_state="collapsed")
|
|
st.markdown("<h1 style='text-align: center;'>Ultralytics Explorer Demo</h1>", unsafe_allow_html=True)
|
|
|
|
if st.session_state.get("explorer") is None:
|
|
init_explorer_form(data, model)
|
|
return
|
|
|
|
st.button(":arrow_backward: Select Dataset", on_click=reset_explorer)
|
|
exp = st.session_state.get("explorer")
|
|
col1, col2 = st.columns([0.75, 0.25], gap="small")
|
|
imgs = []
|
|
if st.session_state.get("error"):
|
|
st.error(st.session_state["error"])
|
|
elif st.session_state.get("imgs"):
|
|
imgs = st.session_state.get("imgs")
|
|
else:
|
|
imgs = exp.table.to_lance().to_table(columns=["im_file"]).to_pydict()["im_file"]
|
|
st.session_state["res"] = exp.table.to_arrow()
|
|
total_imgs, selected_imgs = len(imgs), []
|
|
with col1:
|
|
subcol1, subcol2, subcol3, subcol4, subcol5 = st.columns(5)
|
|
with subcol1:
|
|
st.write("Max Images Displayed:")
|
|
with subcol2:
|
|
num = st.number_input(
|
|
"Max Images Displayed",
|
|
min_value=0,
|
|
max_value=total_imgs,
|
|
value=min(500, total_imgs),
|
|
key="num_imgs_displayed",
|
|
label_visibility="collapsed",
|
|
)
|
|
with subcol3:
|
|
st.write("Start Index:")
|
|
with subcol4:
|
|
start_idx = st.number_input(
|
|
"Start Index",
|
|
min_value=0,
|
|
max_value=total_imgs,
|
|
value=0,
|
|
key="start_index",
|
|
label_visibility="collapsed",
|
|
)
|
|
with subcol5:
|
|
reset = st.button("Reset", use_container_width=False, key="reset")
|
|
if reset:
|
|
st.session_state["imgs"] = None
|
|
st.experimental_rerun()
|
|
|
|
query_form()
|
|
ai_query_form()
|
|
if total_imgs:
|
|
labels, boxes, masks, kpts, classes = None, None, None, None, None
|
|
task = exp.model.task
|
|
if st.session_state.get("display_labels"):
|
|
labels = st.session_state.get("res").to_pydict()["labels"][start_idx : start_idx + num]
|
|
boxes = st.session_state.get("res").to_pydict()["bboxes"][start_idx : start_idx + num]
|
|
masks = st.session_state.get("res").to_pydict()["masks"][start_idx : start_idx + num]
|
|
kpts = st.session_state.get("res").to_pydict()["keypoints"][start_idx : start_idx + num]
|
|
classes = st.session_state.get("res").to_pydict()["cls"][start_idx : start_idx + num]
|
|
imgs_displayed = imgs[start_idx : start_idx + num]
|
|
selected_imgs = image_select(
|
|
f"Total samples: {total_imgs}",
|
|
images=imgs_displayed,
|
|
use_container_width=False,
|
|
# indices=[i for i in range(num)] if select_all else None,
|
|
labels=labels,
|
|
classes=classes,
|
|
bboxes=boxes,
|
|
masks=masks if task == "segment" else None,
|
|
kpts=kpts if task == "pose" else None,
|
|
)
|
|
|
|
with col2:
|
|
similarity_form(selected_imgs)
|
|
st.checkbox("Labels", value=False, key="display_labels")
|
|
utralytics_explorer_docs_callback()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
kwargs = dict(zip(sys.argv[1::2], sys.argv[2::2]))
|
|
layout(**kwargs)
|