ultralytics 8.0.237__py3-none-any.whl → 8.0.239__py3-none-any.whl

ultralytics/data/explorer/gui/dash.py

@@ -1,87 +1,122 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
 import time
 from threading import Thread
 
+import pandas as pd
+
 from ultralytics import Explorer
-from ultralytics.utils import ROOT
+from ultralytics.utils import ROOT, SETTINGS
 from ultralytics.utils.checks import check_requirements
 
-check_requirements('streamlit')
-check_requirements('streamlit-select>=0.2')
+check_requirements(("streamlit>=1.29.0", "streamlit-select>=0.2"))
+
 import streamlit as st
 from streamlit_select import image_select
 
 
 def _get_explorer():
-    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')})
+    """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")}
+    )
     thread.start()
-    progress_bar = st.progress(0, text='Creating embeddings table...')
+    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}%')
+        progress_bar.progress(exp.progress, text=f"Progress: {exp.progress * 100}%")
     thread.join()
-    st.session_state['explorer'] = exp
+    st.session_state["explorer"] = exp
     progress_bar.empty()
 
 
 def init_explorer_form():
-    datasets = ROOT / 'cfg' / 'datasets'
-    ds = [d.name for d in datasets.glob('*.yaml')]
+    """Initializes an Explorer instance and creates embeddings table with progress tracking."""
+    datasets = ROOT / "cfg" / "datasets"
+    ds = [d.name for d in datasets.glob("*.yaml")]
     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']
-    with st.form(key='explorer_init_form'):
+        "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",
+    ]
+    with st.form(key="explorer_init_form"):
         col1, col2 = st.columns(2)
         with col1:
-            dataset = st.selectbox('Select dataset', ds, key='dataset', index=ds.index('coco128.yaml'))
+            st.selectbox("Select dataset", ds, key="dataset", index=ds.index("coco128.yaml"))
         with col2:
-            model = st.selectbox('Select model', models, key='model')
-        st.checkbox('Force recreate embeddings', key='force_recreate_embeddings')
+            st.selectbox("Select model", models, key="model")
+        st.checkbox("Force recreate embeddings", key="force_recreate_embeddings")
 
-        st.form_submit_button('Explore', on_click=_get_explorer)
+        st.form_submit_button("Explore", on_click=_get_explorer)
 
 
 def query_form():
-    with st.form('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:
-            query = st.text_input('Query', '', label_visibility='collapsed', key='query')
+            st.text_input("Query", "Show images with 1 person and 1 dog", label_visibility="collapsed", key="ai_query")
         with col2:
-            st.form_submit_button('Query', on_click=run_sql_query)
+            st.form_submit_button("Ask AI", on_click=run_ai_query)
 
 
 def find_similar_imgs(imgs):
-    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
+    """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
 
 
 def similarity_form(selected_imgs):
-    st.write('Similarity Search')
-    with st.form('similarity_form'):
+    """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:
-            limit = st.number_input('limit',
-                                    min_value=None,
-                                    max_value=None,
-                                    value=25,
-                                    label_visibility='collapsed',
-                                    key='limit')
+            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.write("Selected: ", len(selected_imgs))
             st.form_submit_button(
-                'Search',
+                "Search",
                 disabled=disabled,
                 on_click=find_similar_imgs,
-                args=(selected_imgs, ),
+                args=(selected_imgs,),
             )
         if disabled:
-            st.error('Select at least one image to search.')
+            st.error("Select at least one image to search.")
 
 
 # def persist_reset_form():
@@ -95,74 +130,109 @@ def similarity_form(selected_imgs):
 
 
 def run_sql_query():
-    query = st.session_state.get('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"]
+
+
+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
+    st.session_state["error"] = None
+    query = st.session_state.get("ai_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']
+        exp = st.session_state["explorer"]
+        res = exp.ask_ai(query)
+        if not isinstance(res, pd.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()
 
 
 def reset_explorer():
-    st.session_state['explorer'] = None
-    st.session_state['imgs'] = None
+    """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.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=''>API docs</a> to try examples & learn more</p>",
+            "<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/')
+            help=None,
+        )
+        st.link_button("Ultrlaytics Explorer API", "https://docs.ultralytics.com/datasets/explorer/")
 
 
 def layout():
-    st.set_page_config(layout='wide', initial_sidebar_state='collapsed')
+    """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:
+    if st.session_state.get("explorer") is None:
         init_explorer_form()
         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 = st.session_state.get('imgs') or exp.table.to_lance().to_table(columns=['im_file']).to_pydict()['im_file']
-    total_imgs = len(imgs)
+    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"])
+    else:
+        imgs = st.session_state.get("imgs") or exp.table.to_lance().to_table(columns=["im_file"]).to_pydict()["im_file"]
+    total_imgs, selected_imgs = len(imgs), []
     with col1:
         subcol1, subcol2, subcol3, subcol4, subcol5 = st.columns(5)
         with subcol1:
-            st.write('Max Images Displayed:')
+            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')
+            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:')
+            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')
+            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')
+            reset = st.button("Reset", use_container_width=False, key="reset")
             if reset:
-                st.session_state['imgs'] = None
+                st.session_state["imgs"] = None
                 st.experimental_rerun()
 
         query_form()
+        ai_query_form()
         if total_imgs:
-            imgs_displayed = imgs[start_idx:start_idx + num]
+            imgs_displayed = imgs[start_idx : start_idx + num]
             selected_imgs = image_select(
-                f'Total samples: {total_imgs}',
+                f"Total samples: {total_imgs}",
                 images=imgs_displayed,
                 use_container_width=False,
                 # indices=[i for i in range(num)] if select_all else None,
@@ -174,5 +244,5 @@ def layout():
         utralytics_explorer_docs_callback()
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     layout()

ultralytics/data/explorer/utils.py

@@ -1,14 +1,22 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+import getpass
 from typing import List
 
 import cv2
 import numpy as np
+import pandas as pd
 
 from ultralytics.data.augment import LetterBox
+from ultralytics.utils import LOGGER as logger
+from ultralytics.utils import SETTINGS
+from ultralytics.utils.checks import check_requirements
 from ultralytics.utils.ops import xyxy2xywh
 from ultralytics.utils.plotting import plot_images
 
 
 def get_table_schema(vector_size):
+    """Extracts and returns the schema of a database table."""
     from lancedb.pydantic import LanceModel, Vector
 
     class Schema(LanceModel):
@@ -24,6 +32,7 @@ def get_table_schema(vector_size):
 
 
 def get_sim_index_schema():
+    """Returns a LanceModel schema for a database table with specified vector size."""
     from lancedb.pydantic import LanceModel
 
     class Schema(LanceModel):
@@ -36,33 +45,35 @@ def get_sim_index_schema():
 
 
 def sanitize_batch(batch, dataset_info):
-    batch['cls'] = batch['cls'].flatten().int().tolist()
-    box_cls_pair = sorted(zip(batch['bboxes'].tolist(), batch['cls']), key=lambda x: x[1])
-    batch['bboxes'] = [box for box, _ in box_cls_pair]
-    batch['cls'] = [cls for _, cls in box_cls_pair]
-    batch['labels'] = [dataset_info['names'][i] for i in batch['cls']]
-    batch['masks'] = batch['masks'].tolist() if 'masks' in batch else [[[]]]
-    batch['keypoints'] = batch['keypoints'].tolist() if 'keypoints' in batch else [[[]]]
-
+    """Sanitizes input batch for inference, ensuring correct format and dimensions."""
+    batch["cls"] = batch["cls"].flatten().int().tolist()
+    box_cls_pair = sorted(zip(batch["bboxes"].tolist(), batch["cls"]), key=lambda x: x[1])
+    batch["bboxes"] = [box for box, _ in box_cls_pair]
+    batch["cls"] = [cls for _, cls in box_cls_pair]
+    batch["labels"] = [dataset_info["names"][i] for i in batch["cls"]]
+    batch["masks"] = batch["masks"].tolist() if "masks" in batch else [[[]]]
+    batch["keypoints"] = batch["keypoints"].tolist() if "keypoints" in batch else [[[]]]
     return batch
 
 
-def plot_similar_images(similar_set, plot_labels=True):
+def plot_query_result(similar_set, plot_labels=True):
     """
     Plot images from the similar set.
 
     Args:
-        similar_set (list): Pyarrow table containing the similar data points
+        similar_set (list): Pyarrow or pandas object containing the similar data points
         plot_labels (bool): Whether to plot labels or not
     """
-    similar_set = similar_set.to_pydict()
+    similar_set = (
+        similar_set.to_dict(orient="list") if isinstance(similar_set, pd.DataFrame) else similar_set.to_pydict()
+    )
     empty_masks = [[[]]]
     empty_boxes = [[]]
-    images = similar_set.get('im_file', [])
-    bboxes = similar_set.get('bboxes', []) if similar_set.get('bboxes') is not empty_boxes else []
-    masks = similar_set.get('masks') if similar_set.get('masks')[0] != empty_masks else []
-    kpts = similar_set.get('keypoints') if similar_set.get('keypoints')[0] != empty_masks else []
-    cls = similar_set.get('cls', [])
+    images = similar_set.get("im_file", [])
+    bboxes = similar_set.get("bboxes", []) if similar_set.get("bboxes") is not empty_boxes else []
+    masks = similar_set.get("masks") if similar_set.get("masks")[0] != empty_masks else []
+    kpts = similar_set.get("keypoints") if similar_set.get("keypoints")[0] != empty_masks else []
+    cls = similar_set.get("cls", [])
 
     plot_size = 640
     imgs, batch_idx, plot_boxes, plot_masks, plot_kpts = [], [], [], [], []
@@ -93,12 +104,63 @@ def plot_similar_images(similar_set, plot_labels=True):
     batch_idx = np.concatenate(batch_idx, axis=0)
     cls = np.concatenate([np.array(c, dtype=np.int32) for c in cls], axis=0)
 
-    return plot_images(imgs,
-                       batch_idx,
-                       cls,
-                       bboxes=boxes,
-                       masks=masks,
-                       kpts=kpts,
-                       max_subplots=len(images),
-                       save=False,
-                       threaded=False)
+    return plot_images(
+        imgs, batch_idx, cls, bboxes=boxes, masks=masks, kpts=kpts, max_subplots=len(images), save=False, threaded=False
+    )
+
+
+def prompt_sql_query(query):
+    """Plots images with optional labels from a similar data set."""
+    check_requirements("openai>=1.6.1")
+    from openai import OpenAI
+
+    if not SETTINGS["openai_api_key"]:
+        logger.warning("OpenAI API key not found in settings. Please enter your API key below.")
+        openai_api_key = getpass.getpass("OpenAI API key: ")
+        SETTINGS.update({"openai_api_key": openai_api_key})
+    openai = OpenAI(api_key=SETTINGS["openai_api_key"])
+
+    messages = [
+        {
+            "role": "system",
+            "content": """
+                You are a helpful data scientist proficient in SQL. You need to output exactly one SQL query based on
+                the following schema and a user request. You only need to output the format with fixed selection
+                statement that selects everything from "'table'", like `SELECT * from 'table'`
+
+                Schema:
+                im_file: string not null
+                labels: list<item: string> not null
+                child 0, item: string
+                cls: list<item: int64> not null
+                child 0, item: int64
+                bboxes: list<item: list<item: double>> not null
+                child 0, item: list<item: double>
+                    child 0, item: double
+                masks: list<item: list<item: list<item: int64>>> not null
+                child 0, item: list<item: list<item: int64>>
+                    child 0, item: list<item: int64>
+                        child 0, item: int64
+                keypoints: list<item: list<item: list<item: double>>> not null
+                child 0, item: list<item: list<item: double>>
+                    child 0, item: list<item: double>
+                        child 0, item: double
+                vector: fixed_size_list<item: float>[256] not null
+                child 0, item: float
+
+                Some details about the schema:
+                - the "labels" column contains the string values like 'person' and 'dog' for the respective objects
+                    in each image
+                - the "cls" column contains the integer values on these classes that map them the labels
+
+                Example of a correct query:
+                request - Get all data points that contain 2 or more people and at least one dog
+                correct query-
+                SELECT * FROM 'table' WHERE  ARRAY_LENGTH(cls) >= 2  AND ARRAY_LENGTH(FILTER(labels, x -> x = 'person')) >= 2  AND ARRAY_LENGTH(FILTER(labels, x -> x = 'dog')) >= 1;
+             """,
+        },
+        {"role": "user", "content": f"{query}"},
+    ]
+
+    response = openai.chat.completions.create(model="gpt-3.5-turbo", messages=messages)
+    return response.choices[0].message.content