Repository: ultralytics/ultralytics
Pull request: `ultralytics 8.0.236` dataset semantic & SQL search API


copilot:all


## 🛠️ PR Summary

<sub>Made with ❤️ by [Ultralytics Actions](https://github.com/ultralytics/actions)<sub>

### 🌟 Summary
Introducing expanded dataset exploration features for Ultralytics with an API and GUI support.

### 📊 Key Changes
- Added `explorer` dependency with vector search, SQL queries, and GUI elements.
- New `test_explorer.py` with tests for dataset exploration functionalities.
- `Explorer` class in `__init__.py` for the easy creation of dataset explorations.
- Updated `cfg/__init__.py` with an entry point for the dataset explorer.
- Implemented `ExplorerDataset` within `explorer.py` for loading images without alteration.
- Enabled semantic dataset search with the addition of `plot_similar` and `get_similar` methods.
- Introduced similarity index computation with the `similarity_index` method.
- Integrated GUI demo inside `dash.py`, providing interactive dataset exploration.
- Incorporated contribution to plotting within `plotting.py`, enabling visualizations of dataset insights.

### 🎯 Purpose & Impact
- **Dataset Insight**: Offers intuitive data analysis, helping users understand CV datasets better.
- **Efficiency**: Save time on data analysis thanks to seamless SQL, vector similarity, and semantic searches.
- **Enhanced UX**: Users can interactively explore datasets via included GUI, improving the overall experience.
- **Visualization Tools**: Rich visual representation of data made possible, aiding in pattern recognition and decision-making.

Problem statement:
Add new feature - dataset explorer
## Summary
Dataset explorer and utilities for Ultralytics

## Features

### Embedding

- Allows user to calculate embeddings as a top level operation in YOLO models.
- Start with detection/segmentation backbones for now, pose, etc. are not that relevant for embeddings.

<p><details>
<summary>Example</summary>

```py
from ultralytics import yolo

model = yolo()
embeddings = model.embed(inputs)
```
</details>

### Data explorer dashboard
  - Similarity search
    - Single image
    - Multiple images
  - SQL query/filtering

  <p><details>
  <summary>Example</summary>

  ```sh
  # CLI example
  yolo explorer {DATASET_NAME}
  ```
  ```py
  from ultralytics import Explorer
  
  exp = Explorer(“coco128.yaml”)
  exp.dashboard()
  ```
  </details>

### Advanced features

  - Search across multiple (selected) datasets
  - Delete redundant data or enrich primary dataset from secondary data sources
  - Upload to HUB and export new dataset

  <p><details>
  <summary>Example</summary>

  ```sh
  # CLI example
  yolo explorer {PRIMARY_DATASET} {SECONDARY_DATA_1} {SECONDARY_DATA_2} …
  ```
  ```py
  from ultralytics import Explorer

  exp_primary = Explorer(“coco128.yaml”) # Core dataset to clean or add to with images from other sources
  
  exp_1 = Explorer(“VOC.yaml”)
  exp_2 = Explorer(“coco.yaml”) 
  
  exp.dashboard(sources=[exp_1, exp_2])
  ```
  </details>

## Tasks
- [ ] Create dashboard
  - [ ] use Ultralytics branding
- [ ] Create utility functions/classes
- [ ] ...
- [ ] write unit test
- [ ] create Ultralytics Documentation

Hints from issue discussion:
@AyushExel I added the information you shared to this issue and cleaned it up some. I also tried to include sub-tasks, but wasn't certain on what else would be needed. Please reference this issue with your PR and any important commits.
PR #6971 opened to add this feature 
Cool. Advanced features are out of scope for now. Let's only plan on getting these before the anniversary:
* Embedding as a top level model feature
* Explorer pythonic API
* Explorer Dashboard

Commit message: Merge remote-tracking branch 'refs/remotes/origin/sem_search' into sem_search

Patch to review:
diff --git a/ultralytics/data/explorer/__init__.py b/ultralytics/data/explorer/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/ultralytics/data/explorer/explorer.py b/ultralytics/data/explorer/explorer.py
new file mode 100644
index 00000000000..0e80e2b0cb5
--- /dev/null
+++ b/ultralytics/data/explorer/explorer.py
@@ -0,0 +1,158 @@
+from pathlib import Path
+from typing import List
+
+import cv2
+import numpy as np
+import pyarrow as pa
+import torch
+from tqdm import tqdm
+
+from ultralytics import YOLO
+from ultralytics.data.augment import Format
+from ultralytics.data.dataset import YOLODataset
+from ultralytics.data.utils import check_det_dataset
+from ultralytics.utils import logger
+from ultralytics.utils.checks import check_requirements
+
+from .utils import sanitize_batch
+
+check_requirements('lancedb')
+import lancedb
+
+
+class ExplorerDataset(YOLODataset):
+
+    def __init__(self, *args, data=None, **kwargs):
+        super().__init__(*args, data=data, **kwargs)
+
+    # NOTE: Load the image directly without any resize operations.
+    def load_image(self, i):
+        """Loads 1 image from dataset index 'i', returns (im, resized hw)."""
+        im, f, fn = self.ims[i], self.im_files[i], self.npy_files[i]
+        if im is None:  # not cached in RAM
+            if fn.exists():  # load npy
+                im = np.load(fn)
+            else:  # read image
+                im = cv2.imread(f)  # BGR
+                if im is None:
+                    raise FileNotFoundError(f'Image Not Found {f}')
+            h0, w0 = im.shape[:2]  # orig hw
+            return im, (h0, w0), im.shape[:2]
+
+        return self.ims[i], self.im_hw0[i], self.im_hw[i]
+
+    def build_transforms(self, hyp=None):
+        transforms = Format(
+            bbox_format='xyxy',
+            normalize=False,
+            return_mask=self.use_segments,
+            return_keypoint=self.use_keypoints,
+            batch_idx=True,
+            mask_ratio=hyp.mask_ratio,
+            mask_overlap=hyp.overlap_mask,
+        )
+        return transforms
+
+
+class Explorer:
+
+    def __init__(self, data='coco128.yaml', model='yolov8n.pt', uri='~/ultralytics/explorer') -> None:
+        self.connection = lancedb.connect(uri)
+        self.table_name = Path(data).name
+        self.model = YOLO(model)
+        self.data = data  # None
+        self.choice_set = None
+
+        self.table = None
+
+    def create_embeddings_table(self, force=False, split='train', verbose=False):
+        if (self.table is not None and not force):
+            logger.info('Table already exists. Reusing it. Pass force=True to overwrite it.')
+            return
+        if self.table_name in self.connection.table_names():
+            logger.info(f'Table {self.table_name} already exists. Reusing it. Pass force=True to overwrite it.')
+            self.table = self.connection.open_table(self.table_name)
+            return
+        if self.data is None:
+            raise ValueError('Data must be provided to create embeddings table')
+
+        data_info = check_det_dataset(self.data)
+        if split not in data_info:
+            raise ValueError(
+                f'Split {split} is not found in the dataset. Available keys in the dataset are {list(data_info.keys())}'
+            )
+
+        choice_set = data_info[split]
+        choice_set = choice_set if isinstance(choice_set, list) else [choice_set]
+        self.choice_set = choice_set
+
+        dataset = ExplorerDataset(img_path=choice_set, data=data_info, augment=False, cache=False)
+
+        # Create the table schema
+        schema = pa.schema([
+            pa.field('im_file', pa.string()),
+            pa.field('labels', pa.list_(pa.string())),
+            pa.field('bboxes', pa.list_(pa.list_(pa.float32()))),
+            pa.field('cls', pa.list_(pa.int32())),
+            pa.field('vector', pa.list_(pa.float32(),
+                                        self.model.embed(dataset[0]['im_file'])[0].shape[0])), ])
+        table = self.connection.create_table(self.table_name, schema=schema, mode='overwrite')
+        table.add(
+            self._yeild_batches(dataset,
+                                data_info,
+                                self.model,
+                                exclude_keys=['img', 'ratio_pad', 'resized_shape', 'ori_shape', 'batch_idx']))
+
+        self.table = table
+
+    @staticmethod
+    def _yeild_batches(dataset, data_info, model, exclude_keys: List):
+        # Implement Batching
+        for i in tqdm(range(len(dataset))):
+            batch = dataset[i]
+            for k in exclude_keys:
+                batch.pop(k, None)
+            batch = sanitize_batch(batch, data_info)
+            batch['vector'] = model.embed(batch['im_file'], verbose=False)[0].detach().tolist()
+            yield [batch]
+
+    def query(self, img, limit=25):
+        """
+        Query the table for similar images. Accepts a single image or a list of images.
+
+        Args:
+            img (str or list): Path to the image or a list of paths to the images.
+            limit (int): Number of results to return.
+
+        Returns:
+            An arrow table containing the results.
+        """
+        if self.table is None:
+            raise ValueError('Table is not created. Please create the table first.')
+        if isinstance(img, str):
+            img = [img]
+        elif isinstance(img, list):
+            pass
+        else:
+            raise ValueError(f'img must be a string or a list of strings. Got {type(img)}')
+        embeds = self.model.embed(img)
+        # Get avg if multiple images are passed (len > 1)
+        embeds = torch.mean(torch.stack(embeds))
+
+        query = self.table.query(embeds).limit(limit).to_arrow()
+        return query
+
+    def sql_query(self, query):
+        """
+        Run a SQL-Like query on the table. Utilizes LanceDB predicate pushdown.
+
+        Args:
+            query (str): SQL query to run.
+
+        Returns:
+            An arrow table containing the results.
+        """
+        if self.table is None:
+            raise ValueError('Table is not created. Please create the table first.')
+
+        return self.table.to_lance.to_table(filter=query).to_arrow()
diff --git a/ultralytics/data/explorer/gui/__init__.py b/ultralytics/data/explorer/gui/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/ultralytics/data/explorer/utils.py b/ultralytics/data/explorer/utils.py
new file mode 100644
index 00000000000..209dae743e1
--- /dev/null
+++ b/ultralytics/data/explorer/utils.py
@@ -0,0 +1,7 @@
+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']]
+    return batch


Code review comments:
diff_hunk: @@ -0,0 +1,158 @@
+from pathlib import Path
+from typing import List
+
+import cv2
+import numpy as np
+import pyarrow as pa
+import torch
+from tqdm import tqdm
+
+from ultralytics import YOLO
+from ultralytics.data.augment import Format
+from ultralytics.data.dataset import YOLODataset
+from ultralytics.data.utils import check_det_dataset
+from ultralytics.utils import logger
+from ultralytics.utils.checks import check_requirements
+
+from .utils import sanitize_batch
+
+check_requirements('lancedb')
+import lancedb
+
+
+class ExplorerDataset(YOLODataset):
+
+    def __init__(self, *args, data=None, **kwargs):
+        super().__init__(*args, data=data, **kwargs)
+
+    # NOTE: Load the image directly without any resize operations.
+    def load_image(self, i):
+        """Loads 1 image from dataset index 'i', returns (im, resized hw)."""
+        im, f, fn = self.ims[i], self.im_files[i], self.npy_files[i]
+        if im is None:  # not cached in RAM
+            if fn.exists():  # load npy
+                im = np.load(fn)
+            else:  # read image
+                im = cv2.imread(f)  # BGR
+                if im is None:
+                    raise FileNotFoundError(f'Image Not Found {f}')
+            h0, w0 = im.shape[:2]  # orig hw
+            return im, (h0, w0), im.shape[:2]
+
+        return self.ims[i], self.im_hw0[i], self.im_hw[i]
+
+    def build_transforms(self, hyp=None):
+        transforms = Format(
+            bbox_format='xyxy',
+            normalize=False,
+            return_mask=self.use_segments,
+            return_keypoint=self.use_keypoints,
+            batch_idx=True,
+            mask_ratio=hyp.mask_ratio,
+            mask_overlap=hyp.overlap_mask,
+        )
+        return transforms
+
+
+class Explorer:
+
+    def __init__(self, data='coco128.yaml', model='yolov8n.pt', uri='~/ultralytics/explorer') -> None:
+        self.connection = lancedb.connect(uri)
+        self.table_name = Path(data).name
+        self.model = YOLO(model)
+        self.data = data  # None
+        self.choice_set = None
+
+        self.table = None
+
+    def create_embeddings_table(self, force=False, split='train', verbose=False):
+        if (self.table is not None and not force):
+            logger.info('Table already exists. Reusing it. Pass force=True to overwrite it.')
+            return
+        if self.table_name in self.connection.table_names():
+            logger.info(f'Table {self.table_name} already exists. Reusing it. Pass force=True to overwrite it.')
+            self.table = self.connection.open_table(self.table_name)
+            return
+        if self.data is None:
+            raise ValueError('Data must be provided to create embeddings table')
+
+        data_info = check_det_dataset(self.data)
+        if split not in data_info:
+            raise ValueError(
+                f'Split {split} is not found in the dataset. Available keys in the dataset are {list(data_info.keys())}'
+            )
+
+        choice_set = data_info[split]
+        choice_set = choice_set if isinstance(choice_set, list) else [choice_set]
+        self.choice_set = choice_set
+
+        dataset = ExplorerDataset(img_path=choice_set, data=data_info, augment=False, cache=False)
+
+        # Create the table schema
+        schema = pa.schema([
+            pa.field('im_file', pa.string()),
+            pa.field('labels', pa.list_(pa.string())),
+            pa.field('bboxes', pa.list_(pa.list_(pa.float32()))),
+            pa.field('cls', pa.list_(pa.int32())),
+            pa.field('vector', pa.list_(pa.float32(),
+                                        self.model.embed(dataset[0]['im_file'])[0].shape[0])), ])
+        table = self.connection.create_table(self.table_name, schema=schema, mode='overwrite')
+        table.add(
+            self._yeild_batches(dataset,
+                                data_info,
+                                self.model,
+                                exclude_keys=['img', 'ratio_pad', 'resized_shape', 'ori_shape', 'batch_idx']))
+
+        self.table = table
+
+    @staticmethod
+    def _yeild_batches(dataset, data_info, model, exclude_keys: List):
path: ultralytics/data/explorer/explorer.py
text: @user1:
Is there a typo? Did you mean `_yield_batches`?

@author:
Yes. This is still is WIP