diff --git a/.github/workflows/docs.yml b/.github/workflows/publish-dev-docs.yml similarity index 100% rename from .github/workflows/docs.yml rename to .github/workflows/publish-dev-docs.yml diff --git a/.github/workflows/publish-release-docs.yml b/.github/workflows/publish-release-docs.yml new file mode 100644 index 00000000..22710033 --- /dev/null +++ b/.github/workflows/publish-release-docs.yml @@ -0,0 +1,55 @@ +name: Supervision Release Documentation Workflow 📚 +on: + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.event_name == 'push' && github.ref}} + cancel-in-progress: true + +permissions: + contents: write + pages: write + pull-requests: write + + +jobs: + doc-build-deploy: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.10"] + steps: + - name: 🛎️ Checkout + uses: actions/checkout@v4 + with: + fetch-depth: 0 + ref: ${{ github.head_ref }} + + - name: 🐍 Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + - name: 📦 Install mkdocs-material + run: pip install "mkdocs-material" + - name: 📦 Install mkdocstrings[python] + run: pip install "mkdocstrings[python]" + - name: 📦 Install mkdocs-material[imaging] + run: pip install "mkdocs-material[imaging]" + - name: 📦 Install mike + run: pip install "mike" + - name: 📦 Install mkdocs-git-revision-date-localized-plugin + run: pip install "mkdocs-git-revision-date-localized-plugin" + - name: 📦 Install JupyterLab + run: pip install jupyterlab + - name: 📦 Install mkdocs-jupyter + run: pip install mkdocs-jupyter + - name: 📦 Install mkdocs-git-committers-plugin-2 + run: pip install mkdocs-git-committers-plugin-2 + - name: ⚙️ Configure git for github-actions 👷 + run: | + git config --global user.name "github-actions[bot]" + git config --global user.email "41898282+github-actions[bot]@users.noreply.github.com" + - name: 🚀 Deploy MkDoc-Material 📚 + run: | + latest_tag=$(git describe --tags `git rev-list --tags --max-count=1`) + MKDOCS_GIT_COMMITTERS_APIKEY=${{ secrets.GITHUB_TOKEN }} mike deploy --push --update-aliases $latest_tag latest diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 04f21d15..b016ca8d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -7,7 +7,7 @@ ci: repos: - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.5.0 + rev: v4.6.0 hooks: - id: end-of-file-fixer - id: trailing-whitespace @@ -27,7 +27,6 @@ repos: - id: mixed-line-ending - - repo: https://github.com/PyCQA/bandit rev: '1.7.8' hooks: @@ -46,7 +45,7 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.3.4 + rev: v0.3.7 hooks: - id: ruff args: [--fix, --exit-non-zero-on-fix] diff --git a/README.md b/README.md index 7d280b57..beed655c 100644 --- a/README.md +++ b/README.md @@ -217,19 +217,23 @@ len(dataset) ## 🎬 tutorials -

-Speed Estimation & Vehicle Tracking | Computer Vision | Open Source -Speed Estimation & Vehicle Tracking | Computer Vision | Open Source -

Created: 11 Jan 2024 | Updated: 11 Jan 2024
-
Learn how to track and estimate the speed of vehicles using YOLO, ByteTrack, and Roboflow Inference. This comprehensive tutorial covers object detection, multi-object tracking, filtering detections, perspective transformation, speed estimation, visualization improvements, and more.

+Want to learn how to use Supervision? Explore our [how-to guides](https://supervision.roboflow.com/develop/how_to/detect_and_annotate/), [end-to-end examples](https://github.com/roboflow/supervision/tree/develop/examples), and [cookbooks](https://supervision.roboflow.com/develop/cookbooks/)!

-Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking -Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking -

Created: 6 Sep 2023 | Updated: 6 Sep 2023
-
In this video, we explore real-time traffic analysis using YOLOv8 and ByteTrack to detect and track vehicles on aerial images. Harnessing the power of Python and Supervision, we delve deep into assigning cars to specific entry zones and understanding their direction of movement. By visualizing their paths, we gain insights into traffic flow across bustling roundabouts...

+Dwell Time Analysis with Computer Vision | Real-Time Stream Processing +Dwell Time Analysis with Computer Vision | Real-Time Stream Processing +
Created: 5 Apr 2024
+
Learn how to use computer vision to analyze wait times and optimize processes. This tutorial covers object detection, tracking, and calculating time spent in designated zones. Use these techniques to improve customer experience in retail, traffic management, or other scenarios.

+ +
+ +

+Speed Estimation & Vehicle Tracking | Computer Vision | Open Source +Speed Estimation & Vehicle Tracking | Computer Vision | Open Source +

Created: 11 Jan 2024
+
Learn how to track and estimate the speed of vehicles using YOLO, ByteTrack, and Roboflow Inference. This comprehensive tutorial covers object detection, multi-object tracking, filtering detections, perspective transformation, speed estimation, visualization improvements, and more.

## 💜 built with supervision diff --git a/docs/how_to/detect_and_annotate.md b/docs/how_to/detect_and_annotate.md index 221d9ca1..adea95cf 100644 --- a/docs/how_to/detect_and_annotate.md +++ b/docs/how_to/detect_and_annotate.md @@ -15,7 +15,7 @@ source image. ![basic-annotation](https://media.roboflow.com/supervision_detect_and_annotate_example_1.png) -## Run Inference +## Run Detection First, you'll need to obtain predictions from your object detection or segmentation model. diff --git a/docs/how_to/save_detections.md b/docs/how_to/save_detections.md new file mode 100644 index 00000000..94de6c61 --- /dev/null +++ b/docs/how_to/save_detections.md @@ -0,0 +1,299 @@ +--- +comments: true +status: new +--- + +# Save Detections + +Supervision enables an easy way to save detections in .CSV and .JSON files for offline +processing. This guide demonstrates how to perform video inference using the +[Inference](https://github.com/roboflow/inference), +[Ultralytics](https://github.com/ultralytics/ultralytics) or +[Transformers](https://github.com/huggingface/transformers) packages and save their results with +[`sv.CSVSink`](/latest/detection/tools/save_detections/#supervision.detection.tools.csv_sink.CSVSink) and +[`sv.JSONSink`](/latest/detection/tools/save_detections/#supervision.detection.tools.csv_sink.JSONSink). + +## Run Detection + +First, you'll need to obtain predictions from your object detection or segmentation +model. You can learn more on this topic in our +[How to Detect and Annotate](/latest/how_to/detect_and_annotate.md) guide. + +=== "Inference" + + ```python + import supervision as sv + from inference import get_model + + model = get_model(model_id="yolov8n-640") + frames_generator = sv.get_video_frames_generator() + + for frame in frames_generator: + + results = model.infer(image)[0] + detections = sv.Detections.from_inference(results) + ``` + +=== "Ultralytics" + + ```python + import supervision as sv + from ultralytics import YOLO + + model = YOLO("yolov8n.pt") + frames_generator = sv.get_video_frames_generator() + + for frame in frames_generator: + + results = model(frame)[0] + detections = sv.Detections.from_ultralytics(results) + ``` + +=== "Transformers" + + ```python + import torch + import supervision as sv + from transformers import DetrImageProcessor, DetrForObjectDetection + + processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50") + model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50") + frames_generator = sv.get_video_frames_generator() + + for frame in frames_generator: + + frame = sv.cv2_to_pillow(frame) + inputs = processor(images=frame, return_tensors="pt") + + with torch.no_grad(): + outputs = model(**inputs) + + width, height = frame.size + target_size = torch.tensor([[height, width]]) + results = processor.post_process_object_detection( + outputs=outputs, target_sizes=target_size)[0] + detections = sv.Detections.from_transformers(results) + ``` + +## Save Detections as CSV + +To save detections to a `.CSV` file, open our +[`sv.CSVSink`](/latest/detection/tools/save_detections/#supervision.detection.tools.csv_sink.CSVSink) +and then pass the +[`sv.Detections`](/latest/detection/core/#supervision.detection.core.Detections) +object resulting from the inference to it. Its fields are parsed and saved on disk. + +=== "Inference" + + ```{ .py hl_lines="7 12" } + import supervision as sv + from inference import get_model + + model = get_model(model_id="yolov8n-640") + frames_generator = sv.get_video_frames_generator() + + with sv.CSVSink() as sink: + for frame in frames_generator: + + results = model.infer(image)[0] + detections = sv.Detections.from_inference(results) + sink.append(detections, {}) + ``` + +=== "Ultralytics" + + ```{ .py hl_lines="7 12" } + import supervision as sv + from ultralytics import YOLO + + model = YOLO("yolov8n.pt") + frames_generator = sv.get_video_frames_generator() + + with sv.CSVSink() as sink: + for frame in frames_generator: + + results = model(frame)[0] + detections = sv.Detections.from_ultralytics(results) + sink.append(detections, {}) + ``` + +=== "Transformers" + + ```{ .py hl_lines="9 23" } + import torch + import supervision as sv + from transformers import DetrImageProcessor, DetrForObjectDetection + + processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50") + model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50") + frames_generator = sv.get_video_frames_generator() + + with sv.CSVSink() as sink: + for frame in frames_generator: + + frame = sv.cv2_to_pillow(frame) + inputs = processor(images=frame, return_tensors="pt") + + with torch.no_grad(): + outputs = model(**inputs) + + width, height = frame.size + target_size = torch.tensor([[height, width]]) + results = processor.post_process_object_detection( + outputs=outputs, target_sizes=target_size)[0] + detections = sv.Detections.from_transformers(results) + sink.append(detections, {}) + ``` + +| x_min | y_min | x_max | y_max | class_id | confidence | tracker_id | class_name | +|---------|----------|---------|----------|----------|------------|------------|------------| +| 2941.14 | 1269.31 | 3220.77 | 1500.67 | 2 | 0.8517 | | car | +| 944.889 | 899.641 | 1235.42 | 1308.80 | 7 | 0.6752 | | truck | +| 1439.78 | 1077.79 | 1621.27 | 1231.40 | 2 | 0.6450 | | car | + +## Custom Fields + +Besides regular fields in +[`sv.Detections`](/latest/detection/core/#supervision.detection.core.Detections), +[`sv.CSVSink`](/latest/detection/tools/save_detections/#supervision.detection.tools.csv_sink.CSVSink) +also allows you to add custom information to each row, which can be passed via the +`custom_data` dictionary. Let's utilize this feature to save information about the +frame index from which the detections originate. + +=== "Inference" + + ```{ .py hl_lines="8 12" } + import supervision as sv + from inference import get_model + + model = get_model(model_id="yolov8n-640") + frames_generator = sv.get_video_frames_generator() + + with sv.CSVSink() as sink: + for frame_index, frame in enumerate(frames_generator): + + results = model.infer(image)[0] + detections = sv.Detections.from_inference(results) + sink.append(detections, {"frame_index": frame_index}) + ``` + +=== "Ultralytics" + + ```{ .py hl_lines="8 12" } + import supervision as sv + from ultralytics import YOLO + + model = YOLO("yolov8n.pt") + frames_generator = sv.get_video_frames_generator() + + with sv.CSVSink() as sink: + for frame_index, frame in enumerate(frames_generator): + + results = model(frame)[0] + detections = sv.Detections.from_ultralytics(results) + sink.append(detections, {"frame_index": frame_index}) + ``` + +=== "Transformers" + + ```{ .py hl_lines="10 23" } + import torch + import supervision as sv + from transformers import DetrImageProcessor, DetrForObjectDetection + + processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50") + model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50") + frames_generator = sv.get_video_frames_generator() + + with sv.CSVSink() as sink: + for frame_index, frame in enumerate(frames_generator): + + frame = sv.cv2_to_pillow(frame) + inputs = processor(images=frame, return_tensors="pt") + + with torch.no_grad(): + outputs = model(**inputs) + + width, height = frame.size + target_size = torch.tensor([[height, width]]) + results = processor.post_process_object_detection( + outputs=outputs, target_sizes=target_size)[0] + detections = sv.Detections.from_transformers(results) + sink.append(detections, {"frame_index": frame_index}) + ``` + +| x_min | y_min | x_max | y_max | class_id | confidence | tracker_id | class_name | frame_index | +|---------|----------|---------|----------|----------|------------|------------|------------|-------------| +| 2941.14 | 1269.31 | 3220.77 | 1500.67 | 2 | 0.8517 | | car | 0 | +| 944.889 | 899.641 | 1235.42 | 1308.80 | 7 | 0.6752 | | truck | 0 | +| 1439.78 | 1077.79 | 1621.27 | 1231.40 | 2 | 0.6450 | | car | 0 | + +## Save Detections as JSON + +If you prefer to save the result in a `.JSON` file instead of a `.CSV` file, all you +need to do is replace +[`sv.CSVSink`](/latest/detection/tools/save_detections/#supervision.detection.tools.csv_sink.CSVSink) +with +[`sv.JSONSink`](/latest/detection/tools/save_detections/#supervision.detection.tools.csv_sink.JSONSink). + +=== "Inference" + + ```{ .py hl_lines="7" } + import supervision as sv + from inference import get_model + + model = get_model(model_id="yolov8n-640") + frames_generator = sv.get_video_frames_generator() + + with sv.JSONSink() as sink: + for frame_index, frame in enumerate(frames_generator): + + results = model.infer(image)[0] + detections = sv.Detections.from_inference(results) + sink.append(detections, {"frame_index": frame_index}) + ``` + +=== "Ultralytics" + + ```{ .py hl_lines="7" } + import supervision as sv + from ultralytics import YOLO + + model = YOLO("yolov8n.pt") + frames_generator = sv.get_video_frames_generator() + + with sv.JSONSink() as sink: + for frame_index, frame in enumerate(frames_generator): + + results = model(frame)[0] + detections = sv.Detections.from_ultralytics(results) + sink.append(detections, {"frame_index": frame_index}) + ``` + +=== "Transformers" + + ```{ .py hl_lines="9" } + import torch + import supervision as sv + from transformers import DetrImageProcessor, DetrForObjectDetection + + processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50") + model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50") + frames_generator = sv.get_video_frames_generator() + + with sv.JSONSink() as sink: + for frame_index, frame in enumerate(frames_generator): + + frame = sv.cv2_to_pillow(frame) + inputs = processor(images=frame, return_tensors="pt") + + with torch.no_grad(): + outputs = model(**inputs) + + width, height = frame.size + target_size = torch.tensor([[height, width]]) + results = processor.post_process_object_detection( + outputs=outputs, target_sizes=target_size)[0] + detections = sv.Detections.from_transformers(results) + sink.append(detections, {"frame_index": frame_index}) + ``` diff --git a/docs/utils/image.md b/docs/utils/image.md index 8a6768d7..87fd1943 100644 --- a/docs/utils/image.md +++ b/docs/utils/image.md @@ -5,18 +5,18 @@ status: new # Image Utils -
-

ImageSink

-
- -:::supervision.utils.image.ImageSink -

crop_image

:::supervision.utils.image.crop_image +
+

scale_image

+
+ +:::supervision.utils.image.scale_image +

resize_image

@@ -24,7 +24,19 @@ status: new :::supervision.utils.image.resize_image
-

place_image

+

letterbox_image

-:::supervision.utils.image.place_image +:::supervision.utils.image.letterbox_image + +
+

overlay_image

+
+ +:::supervision.utils.image.overlay_image + +
+

ImageSink

+
+ +:::supervision.utils.image.ImageSink diff --git a/docs/utils/iterables.md b/docs/utils/iterables.md new file mode 100644 index 00000000..3ccf98e5 --- /dev/null +++ b/docs/utils/iterables.md @@ -0,0 +1,18 @@ +--- +comments: true +status: new +--- + +# Iterables Utils + +
+

create_batches

+
+ +:::supervision.utils.iterables.create_batches + +
+

fill

+
+ +:::supervision.utils.iterables.fill diff --git a/docs/utils/notebook.md b/docs/utils/notebook.md index d58bbcfc..d09b9bbd 100644 --- a/docs/utils/notebook.md +++ b/docs/utils/notebook.md @@ -11,7 +11,7 @@ comments: true :::supervision.utils.notebook.plot_image
-

## plot_images_grid +

plot_images_grid

diff --git a/examples/README.md b/examples/README.md index 81c8054b..b723a247 100644 --- a/examples/README.md +++ b/examples/README.md @@ -1,8 +1,8 @@ # Examples -This repository is packed with real-world use-cases, provided through Python scripts or -interactive notebooks. Browse through to understand how the Supervision library -interfaces with diverse applications. +Here, you'll find end-to-end examples that show how to solve common computer vision problems using Supervision. + +For more information and examples, visit our [documentation](https://supervision.roboflow.com/develop/annotators/) and explore our [how-to guides](https://supervision.roboflow.com/develop/how_to/detect_and_annotate/) and [cookbooks](https://supervision.roboflow.com/develop/cookbooks/). Join our [Discord](https://discord.com/invite/GbfgXGJ8Bk) and meet other Supervision power users! - [tracking](./tracking) by [@SkalskiP](https://github.com/SkalskiP) - [count people in zone](./count_people_in_zone) by [@SkalskiP](https://github.com/SkalskiP) @@ -10,23 +10,3 @@ interfaces with diverse applications. - [speed estimation](./speed_estimation) by [@SkalskiP](https://github.com/SkalskiP) - [time in zone](./time_in_zone) by [@SkalskiP](https://github.com/SkalskiP) - [heatmap and track](./heatmap_and_track/) by [@HinePo](https://github.com/HinePo) - -## Contributing - -We welcome contributions from the community in the form of examples, applications, and -guides. To contribute, please follow these steps: - -1. Create a pull request (PR) with the `[Example]` prefix in the title, adding your -project folder to the `examples/` directory in the repository. -2. Confirm your project aligns with the following standards: - - Incorporates the `supervision` package. - - Provides a `README.md` file, detailing the instructions to execute the project. - - Showcases visual results, demonstrating the app's functionality. - - Avoids adding large assets or dependencies unless absolutely necessary. - - The contributor is expected to provide support for issues related to their -examples. - - In case the presented model has licensing complications, kindly specify them to -circumvent potential misunderstandings. - -For inquiries or concerns about these prerequisites, feel free to raise a PR. We are -committed to assist and guide you. diff --git a/examples/count_people_in_zone/README.md b/examples/count_people_in_zone/README.md index 2e14a80b..d9d3cdb2 100644 --- a/examples/count_people_in_zone/README.md +++ b/examples/count_people_in_zone/README.md @@ -1,5 +1,8 @@ # count people in zone +[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-detect-and-count-objects-in-polygon-zone.ipynb) +[![YouTube](https://badges.aleen42.com/src/youtube.svg)](https://www.youtube.com/watch?v=l_kf9CfZ_8M) + ## 👋 hello This demo is a video analysis tool that counts and highlights objects in specific zones diff --git a/examples/time_in_zone/README.md b/examples/time_in_zone/README.md index 2bfd6c7a..98587999 100644 --- a/examples/time_in_zone/README.md +++ b/examples/time_in_zone/README.md @@ -1,5 +1,7 @@ # time in zone +[![YouTube](https://badges.aleen42.com/src/youtube.svg)](https://www.youtube.com/watch?v=hAWpsIuem10) + ## 👋 hello Practical demonstration on leveraging computer vision for analyzing wait times and diff --git a/mkdocs.yml b/mkdocs.yml index d325eb65..dd6feb8d 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -38,9 +38,11 @@ nav: - Home: index.md - How to: - Detect and Annotate: how_to/detect_and_annotate.md + - Save Detections: how_to/save_detections.md + - Filter Detections: how_to/filter_detections.md - Detect Small Objects: how_to/detect_small_objects.md - Track Objects: how_to/track_objects.md - - Filter Detections: how_to/filter_detections.md + - API: - Annotators: annotators.md - Classifications: @@ -66,6 +68,7 @@ nav: - Utils: - Video: utils/video.md - Image: utils/image.md + - Iterables: utils/iterables.md - Notebook: utils/notebook.md - File: utils/file.md - Assets: assets.md diff --git a/poetry.lock b/poetry.lock index 5fa0fbb2..95cbef01 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1122,13 +1122,13 @@ license = ["ukkonen"] [[package]] name = "idna" -version = "3.6" +version = "3.7" description = "Internationalized Domain Names in Applications (IDNA)" optional = false python-versions = ">=3.5" files = [ - {file = "idna-3.6-py3-none-any.whl", hash = "sha256:c05567e9c24a6b9faaa835c4821bad0590fbb9d5779e7caa6e1cc4978e7eb24f"}, - {file = "idna-3.6.tar.gz", hash = "sha256:9ecdbbd083b06798ae1e86adcbfe8ab1479cf864e4ee30fe4e46a003d12491ca"}, + {file = "idna-3.7-py3-none-any.whl", hash = "sha256:82fee1fc78add43492d3a1898bfa6d8a904cc97d8427f683ed8e798d07761aa0"}, + {file = "idna-3.7.tar.gz", hash = "sha256:028ff3aadf0609c1fd278d8ea3089299412a7a8b9bd005dd08b9f8285bcb5cfc"}, ] [[package]] @@ -2162,13 +2162,12 @@ pytz = "*" [[package]] name = "mkdocs-jupyter" -version = "0.24.3" +version = "0.24.7" description = "Use Jupyter in mkdocs websites" optional = false -python-versions = ">=3.7" +python-versions = ">=3.8" files = [ - {file = "mkdocs_jupyter-0.24.3-py3-none-any.whl", hash = "sha256:904262a8678a5e5920b7c3c03b5010b36301a69d0a38f2fcbf430493adf6879e"}, - {file = "mkdocs_jupyter-0.24.3.tar.gz", hash = "sha256:3d81da9aea27480e93bab22438910c4f0b9630613e74f85b576590d78e0e8b14"}, + {file = "mkdocs_jupyter-0.24.7-py3-none-any.whl", hash = "sha256:893d04bea1e007479a46e4e72852cd4d280c4d358ce4a0445250f3f80c639723"}, ] [package.dependencies] @@ -2179,18 +2178,15 @@ mkdocs-material = ">9.0.0" nbconvert = ">=7.2.9,<8" pygments = ">2.12.0" -[package.extras] -test = ["coverage[toml]", "pymdown-extensions", "pytest", "pytest-cov"] - [[package]] name = "mkdocs-material" -version = "9.5.17" +version = "9.5.18" description = "Documentation that simply works" optional = false python-versions = ">=3.8" files = [ - {file = "mkdocs_material-9.5.17-py3-none-any.whl", hash = "sha256:14a2a60119a785e70e765dd033e6211367aca9fc70230e577c1cf6a326949571"}, - {file = "mkdocs_material-9.5.17.tar.gz", hash = "sha256:06ae1275a72db1989cf6209de9e9ecdfbcfdbc24c58353877b2bb927dbe413e4"}, + {file = "mkdocs_material-9.5.18-py3-none-any.whl", hash = "sha256:1e0e27fc9fe239f9064318acf548771a4629d5fd5dfd45444fd80a953fe21eb4"}, + {file = "mkdocs_material-9.5.18.tar.gz", hash = "sha256:a43f470947053fa2405c33995f282d24992c752a50114f23f30da9d8d0c57e62"}, ] [package.dependencies] @@ -2226,13 +2222,13 @@ files = [ [[package]] name = "mkdocstrings" -version = "0.24.2" +version = "0.24.3" description = "Automatic documentation from sources, for MkDocs." optional = false python-versions = ">=3.8" files = [ - {file = "mkdocstrings-0.24.2-py3-none-any.whl", hash = "sha256:61440b77542170238099a7d87882c3417897771950e3aafe6e22abff3d1c51fb"}, - {file = "mkdocstrings-0.24.2.tar.gz", hash = "sha256:b91b9cdd9490ef2e8957000bff1d34a4b308b9cd57b10f26169f085def4c6a92"}, + {file = "mkdocstrings-0.24.3-py3-none-any.whl", hash = "sha256:5c9cf2a32958cd161d5428699b79c8b0988856b0d4a8c5baf8395fc1bf4087c3"}, + {file = "mkdocstrings-0.24.3.tar.gz", hash = "sha256:f327b234eb8d2551a306735436e157d0a22d45f79963c60a8b585d5f7a94c1d2"}, ] [package.dependencies] @@ -2470,13 +2466,13 @@ setuptools = "*" [[package]] name = "notebook" -version = "7.1.2" +version = "7.1.3" description = "Jupyter Notebook - A web-based notebook environment for interactive computing" optional = false python-versions = ">=3.8" files = [ - {file = "notebook-7.1.2-py3-none-any.whl", hash = "sha256:fc6c24b9aef18d0cd57157c9c47e95833b9b0bdc599652639acf0bdb61dc7d5f"}, - {file = "notebook-7.1.2.tar.gz", hash = "sha256:efc2c80043909e0faa17fce9e9b37c059c03af0ec99a4d4db84cb21d9d2e936a"}, + {file = "notebook-7.1.3-py3-none-any.whl", hash = "sha256:919b911e59f41f6e3857ce93c9d93535ba66bb090059712770e5968c07e1004d"}, + {file = "notebook-7.1.3.tar.gz", hash = "sha256:41fcebff44cf7bb9377180808bcbae066629b55d8c7722f1ebbe75ca44f9cfc1"}, ] [package.dependencies] @@ -3647,28 +3643,28 @@ files = [ [[package]] name = "ruff" -version = "0.3.5" +version = "0.4.1" description = "An extremely fast Python linter and code formatter, written in Rust." optional = false python-versions = ">=3.7" files = [ - {file = "ruff-0.3.5-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:aef5bd3b89e657007e1be6b16553c8813b221ff6d92c7526b7e0227450981eac"}, - {file = "ruff-0.3.5-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:89b1e92b3bd9fca249153a97d23f29bed3992cff414b222fcd361d763fc53f12"}, - {file = "ruff-0.3.5-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5e55771559c89272c3ebab23326dc23e7f813e492052391fe7950c1a5a139d89"}, - {file = "ruff-0.3.5-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:dabc62195bf54b8a7876add6e789caae0268f34582333cda340497c886111c39"}, - {file = "ruff-0.3.5-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3a05f3793ba25f194f395578579c546ca5d83e0195f992edc32e5907d142bfa3"}, - {file = "ruff-0.3.5-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:dfd3504e881082959b4160ab02f7a205f0fadc0a9619cc481982b6837b2fd4c0"}, - {file = "ruff-0.3.5-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:87258e0d4b04046cf1d6cc1c56fadbf7a880cc3de1f7294938e923234cf9e498"}, - {file = "ruff-0.3.5-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:712e71283fc7d9f95047ed5f793bc019b0b0a29849b14664a60fd66c23b96da1"}, - {file = "ruff-0.3.5-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a532a90b4a18d3f722c124c513ffb5e5eaff0cc4f6d3aa4bda38e691b8600c9f"}, - {file = "ruff-0.3.5-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:122de171a147c76ada00f76df533b54676f6e321e61bd8656ae54be326c10296"}, - {file = "ruff-0.3.5-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:d80a6b18a6c3b6ed25b71b05eba183f37d9bc8b16ace9e3d700997f00b74660b"}, - {file = "ruff-0.3.5-py3-none-musllinux_1_2_i686.whl", hash = "sha256:a7b6e63194c68bca8e71f81de30cfa6f58ff70393cf45aab4c20f158227d5936"}, - {file = "ruff-0.3.5-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:a759d33a20c72f2dfa54dae6e85e1225b8e302e8ac655773aff22e542a300985"}, - {file = "ruff-0.3.5-py3-none-win32.whl", hash = "sha256:9d8605aa990045517c911726d21293ef4baa64f87265896e491a05461cae078d"}, - {file = "ruff-0.3.5-py3-none-win_amd64.whl", hash = "sha256:dc56bb16a63c1303bd47563c60482a1512721053d93231cf7e9e1c6954395a0e"}, - {file = "ruff-0.3.5-py3-none-win_arm64.whl", hash = "sha256:faeeae9905446b975dcf6d4499dc93439b131f1443ee264055c5716dd947af55"}, - {file = "ruff-0.3.5.tar.gz", hash = "sha256:a067daaeb1dc2baf9b82a32dae67d154d95212080c80435eb052d95da647763d"}, + {file = "ruff-0.4.1-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:2d9ef6231e3fbdc0b8c72404a1a0c46fd0dcea84efca83beb4681c318ea6a953"}, + {file = "ruff-0.4.1-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:9485f54a7189e6f7433e0058cf8581bee45c31a25cd69009d2a040d1bd4bfaef"}, + {file = "ruff-0.4.1-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d2921ac03ce1383e360e8a95442ffb0d757a6a7ddd9a5be68561a671e0e5807e"}, + {file = "ruff-0.4.1-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:eec8d185fe193ad053eda3a6be23069e0c8ba8c5d20bc5ace6e3b9e37d246d3f"}, + {file = "ruff-0.4.1-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:baa27d9d72a94574d250f42b7640b3bd2edc4c58ac8ac2778a8c82374bb27984"}, + {file = "ruff-0.4.1-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:f1ee41580bff1a651339eb3337c20c12f4037f6110a36ae4a2d864c52e5ef954"}, + {file = "ruff-0.4.1-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:0926cefb57fc5fced629603fbd1a23d458b25418681d96823992ba975f050c2b"}, + {file = "ruff-0.4.1-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2c6e37f2e3cd74496a74af9a4fa67b547ab3ca137688c484749189bf3a686ceb"}, + {file = "ruff-0.4.1-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:efd703a5975ac1998c2cc5e9494e13b28f31e66c616b0a76e206de2562e0843c"}, + {file = "ruff-0.4.1-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:b92f03b4aa9fa23e1799b40f15f8b95cdc418782a567d6c43def65e1bbb7f1cf"}, + {file = "ruff-0.4.1-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:1c859f294f8633889e7d77de228b203eb0e9a03071b72b5989d89a0cf98ee262"}, + {file = "ruff-0.4.1-py3-none-musllinux_1_2_i686.whl", hash = "sha256:b34510141e393519a47f2d7b8216fec747ea1f2c81e85f076e9f2910588d4b64"}, + {file = "ruff-0.4.1-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:6e68d248ed688b9d69fd4d18737edcbb79c98b251bba5a2b031ce2470224bdf9"}, + {file = "ruff-0.4.1-py3-none-win32.whl", hash = "sha256:b90506f3d6d1f41f43f9b7b5ff845aeefabed6d2494307bc7b178360a8805252"}, + {file = "ruff-0.4.1-py3-none-win_amd64.whl", hash = "sha256:c7d391e5936af5c9e252743d767c564670dc3889aff460d35c518ee76e4b26d7"}, + {file = "ruff-0.4.1-py3-none-win_arm64.whl", hash = "sha256:a1eaf03d87e6a7cd5e661d36d8c6e874693cb9bc3049d110bc9a97b350680c43"}, + {file = "ruff-0.4.1.tar.gz", hash = "sha256:d592116cdbb65f8b1b7e2a2b48297eb865f6bdc20641879aa9d7b9c11d86db79"}, ] [[package]] diff --git a/supervision/__init__.py b/supervision/__init__.py index 35f525b0..f8e7a832 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -71,8 +71,17 @@ from supervision.geometry.core import Point, Position, Rect from supervision.geometry.utils import get_polygon_center from supervision.metrics.detection import ConfusionMatrix, MeanAveragePrecision from supervision.tracker.byte_tracker.core import ByteTrack +from supervision.utils.conversion import cv2_to_pillow, pillow_to_cv2 from supervision.utils.file import list_files_with_extensions -from supervision.utils.image import ImageSink, crop_image, place_image, resize_image +from supervision.utils.image import ( + ImageSink, + create_tiles, + crop_image, + letterbox_image, + overlay_image, + resize_image, + scale_image, +) from supervision.utils.notebook import plot_image, plot_images_grid from supervision.utils.video import ( FPSMonitor, diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index a7700117..ac901862 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -5,19 +5,15 @@ import cv2 import numpy as np from supervision.annotators.base import BaseAnnotator, ImageType -from supervision.annotators.utils import ( - ColorLookup, - Trace, - resolve_color, - scene_to_annotator_img_type, -) +from supervision.annotators.utils import ColorLookup, Trace, resolve_color from supervision.config import CLASS_NAME_DATA_FIELD, ORIENTED_BOX_COORDINATES from supervision.detection.core import Detections from supervision.detection.utils import clip_boxes, mask_to_polygons from supervision.draw.color import Color, ColorPalette from supervision.draw.utils import draw_polygon from supervision.geometry.core import Position -from supervision.utils.image import crop_image, place_image, resize_image +from supervision.utils.conversion import convert_for_annotation_method +from supervision.utils.image import crop_image, overlay_image, scale_image class BoundingBoxAnnotator(BaseAnnotator): @@ -43,7 +39,7 @@ class BoundingBoxAnnotator(BaseAnnotator): self.thickness: int = thickness self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -124,7 +120,7 @@ class OrientedBoxAnnotator(BaseAnnotator): self.thickness: int = thickness self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -212,7 +208,7 @@ class MaskAnnotator(BaseAnnotator): self.opacity = opacity self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -299,7 +295,7 @@ class PolygonAnnotator(BaseAnnotator): self.thickness: int = thickness self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -385,7 +381,7 @@ class ColorAnnotator(BaseAnnotator): self.color_lookup: ColorLookup = color_lookup self.opacity = opacity - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -479,7 +475,7 @@ class HaloAnnotator(BaseAnnotator): self.color_lookup: ColorLookup = color_lookup self.kernel_size: int = kernel_size - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -577,7 +573,7 @@ class EllipseAnnotator(BaseAnnotator): self.end_angle: int = end_angle self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -668,7 +664,7 @@ class BoxCornerAnnotator(BaseAnnotator): self.corner_length: int = corner_length self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -756,7 +752,7 @@ class CircleAnnotator(BaseAnnotator): self.thickness: int = thickness self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -846,7 +842,7 @@ class DotAnnotator(BaseAnnotator): self.position: Position = position self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -993,7 +989,7 @@ class LabelAnnotator: center_y + text_h // 2, ) - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -1164,7 +1160,7 @@ class BlurAnnotator(BaseAnnotator): """ self.kernel_size: int = kernel_size - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -1249,7 +1245,7 @@ class TraceAnnotator: self.thickness = thickness self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -1356,7 +1352,7 @@ class HeatMapAnnotator: self.top_hue = top_hue self.low_hue = low_hue - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate(self, scene: ImageType, detections: Detections) -> ImageType: """ Annotates the scene with a heatmap based on the provided detections. @@ -1432,7 +1428,7 @@ class PixelateAnnotator(BaseAnnotator): """ self.pixel_size: int = pixel_size - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -1520,7 +1516,7 @@ class TriangleAnnotator(BaseAnnotator): self.position: Position = position self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -1616,7 +1612,7 @@ class RoundBoxAnnotator(BaseAnnotator): raise ValueError("roundness attribute must be float between (0, 1.0]") self.roundness: float = roundness - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -1753,7 +1749,7 @@ class PercentageBarAnnotator(BaseAnnotator): if border_thickness is None: self.border_thickness = int(0.15 * self.height) - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -1926,7 +1922,7 @@ class CropAnnotator(BaseAnnotator): self.border_thickness: int = border_thickness self.border_color_lookup: ColorLookup = border_color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -1969,7 +1965,7 @@ class CropAnnotator(BaseAnnotator): crop_image(image=scene, xyxy=xyxy) for xyxy in detections.xyxy.astype(int) ] resized_crops = [ - resize_image(image=crop, scale_factor=self.scale_factor) for crop in crops + scale_image(image=crop, scale_factor=self.scale_factor) for crop in crops ] anchors = detections.get_anchors_coordinates(anchor=self.position).astype(int) @@ -1978,7 +1974,9 @@ class CropAnnotator(BaseAnnotator): (x1, y1), (x2, y2) = self.calculate_crop_coordinates( anchor=anchor, crop_wh=crop_wh, position=self.position ) - scene = place_image(scene=scene, image=resized_crop, anchor=(x1, y1)) + scene = overlay_image( + scene=scene, inserted_image=resized_crop, anchor=(x1, y1) + ) color = resolve_color( color=self.border_color, detections=detections, diff --git a/supervision/annotators/utils.py b/supervision/annotators/utils.py index 6f9cd9db..e206c8cb 100644 --- a/supervision/annotators/utils.py +++ b/supervision/annotators/utils.py @@ -1,12 +1,8 @@ from enum import Enum -from functools import wraps from typing import Optional, Union -import cv2 import numpy as np -from PIL import Image -from supervision.annotators.base import ImageType from supervision.detection.core import Detections from supervision.draw.color import Color, ColorPalette from supervision.geometry.core import Position @@ -123,33 +119,3 @@ class Trace: def get(self, tracker_id: int) -> np.ndarray: return self.xy[self.tracker_id == tracker_id] - - -def pillow_to_cv2(image: Image.Image) -> np.ndarray: - scene = np.array(image) - scene = cv2.cvtColor(scene, cv2.COLOR_RGB2BGR) - return scene - - -def scene_to_annotator_img_type(annotate_func): - """ - Decorates `BaseAnnotator.annotate` implementations, converts scene to - an image type used internally by the annotators, converts back when annotation - is complete. - """ - - @wraps(annotate_func) - def wrapper(self, scene: ImageType, *args, **kwargs): - if isinstance(scene, np.ndarray): - return annotate_func(self, scene, *args, **kwargs) - - if isinstance(scene, Image.Image): - scene = pillow_to_cv2(scene) - annotated = annotate_func(self, scene, *args, **kwargs) - annotated = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB) - annotated = Image.fromarray(annotated) - return annotated - - raise ValueError(f"Unsupported image type: {type(scene)}") - - return wrapper diff --git a/supervision/assets/list.py b/supervision/assets/list.py index 83518cdb..77b1e461 100644 --- a/supervision/assets/list.py +++ b/supervision/assets/list.py @@ -18,6 +18,7 @@ class VideoAssets(Enum): | `SUBWAY` | `subway.mp4` | [Link](https://media.roboflow.com/supervision/video-examples/subway.mp4) | | `MARKET_SQUARE` | `market-square.mp4` | [Link](https://media.roboflow.com/supervision/video-examples/market-square.mp4) | | `PEOPLE_WALKING` | `people-walking.mp4` | [Link](https://media.roboflow.com/supervision/video-examples/people-walking.mp4) | + | `BEACH` | `beach-1.mp4` | [Link](https://media.roboflow.com/supervision/video-examples/beach-1.mp4) | """ # noqa: E501 // docs VEHICLES = "vehicles.mp4" @@ -27,6 +28,7 @@ class VideoAssets(Enum): SUBWAY = "subway.mp4" MARKET_SQUARE = "market-square.mp4" PEOPLE_WALKING = "people-walking.mp4" + BEACH = "beach-1.mp4" @classmethod def list(cls): @@ -62,4 +64,8 @@ VIDEO_ASSETS: Dict[str, Tuple[str, str]] = { f"{BASE_VIDEO_URL}{VideoAssets.PEOPLE_WALKING.value}", "0574c053c8686c3f1dc0aa3743e45cb9", ), + VideoAssets.BEACH.value: ( + f"{BASE_VIDEO_URL}{VideoAssets.BEACH.value}", + "4175d42fec4d450ed081523fd39e0cf8", + ), } diff --git a/supervision/detection/annotate.py b/supervision/detection/annotate.py index a8df0a57..f32c7344 100644 --- a/supervision/detection/annotate.py +++ b/supervision/detection/annotate.py @@ -3,9 +3,9 @@ from typing import List, Optional, Union import cv2 from supervision.annotators.base import ImageType -from supervision.annotators.utils import scene_to_annotator_img_type from supervision.detection.core import Detections from supervision.draw.color import Color, ColorPalette +from supervision.utils.conversion import convert_for_annotation_method from supervision.utils.internal import deprecated @@ -46,7 +46,7 @@ class BoxAnnotator: self.text_thickness: int = text_thickness self.text_padding: int = text_padding - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, diff --git a/supervision/detection/tools/polygon_zone.py b/supervision/detection/tools/polygon_zone.py index 00b746aa..a1997212 100644 --- a/supervision/detection/tools/polygon_zone.py +++ b/supervision/detection/tools/polygon_zone.py @@ -1,3 +1,4 @@ +import warnings from dataclasses import replace from typing import Iterable, Optional, Tuple @@ -11,7 +12,7 @@ from supervision.draw.color import Color from supervision.draw.utils import draw_polygon, draw_text from supervision.geometry.core import Position from supervision.geometry.utils import get_polygon_center -from supervision.utils.internal import deprecated_parameter +from supervision.utils.internal import SupervisionWarnings, deprecated_parameter class PolygonZone: @@ -21,7 +22,6 @@ class PolygonZone: Attributes: polygon (np.ndarray): A polygon represented by a numpy array of shape `(N, 2)`, containing the `x`, `y` coordinates of the points. - frame_resolution_wh (Tuple[int, int]): The frame resolution (width, height) triggering_anchors (Iterable[sv.Position]): A list of positions specifying which anchors of the detections bounding box to consider when deciding on whether the detection fits within the PolygonZone @@ -41,18 +41,26 @@ class PolygonZone: def __init__( self, polygon: npt.NDArray[np.int64], - frame_resolution_wh: Tuple[int, int], + frame_resolution_wh: Optional[Tuple[int, int]] = None, triggering_anchors: Iterable[Position] = (Position.BOTTOM_CENTER,), ): + if frame_resolution_wh is not None: + warnings.warn( + "The `frame_resolution_wh` parameter is no longer required and will be " + "dropped in version supervision-0.24.0. The mask resolution is now " + "calculated automatically based on the polygon coordinates.", + category=SupervisionWarnings, + ) + self.polygon = polygon.astype(int) - self.frame_resolution_wh = frame_resolution_wh self.triggering_anchors = triggering_anchors self.current_count = 0 - width, height = frame_resolution_wh + x_max, y_max = np.max(polygon, axis=0) + self.frame_resolution_wh = (x_max + 1, y_max + 1) self.mask = polygon_to_mask( - polygon=polygon, resolution_wh=(width + 1, height + 1) + polygon=polygon, resolution_wh=(x_max + 2, y_max + 2) ) def trigger(self, detections: Detections) -> npt.NDArray[np.bool_]: diff --git a/supervision/draw/color.py b/supervision/draw/color.py index 149908d8..b195cffe 100644 --- a/supervision/draw/color.py +++ b/supervision/draw/color.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import List, Tuple +from typing import List, Tuple, Union import matplotlib.pyplot as plt @@ -104,10 +104,13 @@ class Color: Create a Color instance from a hex string. Args: - color_hex (str): Hex string of the color. + color_hex (str): The hex string representing the color. This string can + start with '#' followed by either 3 or 6 hexadecimal characters. In + case of 3 characters, each character is repeated to form the full + 6-character hex code. Returns: - Color: Instance representing the color. + Color: An instance representing the color. Example: ```python @@ -115,6 +118,9 @@ class Color: sv.Color.from_hex('#ff00ff') # Color(r=255, g=0, b=255) + + sv.Color.from_hex('#f0f') + # Color(r=255, g=0, b=255) ``` """ _validate_color_hex(color_hex) @@ -124,6 +130,52 @@ class Color: r, g, b = (int(color_hex[i : i + 2], 16) for i in range(0, 6, 2)) return cls(r, g, b) + @classmethod + def from_rgb_tuple(cls, color_tuple: Tuple[int, int, int]) -> Color: + """ + Create a Color instance from an RGB tuple. + + Args: + color_tuple (Tuple[int, int, int]): A tuple representing the color in RGB + format, where each element is an integer in the range 0-255. + + Returns: + Color: An instance representing the color. + + Example: + ```python + import supervision as sv + + sv.Color.from_rgb_tuple((255, 255, 0)) + # Color(r=255, g=255, b=0) + ``` + """ + r, g, b = color_tuple + return cls(r=r, g=g, b=b) + + @classmethod + def from_bgr_tuple(cls, color_tuple: Tuple[int, int, int]) -> Color: + """ + Create a Color instance from a BGR tuple. + + Args: + color_tuple (Tuple[int, int, int]): A tuple representing the color in BGR + format, where each element is an integer in the range 0-255. + + Returns: + Color: An instance representing the color. + + Example: + ```python + import supervision as sv + + sv.Color.from_bgr_tuple((0, 255, 255)) + # Color(r=255, g=255, b=0) + ``` + """ + b, g, r = color_tuple + return cls(r=r, g=g, b=b) + def as_hex(self) -> str: """ Converts the Color instance to a hex string. @@ -396,3 +448,19 @@ class ColorPalette: raise ValueError("idx argument should not be negative") idx = idx % len(self.colors) return self.colors[idx] + + +def unify_to_bgr(color: Union[Tuple[int, int, int], Color]) -> Tuple[int, int, int]: + """ + Converts a color input in multiple formats to a standardized BGR format. + + Args: + color (Union[Tuple[int, int, int], Color]): The color input to be converted, + which can be either a tuple of RGB values or an instance of a Color class. + + Returns: + Tuple[int, int, int]: The color in BGR format as a tuple of three integers. + """ + if issubclass(type(color), Color): + return color.as_bgr() + return color diff --git a/supervision/utils/conversion.py b/supervision/utils/conversion.py new file mode 100644 index 00000000..8ddce969 --- /dev/null +++ b/supervision/utils/conversion.py @@ -0,0 +1,103 @@ +from functools import wraps +from typing import List + +import cv2 +import numpy as np +from PIL import Image + +from supervision.annotators.base import ImageType + + +def convert_for_annotation_method(annotate_func): + """ + Decorates `BaseAnnotator.annotate` implementations, converts scene to + an image type used internally by the annotators, converts back when annotation + is complete. + """ + + @wraps(annotate_func) + def wrapper(self, scene: ImageType, *args, **kwargs): + if isinstance(scene, np.ndarray): + return annotate_func(self, scene, *args, **kwargs) + + if isinstance(scene, Image.Image): + scene = pillow_to_cv2(scene) + annotated = annotate_func(self, scene, *args, **kwargs) + return cv2_to_pillow(image=annotated) + + raise ValueError(f"Unsupported image type: {type(scene)}") + + return wrapper + + +def convert_for_image_processing(image_processing_fun): + """ + Decorates image processing functions that accept np.ndarray, converting `image` to + np.ndarray, converts back when processing is complete. + """ + + @wraps(image_processing_fun) + def wrapper(image: ImageType, *args, **kwargs): + if isinstance(image, np.ndarray): + return image_processing_fun(image, *args, **kwargs) + + if isinstance(image, Image.Image): + scene = pillow_to_cv2(image) + annotated = image_processing_fun(scene, *args, **kwargs) + return cv2_to_pillow(image=annotated) + + raise ValueError(f"Unsupported image type: {type(image)}") + + return wrapper + + +def images_to_cv2(images: List[ImageType]) -> List[np.ndarray]: + """ + Converts images provided either as Pillow images or OpenCV + images into OpenCV format. + + Args: + images (List[ImageType]): Images to be converted + + Returns: + List[np.ndarray]: List of input images in OpenCV format + (with order preserved). + + """ + result = [] + for image in images: + if issubclass(type(image), Image.Image): + image = pillow_to_cv2(image=image) + result.append(image) + return result + + +def pillow_to_cv2(image: Image.Image) -> np.ndarray: + """ + Converts Pillow image into OpenCV image, handling RGB -> BGR + conversion. + + Args: + image (Image.Image): Pillow image (in RGB format). + + Returns: + (np.ndarray): Input image converted to OpenCV format. + """ + scene = np.array(image) + scene = cv2.cvtColor(scene, cv2.COLOR_RGB2BGR) + return scene + + +def cv2_to_pillow(image: np.ndarray) -> Image.Image: + """ + Converts OpenCV image into Pillow image, handling BGR -> RGB + conversion. + + Args: + image (np.ndarray): OpenCV image (in BGR format). + + Returns: + (Image.Image): Input image converted to Pillow format. + """ + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + return Image.fromarray(image) diff --git a/supervision/utils/image.py b/supervision/utils/image.py index 0b005464..9f2e1783 100644 --- a/supervision/utils/image.py +++ b/supervision/utils/image.py @@ -1,107 +1,351 @@ +import itertools +import math import os import shutil -from typing import Optional, Tuple +from functools import partial +from typing import Callable, List, Literal, Optional, Tuple, Union import cv2 import numpy as np +import numpy.typing as npt + +from supervision.annotators.base import ImageType +from supervision.draw.color import Color, unify_to_bgr +from supervision.draw.utils import calculate_optimal_text_scale, draw_text +from supervision.geometry.core import Point +from supervision.utils.conversion import ( + convert_for_image_processing, + cv2_to_pillow, + images_to_cv2, +) +from supervision.utils.iterables import create_batches, fill + +RelativePosition = Literal["top", "bottom"] + +MAX_COLUMNS_FOR_SINGLE_ROW_GRID = 3 -def crop_image(image: np.ndarray, xyxy: np.ndarray) -> np.ndarray: +@convert_for_image_processing +def crop_image( + image: ImageType, + xyxy: Union[npt.NDArray[int], List[int], Tuple[int, int, int, int]], +) -> ImageType: """ Crops the given image based on the given bounding box. Args: - image (np.ndarray): The image to be cropped, represented as a numpy array. - xyxy (np.ndarray): A numpy array containing the bounding box coordinates - in the format (x1, y1, x2, y2). + image (ImageType): The image to be cropped. `ImageType` is a flexible type, + accepting either `numpy.ndarray` or `PIL.Image.Image`. + xyxy (Union[np.ndarray, List[int], Tuple[int, int, int, int]]): A bounding box + coordinates in the format `(x_min, y_min, x_max, y_max)`, accepted as either + a `numpy.ndarray`, a `list`, or a `tuple`. Returns: - (np.ndarray): The cropped image as a numpy array. + (ImageType): The cropped image. The type is determined by the input type and + may be either a `numpy.ndarray` or `PIL.Image.Image`. + + === "OpenCV" - Examples: ```python + import cv2 import supervision as sv - detection = sv.Detections(...) - with sv.ImageSink(target_dir_path='target/directory/path') as sink: - for xyxy in detection.xyxy: - cropped_image = sv.crop_image(image=image, xyxy=xyxy) - sink.save_image(image=cropped_image) + image = cv2.imread() + image.shape + # (1080, 1920, 3) + + xyxy = [200, 400, 600, 800] + cropped_image = sv.crop_image(image=image, xyxy=xyxy) + cropped_image.shape + # (400, 400, 3) ``` - """ + === "Pillow" + + ```python + from PIL import Image + import supervision as sv + + image = Image.open() + image.size + # (1920, 1080) + + xyxy = [200, 400, 600, 800] + cropped_image = sv.crop_image(image=image, xyxy=xyxy) + cropped_image.size + # (400, 400) + ``` + + ![crop_image](https://media.roboflow.com/supervision-docs/crop-image.png){ align=center width="800" } + """ # noqa E501 // docs + + if isinstance(xyxy, (list, tuple)): + xyxy = np.array(xyxy) xyxy = np.round(xyxy).astype(int) - x1, y1, x2, y2 = xyxy - return image[y1:y2, x1:x2] + x_min, y_min, x_max, y_max = xyxy.flatten() + return image[y_min:y_max, x_min:x_max] -def resize_image(image: np.ndarray, scale_factor: float) -> np.ndarray: +@convert_for_image_processing +def scale_image(image: ImageType, scale_factor: float) -> ImageType: """ - Resizes an image by a given scale factor using cv2.INTER_LINEAR interpolation. + Scales the given image based on the given scale factor. Args: - image (np.ndarray): The input image to be resized. - scale_factor (float): The factor by which the image will be scaled. Scale factor - > 1.0 zooms in, < 1.0 zooms out. + image (ImageType): The image to be scaled. `ImageType` is a flexible type, + accepting either `numpy.ndarray` or `PIL.Image.Image`. + scale_factor (float): The factor by which the image will be scaled. Scale + factor > `1.0` zooms in, < `1.0` zooms out. Returns: - np.ndarray: The resized image. + (ImageType): The scaled image. The type is determined by the input type and + may be either a `numpy.ndarray` or `PIL.Image.Image`. Raises: ValueError: If the scale factor is non-positive. + + === "OpenCV" + + ```python + import cv2 + import supervision as sv + + image = cv2.imread() + image.shape + # (1080, 1920, 3) + + scaled_image = sv.scale_image(image=image, scale_factor=0.5) + scaled_image.shape + # (540, 960, 3) + ``` + + === "Pillow" + + ```python + from PIL import Image + import supervision as sv + + image = Image.open() + image.size + # (1920, 1080) + + scaled_image = sv.scale_image(image=image, scale_factor=0.5) + scaled_image.size + # (960, 540) + ``` """ if scale_factor <= 0: raise ValueError("Scale factor must be positive.") - old_width, old_height = image.shape[1], image.shape[0] - nwe_width = int(old_width * scale_factor) - new_height = int(old_height * scale_factor) - - return cv2.resize(image, (nwe_width, new_height), interpolation=cv2.INTER_LINEAR) + width_old, height_old = image.shape[1], image.shape[0] + width_new = int(width_old * scale_factor) + height_new = int(height_old * scale_factor) + return cv2.resize(image, (width_new, height_new), interpolation=cv2.INTER_LINEAR) -def place_image( - scene: np.ndarray, image: np.ndarray, anchor: Tuple[int, int] -) -> np.ndarray: +@convert_for_image_processing +def resize_image( + image: ImageType, + resolution_wh: Tuple[int, int], + keep_aspect_ratio: bool = False, +) -> ImageType: + """ + Resizes the given image to a specified resolution. Can maintain the original aspect + ratio or resize directly to the desired dimensions. + + Args: + image (ImageType): The image to be resized. `ImageType` is a flexible type, + accepting either `numpy.ndarray` or `PIL.Image.Image`. + resolution_wh (Tuple[int, int]): The target resolution as + `(width, height)`. + keep_aspect_ratio (bool, optional): Flag to maintain the image's original + aspect ratio. Defaults to `False`. + + Returns: + (ImageType): The resized image. The type is determined by the input type and + may be either a `numpy.ndarray` or `PIL.Image.Image`. + + === "OpenCV" + + ```python + import cv2 + import supervision as sv + + image = cv2.imread() + image.shape + # (1080, 1920, 3) + + resized_image = sv.resize_image( + image=image, resolution_wh=(1000, 1000), keep_aspect_ratio=True + ) + resized_image.shape + # (562, 1000, 3) + ``` + + === "Pillow" + + ```python + from PIL import Image + import supervision as sv + + image = Image.open() + image.size + # (1920, 1080) + + resized_image = sv.resize_image( + image=image, resolution_wh=(1000, 1000), keep_aspect_ratio=True + ) + resized_image.size + # (1000, 562) + ``` + + ![resize_image](https://media.roboflow.com/supervision-docs/resize-image.png){ align=center width="800" } + """ # noqa E501 // docs + if keep_aspect_ratio: + image_ratio = image.shape[1] / image.shape[0] + target_ratio = resolution_wh[0] / resolution_wh[1] + if image_ratio >= target_ratio: + width_new = resolution_wh[0] + height_new = int(resolution_wh[0] / image_ratio) + else: + height_new = resolution_wh[1] + width_new = int(resolution_wh[1] * image_ratio) + else: + width_new, height_new = resolution_wh + + return cv2.resize(image, (width_new, height_new), interpolation=cv2.INTER_LINEAR) + + +@convert_for_image_processing +def letterbox_image( + image: ImageType, + resolution_wh: Tuple[int, int], + color: Union[Tuple[int, int, int], Color] = Color.BLACK, +) -> ImageType: + """ + Resizes and pads an image to a specified resolution with a given color, maintaining + the original aspect ratio. + + Args: + image (ImageType): The image to be resized. `ImageType` is a flexible type, + accepting either `numpy.ndarray` or `PIL.Image.Image`. + resolution_wh (Tuple[int, int]): The target resolution as + `(width, height)`. + color (Union[Tuple[int, int, int], Color]): The color to pad with. If tuple + provided it should be in BGR format. + + Returns: + (ImageType): The resized image. The type is determined by the input type and + may be either a `numpy.ndarray` or `PIL.Image.Image`. + + === "OpenCV" + + ```python + import cv2 + import supervision as sv + + image = cv2.imread() + image.shape + # (1080, 1920, 3) + + letterboxed_image = sv.letterbox_image(image=image, resolution_wh=(1000, 1000)) + letterboxed_image.shape + # (1000, 1000, 3) + ``` + + === "Pillow" + + ```python + from PIL import Image + import supervision as sv + + image = Image.open() + image.size + # (1920, 1080) + + letterboxed_image = sv.letterbox_image(image=image, resolution_wh=(1000, 1000)) + letterboxed_image.size + # (1000, 1000) + ``` + + ![letterbox_image](https://media.roboflow.com/supervision-docs/letterbox-image.png){ align=center width="800" } + """ # noqa E501 // docs + color = unify_to_bgr(color=color) + resized_image = resize_image( + image=image, resolution_wh=resolution_wh, keep_aspect_ratio=True + ) + height_new, width_new = resized_image.shape[:2] + padding_top = (resolution_wh[1] - height_new) // 2 + padding_bottom = resolution_wh[1] - height_new - padding_top + padding_left = (resolution_wh[0] - width_new) // 2 + padding_right = resolution_wh[0] - width_new - padding_left + return cv2.copyMakeBorder( + resized_image, + padding_top, + padding_bottom, + padding_left, + padding_right, + cv2.BORDER_CONSTANT, + value=color, + ) + + +def overlay_image( + image: npt.NDArray[np.uint8], + overlay: npt.NDArray[np.uint8], + anchor: Tuple[int, int], +) -> npt.NDArray[np.uint8]: """ Places an image onto a scene at a given anchor point, handling cases where the image's position is partially or completely outside the scene's bounds. Args: - scene (np.ndarray): The background scene onto which the image is placed. - image (np.ndarray): The image to be placed onto the scene. - anchor (Tuple[int, int]): The (x, y) coordinates in the scene where the + image (np.ndarray): The background scene onto which the image is placed. + overlay (np.ndarray): The image to be placed onto the scene. + anchor (Tuple[int, int]): The `(x, y)` coordinates in the scene where the top-left corner of the image will be placed. Returns: - np.ndarray: The modified scene with the image placed at the anchor point, - or unchanged if the image placement is completely outside the scene. - """ - scene_height, scene_width = scene.shape[:2] - image_height, image_width = image.shape[:2] + (np.ndarray): The result image with overlay. + + Examples: + ```python + import cv2 + import numpy as np + import supervision as sv + + image = cv2.imread() + overlay = np.zeros((400, 400, 3), dtype=np.uint8) + result_image = sv.overlay_image(image=image, overlay=overlay, anchor=(200, 400)) + ``` + + ![overlay_image](https://media.roboflow.com/supervision-docs/overlay-image.png){ align=center width="800" } + """ # noqa E501 // docs + scene_height, scene_width = image.shape[:2] + image_height, image_width = overlay.shape[:2] anchor_x, anchor_y = anchor is_out_horizontally = anchor_x + image_width <= 0 or anchor_x >= scene_width is_out_vertically = anchor_y + image_height <= 0 or anchor_y >= scene_height if is_out_horizontally or is_out_vertically: - return scene + return image - start_y = max(anchor_y, 0) - start_x = max(anchor_x, 0) - end_y = min(scene_height, anchor_y + image_height) - end_x = min(scene_width, anchor_x + image_width) + x_min = max(anchor_x, 0) + y_min = max(anchor_y, 0) + x_max = min(scene_width, anchor_x + image_width) + y_max = min(scene_height, anchor_y + image_height) - crop_start_y = max(-anchor_y, 0) - crop_start_x = max(-anchor_x, 0) - crop_end_y = image_height - max((anchor_y + image_height) - scene_height, 0) - crop_end_x = image_width - max((anchor_x + image_width) - scene_width, 0) + crop_x_min = max(-anchor_x, 0) + crop_y_min = max(-anchor_y, 0) + crop_x_max = image_width - max((anchor_x + image_width) - scene_width, 0) + crop_y_max = image_height - max((anchor_y + image_height) - scene_height, 0) - scene[start_y:end_y, start_x:end_x] = image[ - crop_start_y:crop_end_y, crop_start_x:crop_end_x + image[y_min:y_max, x_min:x_max] = overlay[ + crop_y_min:crop_y_max, crop_x_min:crop_x_max ] - return scene + return image class ImageSink: @@ -125,13 +369,13 @@ class ImageSink: ```python import supervision as sv - with sv.ImageSink(target_dir_path='target/directory/path', - overwrite=True) as sink: - for image in sv.get_video_frames_generator( - source_path='source_video.mp4', stride=2): + frames_generator = sv.get_video_frames_generator(, stride=2) + + with sv.ImageSink(target_dir_path=) as sink: + for image in frames_generator: sink.save_image(image=image) ``` - """ + """ # noqa E501 // docs self.target_dir_path = target_dir_path self.overwrite = overwrite @@ -168,3 +412,357 @@ class ImageSink: def __exit__(self, exc_type, exc_value, exc_traceback): pass + + +def create_tiles( + images: List[ImageType], + grid_size: Optional[Tuple[Optional[int], Optional[int]]] = None, + single_tile_size: Optional[Tuple[int, int]] = None, + tile_scaling: Literal["min", "max", "avg"] = "avg", + tile_padding_color: Union[Tuple[int, int, int], Color] = Color.from_hex("#D9D9D9"), + tile_margin: int = 10, + tile_margin_color: Union[Tuple[int, int, int], Color] = Color.from_hex("#BFBEBD"), + return_type: Literal["auto", "cv2", "pillow"] = "auto", + titles: Optional[List[Optional[str]]] = None, + titles_anchors: Optional[Union[Point, List[Optional[Point]]]] = None, + titles_color: Union[Tuple[int, int, int], Color] = Color.from_hex("#262523"), + titles_scale: Optional[float] = None, + titles_thickness: int = 1, + titles_padding: int = 10, + titles_text_font: int = cv2.FONT_HERSHEY_SIMPLEX, + titles_background_color: Union[Tuple[int, int, int], Color] = Color.from_hex( + "#D9D9D9" + ), + default_title_placement: RelativePosition = "top", +) -> ImageType: + """ + Creates tiles mosaic from input images, automating grid placement and + converting images to common resolution maintaining aspect ratio. It is + also possible to render text titles on tiles, using optional set of + parameters specifying text drawing (see parameters description). + + Automated grid placement will try to maintain square shape of grid + (with size being the nearest integer square root of #images), up to two exceptions: + * if there are up to 3 images - images will be displayed in single row + * if square-grid placement causes last row to be empty - number of rows is trimmed + until last row has at least one image + + Args: + images (List[ImageType]): Images to create tiles. Elements can be either + np.ndarray or PIL.Image, common representation will be agreed by the + function. + grid_size (Optional[Tuple[Optional[int], Optional[int]]]): Expected grid + size in format (n_rows, n_cols). If not given - automated grid placement + will be applied. One may also provide only one out of two elements of the + tuple - then grid will be created with either n_rows or n_cols fixed, + leaving the other dimension to be adjusted by the number of images + single_tile_size (Optional[Tuple[int, int]]): sizeof a single tile element + provided in (width, height) format. If not given - size of tile will be + automatically calculated based on `tile_scaling` parameter. + tile_scaling (Literal["min", "max", "avg"]): If `single_tile_size` is not + given - parameter will be used to calculate tile size - using + min / max / avg size of image provided in `images` list. + tile_padding_color (Union[Tuple[int, int, int], sv.Color]): Color to be used in + images letterbox procedure (while standardising tiles sizes) as a padding. + If tuple provided - should be BGR. + tile_margin (int): size of margin between tiles (in pixels) + tile_margin_color (Union[Tuple[int, int, int], sv.Color]): Color of tile margin. + If tuple provided - should be BGR. + return_type (Literal["auto", "cv2", "pillow"]): Parameter dictates the format of + return image. One may choose specific type ("cv2" or "pillow") to enforce + conversion. "auto" mode takes a majority vote between types of elements in + `images` list - resolving draws in favour of OpenCV format. "auto" can be + safely used when all input images are of the same type. + titles (Optional[List[Optional[str]]]): Optional titles to be added to tiles. + Elements of that list may be empty - then specific tile (in order presented + in `images` parameter) will not be filled with title. It is possible to + provide list of titles shorter than `images` - then remaining titles will + be assumed empty. + titles_anchors (Optional[Union[Point, List[Optional[Point]]]]): Parameter to + specify anchor points for titles. It is possible to specify anchor either + globally or for specific tiles (following order of `images`). + If not given (either globally, or for specific element of the list), + it will be calculated automatically based on `default_title_placement`. + titles_color (Union[Tuple[int, int, int], Color]): Color of titles text. + If tuple provided - should be BGR. + titles_scale (Optional[float]): Scale of titles. If not provided - value will + be calculated using `calculate_optimal_text_scale(...)`. + titles_thickness (int): Thickness of titles text. + titles_padding (int): Size of titles padding. + titles_text_font (int): Font to be used to render titles. Must be integer + constant representing OpenCV font. + (See docs: https://docs.opencv.org/4.x/d6/d6e/group__imgproc__draw.html) + titles_background_color (Union[Tuple[int, int, int], Color]): Color of title + text padding. + default_title_placement (Literal["top", "bottom"]): Parameter specifies title + anchor placement in case if explicit anchor is not provided. + + Returns: + ImageType: Image with all input images located in tails grid. The output type is + determined by `return_type` parameter. + + Raises: + ValueError: In case when input images list is empty, provided `grid_size` is too + small to fit all images, `tile_scaling` mode is invalid. + """ + if len(images) == 0: + raise ValueError("Could not create image tiles from empty list of images.") + if return_type == "auto": + return_type = _negotiate_tiles_format(images=images) + tile_padding_color = unify_to_bgr(color=tile_padding_color) + tile_margin_color = unify_to_bgr(color=tile_margin_color) + images = images_to_cv2(images=images) + if single_tile_size is None: + single_tile_size = _aggregate_images_shape(images=images, mode=tile_scaling) + resized_images = [ + letterbox_image( + image=i, resolution_wh=single_tile_size, color=tile_padding_color + ) + for i in images + ] + grid_size = _establish_grid_size(images=images, grid_size=grid_size) + if len(images) > grid_size[0] * grid_size[1]: + raise ValueError( + f"Could not place {len(images)} in grid with size: {grid_size}." + ) + if titles is not None: + titles = fill(sequence=titles, desired_size=len(images), content=None) + titles_anchors = ( + [titles_anchors] + if not issubclass(type(titles_anchors), list) + else titles_anchors + ) + titles_anchors = fill( + sequence=titles_anchors, desired_size=len(images), content=None + ) + titles_color = unify_to_bgr(color=titles_color) + titles_background_color = unify_to_bgr(color=titles_background_color) + tiles = _generate_tiles( + images=resized_images, + grid_size=grid_size, + single_tile_size=single_tile_size, + tile_padding_color=tile_padding_color, + tile_margin=tile_margin, + tile_margin_color=tile_margin_color, + titles=titles, + titles_anchors=titles_anchors, + titles_color=titles_color, + titles_scale=titles_scale, + titles_thickness=titles_thickness, + titles_padding=titles_padding, + titles_text_font=titles_text_font, + titles_background_color=titles_background_color, + default_title_placement=default_title_placement, + ) + if return_type == "pillow": + tiles = cv2_to_pillow(image=tiles) + return tiles + + +def _negotiate_tiles_format(images: List[ImageType]) -> Literal["cv2", "pillow"]: + number_of_np_arrays = sum(issubclass(type(i), np.ndarray) for i in images) + if number_of_np_arrays >= (len(images) // 2): + return "cv2" + return "pillow" + + +def _calculate_aggregated_images_shape( + images: List[np.ndarray], aggregator: Callable[[List[int]], float] +) -> Tuple[int, int]: + height = round(aggregator([i.shape[0] for i in images])) + width = round(aggregator([i.shape[1] for i in images])) + return width, height + + +SHAPE_AGGREGATION_FUN = { + "min": partial(_calculate_aggregated_images_shape, aggregator=np.min), + "max": partial(_calculate_aggregated_images_shape, aggregator=np.max), + "avg": partial(_calculate_aggregated_images_shape, aggregator=np.average), +} + + +def _aggregate_images_shape( + images: List[np.ndarray], mode: Literal["min", "max", "avg"] +) -> Tuple[int, int]: + if mode not in SHAPE_AGGREGATION_FUN: + raise ValueError( + f"Could not aggregate images shape - provided unknown mode: {mode}. " + f"Supported modes: {list(SHAPE_AGGREGATION_FUN.keys())}." + ) + return SHAPE_AGGREGATION_FUN[mode](images) + + +def _establish_grid_size( + images: List[np.ndarray], grid_size: Optional[Tuple[Optional[int], Optional[int]]] +) -> Tuple[int, int]: + if grid_size is None or all(e is None for e in grid_size): + return _negotiate_grid_size(images=images) + if grid_size[0] is None: + return math.ceil(len(images) / grid_size[1]), grid_size[1] + if grid_size[1] is None: + return grid_size[0], math.ceil(len(images) / grid_size[0]) + return grid_size + + +def _negotiate_grid_size(images: List[np.ndarray]) -> Tuple[int, int]: + if len(images) <= MAX_COLUMNS_FOR_SINGLE_ROW_GRID: + return 1, len(images) + nearest_sqrt = math.ceil(np.sqrt(len(images))) + proposed_columns = nearest_sqrt + proposed_rows = nearest_sqrt + while proposed_columns * (proposed_rows - 1) >= len(images): + proposed_rows -= 1 + return proposed_rows, proposed_columns + + +def _generate_tiles( + images: List[np.ndarray], + grid_size: Tuple[int, int], + single_tile_size: Tuple[int, int], + tile_padding_color: Tuple[int, int, int], + tile_margin: int, + tile_margin_color: Tuple[int, int, int], + titles: Optional[List[Optional[str]]], + titles_anchors: List[Optional[Point]], + titles_color: Tuple[int, int, int], + titles_scale: Optional[float], + titles_thickness: int, + titles_padding: int, + titles_text_font: int, + titles_background_color: Tuple[int, int, int], + default_title_placement: RelativePosition, +) -> np.ndarray: + images = _draw_texts( + images=images, + titles=titles, + titles_anchors=titles_anchors, + titles_color=titles_color, + titles_scale=titles_scale, + titles_thickness=titles_thickness, + titles_padding=titles_padding, + titles_text_font=titles_text_font, + titles_background_color=titles_background_color, + default_title_placement=default_title_placement, + ) + rows, columns = grid_size + tiles_elements = list(create_batches(sequence=images, batch_size=columns)) + while len(tiles_elements[-1]) < columns: + tiles_elements[-1].append( + _generate_color_image(shape=single_tile_size, color=tile_padding_color) + ) + while len(tiles_elements) < rows: + tiles_elements.append( + [_generate_color_image(shape=single_tile_size, color=tile_padding_color)] + * columns + ) + return _merge_tiles_elements( + tiles_elements=tiles_elements, + grid_size=grid_size, + single_tile_size=single_tile_size, + tile_margin=tile_margin, + tile_margin_color=tile_margin_color, + ) + + +def _draw_texts( + images: List[np.ndarray], + titles: Optional[List[Optional[str]]], + titles_anchors: List[Optional[Point]], + titles_color: Tuple[int, int, int], + titles_scale: Optional[float], + titles_thickness: int, + titles_padding: int, + titles_text_font: int, + titles_background_color: Tuple[int, int, int], + default_title_placement: RelativePosition, +) -> List[np.ndarray]: + if titles is None: + return images + titles_anchors = _prepare_default_titles_anchors( + images=images, + titles_anchors=titles_anchors, + default_title_placement=default_title_placement, + ) + if titles_scale is None: + image_height, image_width = images[0].shape[:2] + titles_scale = calculate_optimal_text_scale( + resolution_wh=(image_width, image_height) + ) + result = [] + for image, text, anchor in zip(images, titles, titles_anchors): + if text is None: + result.append(image) + continue + processed_image = draw_text( + scene=image, + text=text, + text_anchor=anchor, + text_color=Color.from_bgr_tuple(titles_color), + text_scale=titles_scale, + text_thickness=titles_thickness, + text_padding=titles_padding, + text_font=titles_text_font, + background_color=Color.from_bgr_tuple(titles_background_color), + ) + result.append(processed_image) + return result + + +def _prepare_default_titles_anchors( + images: List[np.ndarray], + titles_anchors: List[Optional[Point]], + default_title_placement: RelativePosition, +) -> List[Point]: + result = [] + for image, anchor in zip(images, titles_anchors): + if anchor is not None: + result.append(anchor) + continue + image_height, image_width = image.shape[:2] + if default_title_placement == "top": + default_anchor = Point(x=image_width / 2, y=image_height * 0.1) + else: + default_anchor = Point(x=image_width / 2, y=image_height * 0.9) + result.append(default_anchor) + return result + + +def _merge_tiles_elements( + tiles_elements: List[List[np.ndarray]], + grid_size: Tuple[int, int], + single_tile_size: Tuple[int, int], + tile_margin: int, + tile_margin_color: Tuple[int, int, int], +) -> np.ndarray: + vertical_padding = ( + np.ones((single_tile_size[1], tile_margin, 3)) * tile_margin_color + ) + merged_rows = [ + np.concatenate( + list( + itertools.chain.from_iterable( + zip(row, [vertical_padding] * grid_size[1]) + ) + )[:-1], + axis=1, + ) + for row in tiles_elements + ] + row_width = merged_rows[0].shape[1] + horizontal_padding = ( + np.ones((tile_margin, row_width, 3), dtype=np.uint8) * tile_margin_color + ) + rows_with_paddings = [] + for row in merged_rows: + rows_with_paddings.append(row) + rows_with_paddings.append(horizontal_padding) + return np.concatenate( + rows_with_paddings[:-1], + axis=0, + ).astype(np.uint8) + + +def _generate_color_image( + shape: Tuple[int, int], color: Tuple[int, int, int] +) -> np.ndarray: + return np.ones(shape[::-1] + (3,), dtype=np.uint8) * color diff --git a/supervision/utils/iterables.py b/supervision/utils/iterables.py new file mode 100644 index 00000000..52bfbeb6 --- /dev/null +++ b/supervision/utils/iterables.py @@ -0,0 +1,70 @@ +from typing import Generator, Iterable, List, TypeVar + +V = TypeVar("V") + + +def create_batches( + sequence: Iterable[V], batch_size: int +) -> Generator[List[V], None, None]: + """ + Provides a generator that yields chunks of the input sequence + of the size specified by the `batch_size` parameter. The last + chunk may be a smaller batch. + + Args: + sequence (Iterable[V]): The sequence to be split into batches. + batch_size (int): The expected size of a batch. + + Returns: + (Generator[List[V], None, None]): A generator that yields chunks + of `sequence` of size `batch_size`, up to the length of + the input `sequence`. + + Examples: + ```python + list(create_batches([1, 2, 3, 4, 5], 2)) + # [[1, 2], [3, 4], [5]] + + list(create_batches("abcde", 3)) + # [['a', 'b', 'c'], ['d', 'e']] + ``` + """ + batch_size = max(batch_size, 1) + current_batch = [] + for element in sequence: + if len(current_batch) == batch_size: + yield current_batch + current_batch = [] + current_batch.append(element) + if current_batch: + yield current_batch + + +def fill(sequence: List[V], desired_size: int, content: V) -> List[V]: + """ + Fill the sequence with padding elements until the sequence reaches + the desired size. + + Args: + sequence (List[V]): The input sequence. + desired_size (int): The expected size of the output list. The + difference between this value and the actual length of `sequence` + (if positive) dictates how many elements will be added as padding. + content (V): The element to be placed at the end of the input + `sequence` as padding. + + Returns: + (List[V]): A padded version of the input `sequence` (if needed). + + Examples: + ```python + fill([1, 2], 4, 0) + # [1, 2, 0, 0] + + fill(['a', 'b'], 3, 'c') + # ['a', 'b', 'c'] + ``` + """ + missing_size = max(0, desired_size - len(sequence)) + sequence.extend([content] * missing_size) + return sequence diff --git a/supervision/utils/notebook.py b/supervision/utils/notebook.py index 159f3c17..19f5eaed 100644 --- a/supervision/utils/notebook.py +++ b/supervision/utils/notebook.py @@ -5,7 +5,7 @@ import matplotlib.pyplot as plt from PIL import Image from supervision.annotators.base import ImageType -from supervision.annotators.utils import pillow_to_cv2 +from supervision.utils.conversion import pillow_to_cv2 def plot_image( diff --git a/test/utils/assets/1.jpg b/test/utils/assets/1.jpg new file mode 100644 index 00000000..ed88f941 Binary files /dev/null and b/test/utils/assets/1.jpg differ diff --git a/test/utils/assets/2.jpg b/test/utils/assets/2.jpg new file mode 100644 index 00000000..cbe01e98 Binary files /dev/null and b/test/utils/assets/2.jpg differ diff --git a/test/utils/assets/3.jpg b/test/utils/assets/3.jpg new file mode 100644 index 00000000..8fc100b9 Binary files /dev/null and b/test/utils/assets/3.jpg differ diff --git a/test/utils/assets/4.jpg b/test/utils/assets/4.jpg new file mode 100644 index 00000000..16467d97 Binary files /dev/null and b/test/utils/assets/4.jpg differ diff --git a/test/utils/assets/5.jpg b/test/utils/assets/5.jpg new file mode 100644 index 00000000..e58fd9e1 Binary files /dev/null and b/test/utils/assets/5.jpg differ diff --git a/test/utils/assets/all_images_tile.png b/test/utils/assets/all_images_tile.png new file mode 100644 index 00000000..ef5066a7 Binary files /dev/null and b/test/utils/assets/all_images_tile.png differ diff --git a/test/utils/assets/all_images_tile_and_custom_colors.png b/test/utils/assets/all_images_tile_and_custom_colors.png new file mode 100644 index 00000000..db8ed13b Binary files /dev/null and b/test/utils/assets/all_images_tile_and_custom_colors.png differ diff --git a/test/utils/assets/all_images_tile_and_custom_colors_and_titles.png b/test/utils/assets/all_images_tile_and_custom_colors_and_titles.png new file mode 100644 index 00000000..8a058f97 Binary files /dev/null and b/test/utils/assets/all_images_tile_and_custom_colors_and_titles.png differ diff --git a/test/utils/assets/all_images_tile_and_custom_grid.png b/test/utils/assets/all_images_tile_and_custom_grid.png new file mode 100644 index 00000000..1407a97a Binary files /dev/null and b/test/utils/assets/all_images_tile_and_custom_grid.png differ diff --git a/test/utils/assets/all_images_tile_and_titles_with_custom_configs.png b/test/utils/assets/all_images_tile_and_titles_with_custom_configs.png new file mode 100644 index 00000000..0da1d2b8 Binary files /dev/null and b/test/utils/assets/all_images_tile_and_titles_with_custom_configs.png differ diff --git a/test/utils/assets/four_images_tile.png b/test/utils/assets/four_images_tile.png new file mode 100644 index 00000000..220df446 Binary files /dev/null and b/test/utils/assets/four_images_tile.png differ diff --git a/test/utils/assets/single_image_tile.png b/test/utils/assets/single_image_tile.png new file mode 100644 index 00000000..434deb00 Binary files /dev/null and b/test/utils/assets/single_image_tile.png differ diff --git a/test/utils/assets/single_image_tile_enforced_grid.png b/test/utils/assets/single_image_tile_enforced_grid.png new file mode 100644 index 00000000..0f8b5ce4 Binary files /dev/null and b/test/utils/assets/single_image_tile_enforced_grid.png differ diff --git a/test/utils/assets/three_images_tile.png b/test/utils/assets/three_images_tile.png new file mode 100644 index 00000000..104bba8f Binary files /dev/null and b/test/utils/assets/three_images_tile.png differ diff --git a/test/utils/assets/two_images_tile.png b/test/utils/assets/two_images_tile.png new file mode 100644 index 00000000..9922c68e Binary files /dev/null and b/test/utils/assets/two_images_tile.png differ diff --git a/test/utils/conftest.py b/test/utils/conftest.py new file mode 100644 index 00000000..cc134f8c --- /dev/null +++ b/test/utils/conftest.py @@ -0,0 +1,99 @@ +import os +from typing import List + +import cv2 +import numpy as np +from _pytest.fixtures import fixture +from PIL import Image + +ASSETS_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "assets")) +ALL_IMAGES_LIST = [os.path.join(ASSETS_DIR, f"{i}.jpg") for i in range(1, 6)] + + +@fixture(scope="function") +def empty_opencv_image() -> np.ndarray: + return np.zeros((128, 128, 3), dtype=np.uint8) + + +@fixture(scope="function") +def empty_pillow_image() -> Image.Image: + return Image.new(mode="RGB", size=(128, 128), color=(0, 0, 0)) + + +@fixture(scope="function") +def all_images() -> List[np.ndarray]: + return [cv2.imread(path) for path in ALL_IMAGES_LIST] + + +@fixture(scope="function") +def one_image() -> np.ndarray: + return cv2.imread(ALL_IMAGES_LIST[0]) + + +@fixture(scope="function") +def two_images() -> List[np.ndarray]: + return [cv2.imread(path) for path in ALL_IMAGES_LIST[:2]] + + +@fixture(scope="function") +def three_images() -> List[np.ndarray]: + return [cv2.imread(path) for path in ALL_IMAGES_LIST[:3]] + + +@fixture(scope="function") +def four_images() -> List[np.ndarray]: + return [cv2.imread(path) for path in ALL_IMAGES_LIST[:4]] + + +@fixture(scope="function") +def all_images_tile() -> np.ndarray: + return cv2.imread(os.path.join(ASSETS_DIR, "all_images_tile.png")) + + +@fixture(scope="function") +def all_images_tile_and_custom_colors() -> np.ndarray: + return cv2.imread(os.path.join(ASSETS_DIR, "all_images_tile_and_custom_colors.png")) + + +@fixture(scope="function") +def all_images_tile_and_custom_grid() -> np.ndarray: + return cv2.imread(os.path.join(ASSETS_DIR, "all_images_tile_and_custom_grid.png")) + + +@fixture(scope="function") +def four_images_tile() -> np.ndarray: + return cv2.imread(os.path.join(ASSETS_DIR, "four_images_tile.png")) + + +@fixture(scope="function") +def single_image_tile() -> np.ndarray: + return cv2.imread(os.path.join(ASSETS_DIR, "single_image_tile.png")) + + +@fixture(scope="function") +def single_image_tile_enforced_grid() -> np.ndarray: + return cv2.imread(os.path.join(ASSETS_DIR, "single_image_tile_enforced_grid.png")) + + +@fixture(scope="function") +def three_images_tile() -> np.ndarray: + return cv2.imread(os.path.join(ASSETS_DIR, "three_images_tile.png")) + + +@fixture(scope="function") +def two_images_tile() -> np.ndarray: + return cv2.imread(os.path.join(ASSETS_DIR, "two_images_tile.png")) + + +@fixture(scope="function") +def all_images_tile_and_custom_colors_and_titles() -> np.ndarray: + return cv2.imread( + os.path.join(ASSETS_DIR, "all_images_tile_and_custom_colors_and_titles.png") + ) + + +@fixture(scope="function") +def all_images_tile_and_titles_with_custom_configs() -> np.ndarray: + return cv2.imread( + os.path.join(ASSETS_DIR, "all_images_tile_and_titles_with_custom_configs.png") + ) diff --git a/test/utils/test_conversion.py b/test/utils/test_conversion.py new file mode 100644 index 00000000..fb3d8faf --- /dev/null +++ b/test/utils/test_conversion.py @@ -0,0 +1,173 @@ +import numpy as np +from PIL import Image, ImageChops + +from supervision.utils.conversion import ( + convert_for_image_processing, + cv2_to_pillow, + images_to_cv2, + pillow_to_cv2, +) + + +def test_convert_for_image_processing_when_pillow_image_submitted( + empty_opencv_image: np.ndarray, empty_pillow_image: Image.Image +) -> None: + # given + param_a_value = 3 + param_b_value = "some" + + @convert_for_image_processing + def my_custom_processing_function( + image: np.ndarray, + param_a: int, + param_b: str, + ) -> np.ndarray: + assert np.allclose( + image, empty_opencv_image + ), "Expected conversion to OpenCV image to happen" + assert ( + param_a == param_a_value + ), f"Parameter a expected to be {param_a_value} in target function" + assert ( + param_b == param_b_value + ), f"Parameter b expected to be {param_b_value} in target function" + return image + + # when + result = my_custom_processing_function( + empty_pillow_image, + param_a_value, + param_b=param_b_value, + ) + + # then + difference = ImageChops.difference(result, empty_pillow_image) + assert difference.getbbox() is None, ( + "Wrapper is expected to convert-back the OpenCV image " + "into Pillow format without changes to content" + ) + + +def test_convert_for_image_processing_when_opencv_image_submitted( + empty_opencv_image: np.ndarray, +) -> None: + # given + param_a_value = 3 + param_b_value = "some" + + @convert_for_image_processing + def my_custom_processing_function( + image: np.ndarray, + param_a: int, + param_b: str, + ) -> np.ndarray: + assert np.allclose( + image, empty_opencv_image + ), "Expected conversion to OpenCV image to happen" + assert ( + param_a == param_a_value + ), f"Parameter a expected to be {param_a_value} in target function" + assert ( + param_b == param_b_value + ), f"Parameter b expected to be {param_b_value} in target function" + return image + + # when + result = my_custom_processing_function( + empty_opencv_image, + param_a_value, + param_b=param_b_value, + ) + + # then + assert ( + result is empty_opencv_image + ), "Expected to return OpenCV image without changes" + + +def test_cv2_to_pillow( + empty_opencv_image: np.ndarray, empty_pillow_image: Image.Image +) -> None: + # when + result = cv2_to_pillow(image=empty_opencv_image) + + # then + difference = ImageChops.difference(result, empty_pillow_image) + assert ( + difference.getbbox() is None + ), "Conversion to PIL.Image expected not to change the content of image" + + +def test_pillow_to_cv2( + empty_opencv_image: np.ndarray, empty_pillow_image: Image.Image +) -> None: + # when + result = pillow_to_cv2(image=empty_pillow_image) + + # then + assert np.allclose( + result, empty_opencv_image + ), "Conversion to OpenCV image expected not to change the content of image" + + +def test_images_to_cv2_when_empty_input_provided() -> None: + # when + result = images_to_cv2(images=[]) + + # then + assert result == [], "Expected empty output when empty input provided" + + +def test_images_to_cv2_when_only_cv2_images_provided( + empty_opencv_image: np.ndarray, +) -> None: + # given + images = [empty_opencv_image] * 5 + + # when + result = images_to_cv2(images=images) + + # then + assert len(result) == 5, "Expected the same number of output element as input ones" + for result_element in result: + assert ( + result_element is empty_opencv_image + ), "Expected CV images not to be touched by conversion" + + +def test_images_to_cv2_when_only_pillow_images_provided( + empty_pillow_image: Image.Image, + empty_opencv_image: np.ndarray, +) -> None: + # given + images = [empty_pillow_image] * 5 + + # when + result = images_to_cv2(images=images) + + # then + assert len(result) == 5, "Expected the same number of output element as input ones" + for result_element in result: + assert np.allclose( + result_element, empty_opencv_image + ), "Output images expected to be equal to empty OpenCV image" + + +def test_images_to_cv2_when_mixed_input_provided( + empty_pillow_image: Image.Image, + empty_opencv_image: np.ndarray, +) -> None: + # given + images = [empty_pillow_image, empty_opencv_image] + + # when + result = images_to_cv2(images=images) + + # then + assert len(result) == 2, "Expected the same number of output element as input ones" + assert np.allclose( + result[0], empty_opencv_image + ), "PIL image should be converted to OpenCV one, equal to example empty image" + assert ( + result[1] is empty_opencv_image + ), "Expected CV images not to be touched by conversion" diff --git a/test/utils/test_image.py b/test/utils/test_image.py new file mode 100644 index 00000000..487434ae --- /dev/null +++ b/test/utils/test_image.py @@ -0,0 +1,244 @@ +from typing import List + +import numpy as np +import pytest +from PIL import Image, ImageChops + +from supervision import Color, Point +from supervision.utils.image import create_tiles, letterbox_image, resize_image + + +def test_resize_image_for_opencv_image() -> None: + # given + image = np.zeros((480, 640, 3), dtype=np.uint8) + expected_result = np.zeros((768, 1024, 3), dtype=np.uint8) + + # when + result = resize_image( + image=image, + resolution_wh=(1024, 1024), + keep_aspect_ratio=True, + ) + + # then + assert np.allclose( + result, expected_result + ), "Expected output shape to be (w, h): (1024, 768)" + + +def test_resize_image_for_pillow_image() -> None: + # given + image = Image.new(mode="RGB", size=(640, 480), color=(0, 0, 0)) + expected_result = Image.new(mode="RGB", size=(1024, 768), color=(0, 0, 0)) + + # when + result = resize_image( + image=image, + resolution_wh=(1024, 1024), + keep_aspect_ratio=True, + ) + + # then + assert result.size == (1024, 768), "Expected output shape to be (w, h): (1024, 768)" + difference = ImageChops.difference(result, expected_result) + assert ( + difference.getbbox() is None + ), "Expected no difference in resized image content as the image is all zeros" + + +def test_letterbox_image_for_opencv_image() -> None: + # given + image = np.zeros((480, 640, 3), dtype=np.uint8) + expected_result = np.concatenate( + [ + np.ones((128, 1024, 3), dtype=np.uint8) * 255, + np.zeros((768, 1024, 3), dtype=np.uint8), + np.ones((128, 1024, 3), dtype=np.uint8) * 255, + ], + axis=0, + ) + + # when + result = letterbox_image( + image=image, resolution_wh=(1024, 1024), color=(255, 255, 255) + ) + + # then + assert np.allclose(result, expected_result), ( + "Expected output shape to be (w, h): " + "(1024, 1024) with padding added top and bottom" + ) + + +def test_letterbox_image_for_pillow_image() -> None: + # given + image = Image.new(mode="RGB", size=(640, 480), color=(0, 0, 0)) + expected_result = Image.fromarray( + np.concatenate( + [ + np.ones((128, 1024, 3), dtype=np.uint8) * 255, + np.zeros((768, 1024, 3), dtype=np.uint8), + np.ones((128, 1024, 3), dtype=np.uint8) * 255, + ], + axis=0, + ) + ) + + # when + result = letterbox_image( + image=image, resolution_wh=(1024, 1024), color=(255, 255, 255) + ) + + # then + assert result.size == ( + 1024, + 1024, + ), "Expected output shape to be (w, h): (1024, 1024)" + difference = ImageChops.difference(result, expected_result) + assert ( + difference.getbbox() is None + ), "Expected padding to be added top and bottom with padding added top and bottom" + + +def test_create_tiles_with_one_image( + one_image: np.ndarray, single_image_tile: np.ndarray +) -> None: + # when + result = create_tiles(images=[one_image], single_tile_size=(240, 240)) + + # # then + assert np.allclose(result, single_image_tile, atol=5.0) + + +def test_create_tiles_with_one_image_and_enforced_grid( + one_image: np.ndarray, single_image_tile_enforced_grid: np.ndarray +) -> None: + # when + result = create_tiles( + images=[one_image], + grid_size=(None, 3), + single_tile_size=(240, 240), + ) + + # then + assert np.allclose(result, single_image_tile_enforced_grid, atol=5.0) + + +def test_create_tiles_with_two_images( + two_images: List[np.ndarray], two_images_tile: np.ndarray +) -> None: + # when + result = create_tiles(images=two_images, single_tile_size=(240, 240)) + + # then + assert np.allclose(result, two_images_tile, atol=5.0) + + +def test_create_tiles_with_three_images( + three_images: List[np.ndarray], three_images_tile: np.ndarray +) -> None: + # when + result = create_tiles(images=three_images, single_tile_size=(240, 240)) + + # then + assert np.allclose(result, three_images_tile, atol=5.0) + + +def test_create_tiles_with_four_images( + four_images: List[np.ndarray], + four_images_tile: np.ndarray, +) -> None: + # when + result = create_tiles(images=four_images, single_tile_size=(240, 240)) + + # then + assert np.allclose(result, four_images_tile, atol=5.0) + + +def test_create_tiles_with_all_images( + all_images: List[np.ndarray], + all_images_tile: np.ndarray, +) -> None: + # when + result = create_tiles(images=all_images, single_tile_size=(240, 240)) + + # then + assert np.allclose(result, all_images_tile, atol=5.0) + + +def test_create_tiles_with_all_images_and_custom_grid( + all_images: List[np.ndarray], all_images_tile_and_custom_grid: np.ndarray +) -> None: + # when + result = create_tiles( + images=all_images, + grid_size=(3, 3), + single_tile_size=(240, 240), + ) + + # then + assert np.allclose(result, all_images_tile_and_custom_grid, atol=5.0) + + +def test_create_tiles_with_all_images_and_custom_colors( + all_images: List[np.ndarray], all_images_tile_and_custom_colors: np.ndarray +) -> None: + # when + result = create_tiles( + images=all_images, + tile_margin_color=(127, 127, 127), + tile_padding_color=(224, 224, 224), + single_tile_size=(240, 240), + ) + + # then + assert np.allclose(result, all_images_tile_and_custom_colors, atol=5.0) + + +def test_create_tiles_with_all_images_and_titles( + all_images: List[np.ndarray], + all_images_tile_and_custom_colors_and_titles: np.ndarray, +) -> None: + # when + result = create_tiles( + images=all_images, + titles=["Image 1", None, "Image 3", "Image 4"], + single_tile_size=(240, 240), + ) + + # then + assert np.allclose(result, all_images_tile_and_custom_colors_and_titles, atol=5.0) + + +def test_create_tiles_with_all_images_and_titles_with_custom_configs( + all_images: List[np.ndarray], + all_images_tile_and_titles_with_custom_configs: np.ndarray, +) -> None: + # when + result = create_tiles( + images=all_images, + titles=["Image 1", None, "Image 3", "Image 4"], + single_tile_size=(240, 240), + titles_anchors=[ + Point(x=200, y=300), + Point(x=300, y=400), + None, + Point(x=300, y=400), + ], + titles_color=Color.RED, + titles_scale=1.5, + titles_thickness=3, + titles_padding=20, + titles_background_color=Color.BLACK, + default_title_placement="bottom", + ) + + # then + assert np.allclose(result, all_images_tile_and_titles_with_custom_configs, atol=5.0) + + +def test_create_tiles_with_all_images_and_custom_grid_to_small_to_fit_images( + all_images: List[np.ndarray], +) -> None: + with pytest.raises(ValueError): + _ = create_tiles(images=all_images, grid_size=(2, 2)) diff --git a/test/utils/test_iterables.py b/test/utils/test_iterables.py new file mode 100644 index 00000000..2d34605c --- /dev/null +++ b/test/utils/test_iterables.py @@ -0,0 +1,43 @@ +import pytest + +from supervision.utils.iterables import create_batches, fill + + +@pytest.mark.parametrize( + "sequence, batch_size, expected", + [ + # Empty sequence, non-zero batch size. Expect empty list. + ([], 4, []), + # Non-zero size sequence, batch size of 0. Each item is its own batch. + ([1, 2, 3], 0, [[1], [2], [3]]), + # Batch size larger than sequence. All items in a single batch. + ([1, 2], 4, [[1, 2]]), + # Batch size evenly divides the sequence. Equal size batches. + ([1, 2, 3, 4], 2, [[1, 2], [3, 4]]), + # Batch size doesn't evenly divide sequence. Last batch smaller. + ([1, 2, 3, 4], 3, [[1, 2, 3], [4]]), + ], +) +def test_create_batches(sequence, batch_size, expected) -> None: + result = list(create_batches(sequence=sequence, batch_size=batch_size)) + assert result == expected + + +@pytest.mark.parametrize( + "sequence, desired_size, content, expected", + [ + # Empty sequence, desired size 0. Expect empty list. + ([], 0, 1, []), + # Empty sequence, non-zero desired size. Filled with padding. + ([], 3, 1, [1, 1, 1]), + # Sequence at desired size. No changes. + ([2, 2, 2], 3, 1, [2, 2, 2]), + # Sequence exceeds desired size. No changes. + ([2, 2, 2, 2], 3, 1, [2, 2, 2, 2]), + # Non-empty sequence, shorter than desired. Padding added. + ([2], 3, 1, [2, 1, 1]), + ], +) +def test_fill(sequence, desired_size, content, expected) -> None: + result = fill(sequence=sequence, desired_size=desired_size, content=content) + assert result == expected