First batch of basic tools:
- 🎨 DEFAULT_COLOR_PALETTE, Color and ColorPalette classes - 📐 initial implementation of Point, Vector and Rect classes - 🎬 VideoInfo and VideoSink classes as well as get_video_frames_generator - 📓 show_frame_in_notebook util
This commit is contained in:
SkalskiP
parent
a4e60f33ba
commit
3dad4f4a4d
|
|
@ -131,3 +131,6 @@ dmypy.json
|
|||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# Notebooks
|
||||
notebooks/
|
||||
|
|
|
|||
156
README.md
156
README.md
|
|
@ -1,40 +1,82 @@
|
|||
# Python Template 🐍
|
||||
A template repo holding our common setup for a python project.
|
||||
<div align="center">
|
||||
<p>
|
||||
<a align="center" href="" target="_blank">
|
||||
<img
|
||||
width="850"
|
||||
src="https://media.roboflow.com/open-source/supervision/roboflow-supervision-banner.png?ik-sdk-version=javascript-1.4.3&updatedAt=1674062891088"
|
||||
>
|
||||
</a>
|
||||
</p>
|
||||
<br>
|
||||
|
||||
## Installation
|
||||
<div align="center">
|
||||
<a href="https://youtube.com/roboflow">
|
||||
<img
|
||||
src="https://media.roboflow.com/notebooks/template/icons/purple/youtube.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672949634652"
|
||||
width="3%"
|
||||
/>
|
||||
</a>
|
||||
<img src="https://github.com/SkalskiP/SkalskiP/blob/master/icons/transparent.png" width="3%"/>
|
||||
<a href="https://roboflow.com">
|
||||
<img
|
||||
src="https://media.roboflow.com/notebooks/template/icons/purple/roboflow-app.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672949746649"
|
||||
width="3%"
|
||||
/>
|
||||
</a>
|
||||
<img src="https://github.com/SkalskiP/SkalskiP/blob/master/icons/transparent.png" width="3%"/>
|
||||
<a href="https://www.linkedin.com/company/roboflow-ai/">
|
||||
<img
|
||||
src="https://media.roboflow.com/notebooks/template/icons/purple/linkedin.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672949633691"
|
||||
width="3%"
|
||||
/>
|
||||
</a>
|
||||
<img src="https://github.com/SkalskiP/SkalskiP/blob/master/icons/transparent.png" width="3%"/>
|
||||
<a href="https://docs.roboflow.com">
|
||||
<img
|
||||
src="https://media.roboflow.com/notebooks/template/icons/purple/knowledge.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672949634511"
|
||||
width="3%"
|
||||
/>
|
||||
</a>
|
||||
<img src="https://github.com/SkalskiP/SkalskiP/blob/master/icons/transparent.png" width="3%"/>
|
||||
<a href="https://disuss.roboflow.com">
|
||||
<img
|
||||
src="https://media.roboflow.com/notebooks/template/icons/purple/forum.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672949633584"
|
||||
width="3%"
|
||||
/>
|
||||
<img src="https://github.com/SkalskiP/SkalskiP/blob/master/icons/transparent.png" width="3%"/>
|
||||
<a href="https://blog.roboflow.com">
|
||||
<img
|
||||
src="https://media.roboflow.com/notebooks/template/icons/purple/blog.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672949633605"
|
||||
width="3%"
|
||||
/>
|
||||
</a>
|
||||
</a>
|
||||
</div>
|
||||
|
||||
You can install the package using pip
|
||||
</div>
|
||||
|
||||
```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.
|
||||
35
setup.py
35
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',
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1 +1 @@
|
|||
__version__ = "0.0.1"
|
||||
__version__ = "0.0.2"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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]
|
||||
|
|
@ -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()
|
||||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
@ -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()
|
||||
|
|
@ -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
|
||||
|
|
@ -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()
|
||||
|
|
@ -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()
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
Loading…
Reference in New Issue