From 3dad4f4a4dc1ec00d183dd87c722735d73dd9c35 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Wed, 18 Jan 2023 19:27:51 +0100 Subject: [PATCH] =?UTF-8?q?First=20batch=20of=20basic=20tools:=20-=20?= =?UTF-8?q?=F0=9F=8E=A8=20DEFAULT=5FCOLOR=5FPALETTE,=20Color=20and=20Color?= =?UTF-8?q?Palette=20classes=20-=20=F0=9F=93=90=20initial=20implementation?= =?UTF-8?q?=20of=20Point,=20Vector=20and=20Rect=20classes=20-=20?= =?UTF-8?q?=F0=9F=8E=AC=20VideoInfo=20and=20VideoSink=20classes=20as=20wel?= =?UTF-8?q?l=20as=20get=5Fvideo=5Fframes=5Fgenerator=20-=20=F0=9F=93=93=20?= =?UTF-8?q?show=5Fframe=5Fin=5Fnotebook=20util?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 3 + README.md | 156 +++++++++--------- setup.py | 35 ++-- supervision/__init__.py | 2 +- supervision/annotators/bbox.py | 11 -- supervision/annotators/dataclasses.py | 24 --- supervision/annotators/utils.py | 10 -- supervision/commons/dataclasses.py | 40 ++--- supervision/{annotators => draw}/__init__.py | 0 supervision/draw/color.py | 113 +++++++++++++ .../consts.py => notebook/__init__.py} | 0 supervision/notebook/utils.py | 22 +++ supervision/video/__init__.py | 0 supervision/video/dataclasses.py | 45 +++++ supervision/video/sink.py | 42 +++++ supervision/video/source.py | 20 +++ tests/__init__.py | 0 tests/commons/__init__.py | 0 tests/commons/test_dataclasses.py | 40 +++++ tests/draw/__init__.py | 0 tests/draw/test_color.py | 34 ++++ 21 files changed, 436 insertions(+), 161 deletions(-) delete mode 100644 supervision/annotators/bbox.py delete mode 100644 supervision/annotators/dataclasses.py delete mode 100644 supervision/annotators/utils.py rename supervision/{annotators => draw}/__init__.py (100%) create mode 100644 supervision/draw/color.py rename supervision/{annotators/consts.py => notebook/__init__.py} (100%) create mode 100644 supervision/notebook/utils.py create mode 100644 supervision/video/__init__.py create mode 100644 supervision/video/dataclasses.py create mode 100644 supervision/video/sink.py create mode 100644 supervision/video/source.py create mode 100644 tests/__init__.py create mode 100644 tests/commons/__init__.py create mode 100644 tests/commons/test_dataclasses.py create mode 100644 tests/draw/__init__.py create mode 100644 tests/draw/test_color.py diff --git a/.gitignore b/.gitignore index ca4e8eb4..44dc84e4 100644 --- a/.gitignore +++ b/.gitignore @@ -131,3 +131,6 @@ dmypy.json # Pyre type checker .pyre/ + +# Notebooks +notebooks/ diff --git a/README.md b/README.md index 23ecca97..76ec98c8 100644 --- a/README.md +++ b/README.md @@ -1,40 +1,82 @@ -# Python Template ๐Ÿ -A template repo holding our common setup for a python project. +
+

+ + + +

+
-## Installation +
+ + + + + + + + + + + + + + + + + + + + + + + +
-You can install the package using pip +
-```bash -pip install -e . +## ๐Ÿ‘‹ hello + +A set of easy-to-use utils that will come in handy in any Computer Vision project. **Supervision** is still in +pre-release stage. ๐Ÿšง Keep your eyes open for potential bugs and be aware that at this stage our API is still fluid +and may change. + +## ๐Ÿ’ป install + +```console +# clone repository and navigate to root directory +git clone git@github.com:roboflow/supervision.git +cd supervision + +# setup python environment and activate it +python3 -m venv venv +source venv/bin/activate + +# install +pip install -e ".[all]" ``` -or for development - -```bash -pip install -e ".[dev]" -``` - -## Structure - -The project has the following structure - -``` -โ”œโ”€โ”€ .github -โ”‚ โ””โ”€โ”€ workflows -โ”‚ โ””โ”€โ”€ test.yml # holds our github action config -โ”œโ”€โ”€ .gitignore -โ”œโ”€โ”€ Makefile -โ”œโ”€โ”€ README.md -โ”œโ”€โ”€ setup.py -โ”œโ”€โ”€ src -โ”‚ โ”œโ”€โ”€ __init__.py -โ”‚ โ”œโ”€โ”€ hello.py -โ””โ”€โ”€ test - โ””โ”€โ”€ test_hello.py -``` - -### Code Quality ๐Ÿงน +## ๐Ÿงน Code Quality We provide two handy commands inside the `Makefile`, namely: @@ -43,54 +85,6 @@ We provide two handy commands inside the `Makefile`, namely: So far, **there is no types checking with mypy**. See [issue](https://github.com/roboflow-ai/template-python/issues/4). -### Tests ๐Ÿงช +## ๐Ÿงช Tests [`pytests`](https://docs.pytest.org/en/7.1.x/) is used to run our tests. - -### Publish on PyPi ๐Ÿš€ - -**Important**: Before publishing, edit `__version__` in [src/__init__](/src/__init__.py) to match the wanted new version. - -We use [`twine`](https://twine.readthedocs.io/en/stable/) to make our life easier. You can publish by using - -``` -export PYPI_USERNAME="you_username" -export PYPI_PASSWORD="your_password" -export PYPI_TEST_PASSWORD="your_password_for_test_pypi" -make publish -e PYPI_USERNAME=$PYPI_USERNAME -e PYPI_PASSWORD=$PYPI_PASSWORD -e PYPI_TEST_PASSWORD=$PYPI_TEST_PASSWORD -``` - -You can also use token for auth, see [pypi doc](https://pypi.org/help/#apitoken). In that case, - -``` -export PYPI_USERNAME="__token__" -export PYPI_PASSWORD="your_token" -export PYPI_TEST_PASSWORD="your_token_for_test_pypi" -make publish -e PYPI_USERNAME=$PYPI_USERNAME -e PYPI_PASSWORD=$PYPI_PASSWORD -e PYPI_TEST_PASSWORD=$PYPI_TEST_PASSWORD -``` - -**Note**: We will try to push to [test pypi](https://test.pypi.org/) before pushing to pypi, to assert everything will work - -### CI/CD ๐Ÿค– - -We use [GitHub actions](https://github.com/features/actions) to automatically run tests and check code quality when a new PR is done on `main`. - -On any pull request, we will check the code quality and tests. - -When a new release is created, we will try to push the new code to PyPi. We use [`twine`](https://twine.readthedocs.io/en/stable/) to make our life easier. - -The **correct steps** to create a new realease are the following: -- edit `__version__` in [src/__init__](/src/__init__.py) to match the wanted new version. -- create a new [`tag`](https://git-scm.com/docs/git-tag) with the release name, e.g. `git tag v0.0.1 && git push origin v0.0.1` or from the GitHub UI. -- create a new release from GitHub UI - -The CI will run when you create the new release. - -# Q&A - -## Why no cookiecutter? -This is a template repo, it's meant to be used inside GitHub upon repo creation. - -## Why reinvent the wheel? - -There are several very good templates on GitHub, I prefer to use code we wrote instead of blinding taking the most starred template and having features we don't need. From experience, it's better to keep it simple and general enough for our specific use cases. \ No newline at end of file diff --git a/setup.py b/setup.py index 6d5f2631..e8860df4 100644 --- a/setup.py +++ b/setup.py @@ -1,32 +1,34 @@ import setuptools from setuptools import find_packages import re +from pathlib import Path -with open('./supervision/__init__.py', 'r') as f: - content = f.read() - # from https://www.py4u.net/discuss/139845 - version = re.search(r'__version__\s*=\s*[\'"]([^\'"]*)[\'"]', content).group(1) +FILE = Path(__file__).resolve() +PARENT = FILE.parent # root directory +README = (PARENT / "README.md").read_text(encoding="utf-8") + + +def get_version(): + file = PARENT / 'supervision/__init__.py' + return re.search(r'^__version__ = [\'"]([^\'"]*)[\'"]', file.read_text(encoding="utf-8"), re.M)[1] -with open('README.md', 'r') as fh: - long_description = fh.read() setuptools.setup( name='supervision', - version=version, + version=get_version(), author='Piotr Skalski', author_email='piotr.skalski92@gmail.com', license='MIT', description='A set of easy-to-use utils that will come in handy in any Computer Vision project', - long_description=long_description, + long_description=README, long_description_content_type='text/markdown', - url='https://github.com/roboflow-ai/supervision', - install_requires=[], - packages=find_packages(exclude=("tests",)), - extras_require={ - 'annotators': [ + url='https://github.com/roboflow/supervision', + install_requires=[ 'numpy', 'opencv-python' - ], + ], + packages=find_packages(exclude=("tests",)), + extras_require={ 'dev': [ 'flake8', 'black==22.3.0', @@ -38,7 +40,6 @@ setuptools.setup( ], }, classifiers=[ - 'Development Status :: 5 - Production/Stable', 'Intended Audience :: Developers', 'Intended Audience :: Science/Research', 'License :: OSI Approved :: BSD License', @@ -46,14 +47,18 @@ setuptools.setup( 'Programming Language :: Python :: 3.7', 'Programming Language :: Python :: 3.8', 'Programming Language :: Python :: 3.9', + 'Programming Language :: Python :: 3.10', 'Programming Language :: Python :: 3 :: Only', 'Topic :: Software Development', 'Topic :: Scientific/Engineering', + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Scientific/Engineering :: Image Recognition", 'Typing :: Typed', 'Operating System :: Microsoft :: Windows', 'Operating System :: POSIX', 'Operating System :: Unix', 'Operating System :: MacOS' ], + keywords="machine-learning, deep-learning, vision, ML, DL, AI, YOLOv5, YOLOv8, Roboflow", python_requires='>=3.7', ) diff --git a/supervision/__init__.py b/supervision/__init__.py index f102a9ca..3b93d0be 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -1 +1 @@ -__version__ = "0.0.1" +__version__ = "0.0.2" diff --git a/supervision/annotators/bbox.py b/supervision/annotators/bbox.py deleted file mode 100644 index 91564ebc..00000000 --- a/supervision/annotators/bbox.py +++ /dev/null @@ -1,11 +0,0 @@ -from typing import List - -import numpy as np - -from supervision.commons.dataclasses import Detection - - -class BoxAnnotator: - - def annotate(self, image: np.ndarray, detections: List[Detection]) -> np.ndarray: - pass diff --git a/supervision/annotators/dataclasses.py b/supervision/annotators/dataclasses.py deleted file mode 100644 index 52b99593..00000000 --- a/supervision/annotators/dataclasses.py +++ /dev/null @@ -1,24 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass - -from typing import Tuple, List - - -@dataclass -class Color: - r: int - g: int - b: int - - @classmethod - def from_hex_string(cls, hex_string: str) -> Color: - pass - - def as_bgr_tuple(self) -> Tuple[int, int, int]: - return self.r, self.g, self.b - - -@dataclass -class ColorPalette: - colors: List[Color] diff --git a/supervision/annotators/utils.py b/supervision/annotators/utils.py deleted file mode 100644 index de26ed41..00000000 --- a/supervision/annotators/utils.py +++ /dev/null @@ -1,10 +0,0 @@ -import numpy as np - -import cv2 - -from supervision.annotators.dataclasses import Color -from supervision.commons.dataclasses import Rect - - -def draw_rect(image: np.ndarray, rect: Rect, color: Color, thickness: int) -> np.ndarray: - cv2.rectangle() diff --git a/supervision/commons/dataclasses.py b/supervision/commons/dataclasses.py index 09c96fd9..81b3f821 100644 --- a/supervision/commons/dataclasses.py +++ b/supervision/commons/dataclasses.py @@ -1,9 +1,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Optional, Tuple - -import numpy as np +from typing import Tuple @dataclass @@ -18,6 +16,19 @@ class Point: return self.x, self.y +@dataclass +class Vector: + start: Point + end: Point + + def is_in(self, point: Point) -> bool: + v1 = Vector(self.start, self.end) + v2 = Vector(self.start, point) + cross_product = (v1.end.x - v1.start.x) * (v2.end.y - v2.start.y) - \ + (v1.end.y - v1.start.y) * (v2.end.x - v2.start.x) + return cross_product < 0 + + @dataclass class Rect: x: float @@ -33,19 +44,10 @@ class Rect: def bottom_right(self) -> Point: return Point(x=self.x + self.width, y=self.y + self.height) - -@dataclass -class Detection: - x_min: float - x_max: float - y_min: float - y_max: float - class_id: int - class_name: Optional[str] - confidence: Optional[float] - mask: Optional[np.ndarray] - contour: Optional[np.ndarray] - - @property - def rect(self) -> Rect: - return Rect(x=self.x_min, y=self.y_min, width=self.x_max - self.x_min, height=self.y_max - self.y_min) + def pad(self, padding) -> Rect: + return Rect( + x=self.x - padding, + y=self.y - padding, + width=self.width + 2 * padding, + height=self.height + 2 * padding, + ) diff --git a/supervision/annotators/__init__.py b/supervision/draw/__init__.py similarity index 100% rename from supervision/annotators/__init__.py rename to supervision/draw/__init__.py diff --git a/supervision/draw/color.py b/supervision/draw/color.py new file mode 100644 index 00000000..84f3fc42 --- /dev/null +++ b/supervision/draw/color.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Tuple, List + +DEFAULT_COLOR_PALETTE = [ + "#e6194b", "#3cb44b", "#ffe119", "#0082c8", "#f58231", "#911eb4", "#46f0f0", "#f032e6", + "#d2f53c", "#fabebe", "#008080", "#e6beff", "#aa6e28", "#fffac8", "#800000", "#aaffc3" +] + + +def _validate_color_hex(color_hex: str): + color_hex = color_hex.lstrip('#') + if not all(c in '0123456789abcdefABCDEF' for c in color_hex): + raise ValueError('Invalid characters in color hash') + if len(color_hex) not in (3, 6): + raise ValueError('Invalid length of color hash') + + +@dataclass +class Color: + r: int + g: int + b: int + + @classmethod + def from_hex(cls, color_hex: str): + """ + Creates a Color instance from a color hex string + + :param color_hex: str : The color hex string in the format of "fff", "ffffff", "#fff", or "#ffffff" + :return: Color : A Color instance representing the color + + Example: + color = Color.from_hex('#ff00ff') + """ + _validate_color_hex(color_hex) + color_hex = color_hex.lstrip('#') + if len(color_hex) == 3: + color_hex = ''.join(c * 2 for c in color_hex) + r, g, b = (int(color_hex[i:i + 2], 16) for i in range(0, 6, 2)) + return cls(r, g, b) + + def as_rgb(self) -> Tuple[int, int, int]: + """ + Returns the color as a tuple of integers in the RGB format + + :return: Tuple[int, int, int] : The color in the RGB format + """ + return self.r, self.g, self.b + + def as_bgr(self) -> Tuple[int, int, int]: + """ + Returns the color as a tuple of integers in the BGR format + + :return: Tuple[int, int, int] : The color in the BGR format + """ + return self.b, self.g, self.r + + @classmethod + def white(cls) -> Color: + return Color.from_hex(color_hex='#ffffff') + + @classmethod + def black(cls) -> Color: + return Color.from_hex(color_hex='#000000') + + @classmethod + def red(cls) -> Color: + return Color.from_hex(color_hex='#ff0000') + + @classmethod + def green(cls) -> Color: + return Color.from_hex(color_hex='#00ff00') + + @classmethod + def blue(cls) -> Color: + return Color.from_hex(color_hex='#0000ff') + + +@dataclass +class ColorPalette: + colors: List[Color] = field(default_factory=lambda: DEFAULT_COLOR_PALETTE) + + @classmethod + def from_hex(cls, color_hex_list: List[str]): + """ + Creates a ColorPalette instance from a list of color hex strings + + :param color_hex_list: List[str] : A list of color hex strings in the format of "fff", "ffffff", "#fff", or "#ffffff" + :return: ColorPalette : A ColorPalette instance representing the color palette + + Example: + color_palette = ColorPalette.from_hex(['#ff0000', '#00ff00', '#0000ff']) + """ + colors = [Color.from_hex(color_hex) for color_hex in color_hex_list] + return cls(colors) + + def by_idx(self, idx: int) -> Color: + """ + Returns the color at a given index in the color palette. + + :param idx: int : The index of the color in the color palette + :return: Color : The color at the given index + + Example: + color_palette = ColorPalette.from_hex(['#ff0000', '#00ff00', '#0000ff']) + color = color_palette.by_idx(1) + """ + if idx < 0: + raise ValueError("idx argument should not be negative") + idx = idx % len(self.colors) + return self.colors[idx] diff --git a/supervision/annotators/consts.py b/supervision/notebook/__init__.py similarity index 100% rename from supervision/annotators/consts.py rename to supervision/notebook/__init__.py diff --git a/supervision/notebook/utils.py b/supervision/notebook/utils.py new file mode 100644 index 00000000..55cce194 --- /dev/null +++ b/supervision/notebook/utils.py @@ -0,0 +1,22 @@ +from typing import Tuple + +import cv2 +import matplotlib.pyplot as plt +import numpy as np + + +def show_frame_in_notebook(frame: np.ndarray, size: Tuple[int, int] = (10, 10), cmap: str = 'gray'): + """ + Display a frame in Jupyter Notebook using Matplotlib + + :param frame: np.ndarray : The frame to be displayed. + :param size: Tuple[int, int] : The size of the plot. default:(10,10) + :param cmap: str : the colormap to use for single channel images. default:gray + """ + if frame.ndim == 2: + plt.figure(figsize=size) + plt.imshow(frame, cmap=cmap) + else: + plt.figure(figsize=size) + plt.imshow(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + plt.show() diff --git a/supervision/video/__init__.py b/supervision/video/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/supervision/video/dataclasses.py b/supervision/video/dataclasses.py new file mode 100644 index 00000000..b588fe6f --- /dev/null +++ b/supervision/video/dataclasses.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Tuple + +import cv2 + + +@dataclass +class VideoInfo: + """ + Data class containing information about the video resolution, fps, and total frame count. + + :param width: int : The width of the video frames in pixels. + :param height: int : The height of the video frames in pixels. + :param fps: int : The frames per second of the video. + :param total_frames: int : The total number of frames in the video. + """ + width: int + height: int + fps: int + total_frames: Optional[int] = None + + @classmethod + def from_video_path(cls, video_path: str) -> VideoInfo: + """ + Returns a VideoInfo data class containing information about the video resolution, fps, and total frame count. + + :param video_path: str : The path of the video file. + :return: VideoInfo : A data class containing information about the video resolution, fps, and total frame count. + """ + video = cv2.VideoCapture(video_path) + if not video.isOpened(): + raise Exception(f"Could not open video at {video_path}") + + width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) + fps = int(video.get(cv2.CAP_PROP_FPS)) + total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) + video.release() + return VideoInfo(width, height, fps, total_frames) + + @property + def resolution(self) -> Tuple[int, int]: + return self.width, self.height diff --git a/supervision/video/sink.py b/supervision/video/sink.py new file mode 100644 index 00000000..9a7304fb --- /dev/null +++ b/supervision/video/sink.py @@ -0,0 +1,42 @@ +import cv2 +import numpy as np + +from supervision.video.dataclasses import VideoInfo + + +class VideoSink: + """ + A context manager that uses OpenCV to save video frames to a file. + + :param output_path: str : The path to the output file where the video will be saved. + :param video_info: VideoInfo : An instance of VideoInfo containing information about the video resolution, fps, and total frame count. + """ + def __init__(self, output_path: str, video_info: VideoInfo): + """ + Initializes the VideoSink with the specified output path and video information. + """ + self.output_path = output_path + self.video_info = video_info + self.fourcc = cv2.VideoWriter_fourcc(*'mp4v') + self.writer = None + + def __enter__(self): + """ + Opens the output file and returns the VideoSink instance. + """ + self.writer = cv2.VideoWriter(self.output_path, self.fourcc, self.video_info.fps, self.video_info.resolution) + return self + + def write_frame(self, frame: np.ndarray): + """ + Writes a frame to the output video file. + + :param frame: np.ndarray : The frame to be written. + """ + self.writer.write(frame) + + def __exit__(self, exc_type, exc_val, exc_tb): + """ + Closes the output file. + """ + self.writer.release() diff --git a/supervision/video/source.py b/supervision/video/source.py new file mode 100644 index 00000000..f3679cd7 --- /dev/null +++ b/supervision/video/source.py @@ -0,0 +1,20 @@ +from typing import Generator + +import cv2 + + +def get_video_frames_generator(video_path: str) -> Generator[int, None, None]: + """ + Returns a generator that yields the frames of the video. + + :param video_path: str : The path of the video file. + :return: Generator[int, None, None] : Generator that yields the frames of the video. + """ + video = cv2.VideoCapture(video_path) + if not video.isOpened(): + raise Exception(f"Could not open video at {video_path}") + success, frame = video.read() + while success: + yield frame + success, frame = video.read() + video.release() diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/commons/__init__.py b/tests/commons/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/commons/test_dataclasses.py b/tests/commons/test_dataclasses.py new file mode 100644 index 00000000..ebb88fc1 --- /dev/null +++ b/tests/commons/test_dataclasses.py @@ -0,0 +1,40 @@ +import pytest + +from supervision.commons.dataclasses import Vector, Point + + +@pytest.mark.parametrize( + 'vector, point, expected_result', + [ + (Vector(start=Point(x=0, y=0), end=Point(x=5, y=5)), Point(x=-1, y=1), False), + (Vector(start=Point(x=0, y=0), end=Point(x=5, y=5)), Point(x=6, y=6), False), + (Vector(start=Point(x=0, y=0), end=Point(x=5, y=5)), Point(x=3, y=6), False), + + (Vector(start=Point(x=5, y=5), end=Point(x=0, y=0)), Point(x=-1, y=1), True), + (Vector(start=Point(x=5, y=5), end=Point(x=0, y=0)), Point(x=6, y=6), False), + (Vector(start=Point(x=5, y=5), end=Point(x=0, y=0)), Point(x=3, y=6), True), + + (Vector(start=Point(x=0, y=0), end=Point(x=1, y=0)), Point(x=0, y=0), False), + (Vector(start=Point(x=0, y=0), end=Point(x=1, y=0)), Point(x=0, y=-1), True), + (Vector(start=Point(x=0, y=0), end=Point(x=1, y=0)), Point(x=0, y=1), False), + + (Vector(start=Point(x=1, y=0), end=Point(x=0, y=0)), Point(x=0, y=0), False), + (Vector(start=Point(x=1, y=0), end=Point(x=0, y=0)), Point(x=0, y=-1), False), + (Vector(start=Point(x=1, y=0), end=Point(x=0, y=0)), Point(x=0, y=1), True), + + (Vector(start=Point(x=1, y=1), end=Point(x=1, y=3)), Point(x=0, y=0), False), + (Vector(start=Point(x=1, y=1), end=Point(x=1, y=3)), Point(x=1, y=4), False), + (Vector(start=Point(x=1, y=1), end=Point(x=1, y=3)), Point(x=2, y=4), True), + + (Vector(start=Point(x=1, y=3), end=Point(x=1, y=1)), Point(x=0, y=0), True), + (Vector(start=Point(x=1, y=3), end=Point(x=1, y=1)), Point(x=1, y=4), False), + (Vector(start=Point(x=1, y=3), end=Point(x=1, y=1)), Point(x=2, y=4), False), + ] +) +def test_vector_is_in( + vector: Vector, + point: Point, + expected_result: bool +) -> None: + result = vector.is_in(point=point) + assert result == expected_result diff --git a/tests/draw/__init__.py b/tests/draw/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/draw/test_color.py b/tests/draw/test_color.py new file mode 100644 index 00000000..5ea48abb --- /dev/null +++ b/tests/draw/test_color.py @@ -0,0 +1,34 @@ +from contextlib import ExitStack as DoesNotRaise +from typing import Optional + +import pytest + +from supervision.draw.color import Color + + +@pytest.mark.parametrize( + 'color_hex, expected_result, exception', + [ + ('fff', Color.white(), DoesNotRaise()), + ('#fff', Color.white(), DoesNotRaise()), + ('ffffff', Color.white(), DoesNotRaise()), + ('#ffffff', Color.white(), DoesNotRaise()), + ('f00', Color.red(), DoesNotRaise()), + ('0f0', Color.green(), DoesNotRaise()), + ('00f', Color.blue(), DoesNotRaise()), + ('#808000', Color(r=128, g=128, b=0), DoesNotRaise()), + ('', None, pytest.raises(ValueError)), + ('00', None, pytest.raises(ValueError)), + ('0000', None, pytest.raises(ValueError)), + ('0000000', None, pytest.raises(ValueError)), + ('ffg', None, pytest.raises(ValueError)), + ] +) +def test_color_from_hex( + color_hex, + expected_result: Optional[Color], + exception: Exception +) -> None: + with exception: + result = Color.from_hex(color_hex=color_hex) + assert result == expected_result