From af82b4e2d232644df1aa94dd5c52ce68053b88cd Mon Sep 17 00:00:00 2001 From: Onuralp SEZER Date: Tue, 25 Jul 2023 22:28:26 +0300 Subject: [PATCH] =?UTF-8?q?fix:=20=F0=9F=90=9B=20200+=20E501=20doc=20error?= =?UTF-8?q?=20lines=20are=20fixed.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Onuralp SEZER --- .pre-commit-config.yaml | 37 ++-- supervision/classification/core.py | 12 +- supervision/dataset/core.py | 181 +++++++++++++------- supervision/dataset/formats/pascal_voc.py | 25 ++- supervision/dataset/formats/yolo.py | 17 +- supervision/detection/annotate.py | 25 ++- supervision/detection/core.py | 147 ++++++++++------ supervision/detection/line_counter.py | 3 +- supervision/detection/tools/polygon_zone.py | 21 ++- supervision/detection/utils.py | 105 ++++++++---- supervision/draw/color.py | 6 +- supervision/draw/utils.py | 16 +- supervision/geometry/utils.py | 11 +- supervision/metrics/detection.py | 67 +++++--- supervision/utils/file.py | 10 +- supervision/utils/image.py | 20 ++- supervision/utils/notebook.py | 9 +- supervision/utils/video.py | 35 ++-- test/dataset/formats/test_coco.py | 9 +- test/dataset/formats/test_yolo.py | 3 +- test/detection/test_utils.py | 6 +- 21 files changed, 506 insertions(+), 259 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5eb674dd..c5374eee 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -51,22 +51,22 @@ repos: - mdformat-black exclude: "README.md|CITATION|CONTRIBUTING.md" - - repo: https://github.com/PyCQA/flake8 - rev: 6.0.0 - hooks: - - id: flake8 - name: PEP8 Max Line Length Check - args: - - --count --max-line-length=88 --exit-zero --ignore=D --extend-ignore=E203,E501,W503 --statistics + # - repo: https://github.com/PyCQA/flake8 + # rev: 6.0.0 + # hooks: + # - id: flake8 + # name: PEP8 Max Line Length Check + # args: + # - --count --max-line-length=88 --exit-zero --ignore=D --extend-ignore=E203,E501,W503 --statistics - - repo: https://github.com/PyCQA/flake8 - rev: 6.0.0 - hooks: - - id: flake8 - name: PEP8 - args: - - --count --max-line-length=88 --select=E9,F63,F7,F82 --show-source --statistics + # - repo: https://github.com/PyCQA/flake8 + # rev: 6.0.0 + # hooks: + # - id: flake8 + # name: PEP8 + # args: + # - --count --max-line-length=88 --select=E9,F63,F7,F82 --show-source --statistics - repo: https://github.com/PyCQA/bandit rev: '1.7.5' @@ -74,10 +74,11 @@ repos: - id: bandit args: ["-c", "pyproject.toml"] additional_dependencies: ["bandit[toml]"] - - repo: https://github.com/PyCQA/autoflake - rev: v2.2.0 - hooks: - - id: autoflake + + # - repo: https://github.com/PyCQA/autoflake + # rev: v2.2.0 + # hooks: + # - id: autoflake - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.0.280 diff --git a/supervision/classification/core.py b/supervision/classification/core.py index 2f654b79..5ecd5b8f 100644 --- a/supervision/classification/core.py +++ b/supervision/classification/core.py @@ -42,10 +42,12 @@ class Classifications: @classmethod def from_yolov8(cls, yolov8_results) -> Classifications: """ - Creates a Classifications instance from a [YOLOv8](https://github.com/ultralytics/ultralytics) inference result. + Creates a Classifications instance from a + [YOLOv8](https://github.com/ultralytics/ultralytics) inference result. Args: - yolov8_results (ultralytics.yolo.engine.results.Results): The output Results instance from YOLOv8 + yolov8_results (ultralytics.yolo.engine.results.Results): + The output Results instance from YOLOv8 Returns: Detections: A new Classifications object. @@ -67,13 +69,15 @@ class Classifications: def get_top_k(self, k: int) -> Tuple[np.ndarray, np.ndarray]: """ - Retrieve the top k class IDs and confidences, ordered in descending order by confidence. + Retrieve the top k class IDs and confidences, + ordered in descending order by confidence. Args: k (int): The number of top class IDs and confidences to retrieve. Returns: - Tuple[np.ndarray, np.ndarray]: A tuple containing the top k class IDs and confidences. + Tuple[np.ndarray, np.ndarray]: A tuple containing + the top k class IDs and confidences. Example: ```python diff --git a/supervision/dataset/core.py b/supervision/dataset/core.py index ce7e0edb..68e258b8 100644 --- a/supervision/dataset/core.py +++ b/supervision/dataset/core.py @@ -55,7 +55,8 @@ class DetectionDataset(BaseDataset): Attributes: classes (List[str]): List containing dataset class names. images (Dict[str, np.ndarray]): Dictionary mapping image name to image. - annotations (Dict[str, Detections]): Dictionary mapping image name to annotations. + annotations (Dict[str, Detections]): Dictionary mapping + image name to annotations. """ classes: List[str] @@ -76,8 +77,9 @@ class DetectionDataset(BaseDataset): Iterate over the images and annotations in the dataset. Yields: - Iterator[Tuple[str, np.ndarray, Detections]]: An iterator that yields tuples containing the image name, - the image data, and its corresponding annotation. + Iterator[Tuple[str, np.ndarray, Detections]]: + An iterator that yields tuples containing the image name, + the image data, and its corresponding annotation. """ for image_name, image in self.images.items(): yield image_name, image, self.annotations.get(image_name, None) @@ -101,22 +103,27 @@ class DetectionDataset(BaseDataset): self, split_ratio=0.8, random_state=None, shuffle: bool = True ) -> Tuple[DetectionDataset, DetectionDataset]: """ - Splits the dataset into two parts (training and testing) using the provided split_ratio. + Splits the dataset into two parts (training and testing) + using the provided split_ratio. Args: - split_ratio (float, optional): The ratio of the training set to the entire dataset. - random_state (int, optional): The seed for the random number generator. This is used for reproducibility. + split_ratio (float, optional): The ratio of the training + set to the entire dataset. + random_state (int, optional): The seed for the random number generator. + This is used for reproducibility. shuffle (bool, optional): Whether to shuffle the data before splitting. Returns: - Tuple[DetectionDataset, DetectionDataset]: A tuple containing the training and testing datasets. + Tuple[DetectionDataset, DetectionDataset]: A tuple containing + the training and testing datasets. Example: ```python >>> import supervision as sv >>> ds = sv.DetectionDataset(...) - >>> train_ds, test_ds = ds.split(split_ratio=0.7, random_state=42, shuffle=True) + >>> train_ds, test_ds = ds.split(split_ratio=0.7, + ... random_state=42, shuffle=True) >>> len(train_ds), len(test_ds) (700, 300) ``` @@ -151,19 +158,27 @@ class DetectionDataset(BaseDataset): approximation_percentage: float = 0.0, ) -> None: """ - Exports the dataset to PASCAL VOC format. This method saves the images and their corresponding annotations in - PASCAL VOC format. + Exports the dataset to PASCAL VOC format. This method saves the images + and their corresponding annotations in PASCAL VOC format. Args: - images_directory_path (Optional[str]): The path to the directory where the images should be saved. + images_directory_path (Optional[str]): The path to the directory + where the images should be saved. If not provided, images will not be saved. - annotations_directory_path (Optional[str]): The path to the directory where the annotations in - PASCAL VOC format should be saved. If not provided, annotations will not be saved. - min_image_area_percentage (float): The minimum percentage of detection area relative to - the image area for a detection to be included. Argument is used only for segmentation datasets. - max_image_area_percentage (float): The maximum percentage of detection area relative to - the image area for a detection to be included. Argument is used only for segmentation datasets. - approximation_percentage (float): The percentage of polygon points to be removed from the input polygon, + annotations_directory_path (Optional[str]): The path to + the directory where the annotations in + PASCAL VOC format should be saved. If not provided, + annotations will not be saved. + min_image_area_percentage (float): The minimum percentage of + detection area relative to + the image area for a detection to be included. + Argument is used only for segmentation datasets. + max_image_area_percentage (float): The maximum percentage + of detection area relative to + the image area for a detection to be included. + Argument is used only for segmentation datasets. + approximation_percentage (float): The percentage of + polygon points to be removed from the input polygon, in the range [0, 1). Argument is used only for segmentation datasets. """ if images_directory_path: @@ -203,11 +218,14 @@ class DetectionDataset(BaseDataset): Creates a Dataset instance from PASCAL VOC formatted data. Args: - images_directory_path (str): The path to the directory containing the images. - annotations_directory_path (str): The path to the directory containing the PASCAL VOC XML annotations. + images_directory_path (str): The path to the + directory containing the images. + annotations_directory_path (str): The path to the directory + containing the PASCAL VOC XML annotations. Returns: - DetectionDataset: A DetectionDataset instance containing the loaded images and annotations. + DetectionDataset: A DetectionDataset instance containing + the loaded images and annotations. Example: ```python @@ -273,13 +291,19 @@ class DetectionDataset(BaseDataset): Creates a Dataset instance from YOLO formatted data. Args: - images_directory_path (str): The path to the directory containing the images. - annotations_directory_path (str): The path to the directory containing the YOLO annotation files. - data_yaml_path (str): The path to the data YAML file containing class information. - force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present. + images_directory_path (str): The path to the + directory containing the images. + annotations_directory_path (str): The path to the directory + containing the YOLO annotation files. + data_yaml_path (str): The path to the data + YAML file containing class information. + force_masks (bool, optional): If True, forces + masks to be loaded for all annotations, + regardless of whether they are present. Returns: - DetectionDataset: A DetectionDataset instance containing the loaded images and annotations. + DetectionDataset: A DetectionDataset instance + containing the loaded images and annotations. Example: ```python @@ -322,23 +346,32 @@ class DetectionDataset(BaseDataset): approximation_percentage: float = 0.0, ) -> None: """ - Exports the dataset to YOLO format. This method saves the images and their corresponding - annotations in YOLO format. + Exports the dataset to YOLO format. This method saves the + images and their corresponding annotations in YOLO format. Args: - images_directory_path (Optional[str]): The path to the directory where the images should be saved. + images_directory_path (Optional[str]): The path to the + directory where the images should be saved. If not provided, images will not be saved. - annotations_directory_path (Optional[str]): The path to the directory where the annotations in - YOLO format should be saved. If not provided, annotations will not be saved. - data_yaml_path (Optional[str]): The path where the data.yaml file should be saved. + annotations_directory_path (Optional[str]): The path to the + directory where the annotations in + YOLO format should be saved. If not provided, + annotations will not be saved. + data_yaml_path (Optional[str]): The path where the data.yaml + file should be saved. If not provided, the file will not be saved. - min_image_area_percentage (float): The minimum percentage of detection area relative to - the image area for a detection to be included. Argument is used only for segmentation datasets. - max_image_area_percentage (float): The maximum percentage of detection area relative to - the image area for a detection to be included. Argument is used only for segmentation datasets. - approximation_percentage (float): The percentage of polygon points to be removed from the input polygon, - in the range [0, 1). This is useful for simplifying the annotations. Argument is used only for - segmentation datasets. + min_image_area_percentage (float): The minimum percentage of + detection area relative to + the image area for a detection to be included. + Argument is used only for segmentation datasets. + max_image_area_percentage (float): The maximum percentage + of detection area relative to + the image area for a detection to be included. + Argument is used only for segmentation datasets. + approximation_percentage (float): The percentage of polygon points to + be removed from the input polygon, in the range [0, 1). + This is useful for simplifying the annotations. + Argument is used only for segmentation datasets. """ if images_directory_path is not None: save_dataset_images( @@ -367,12 +400,16 @@ class DetectionDataset(BaseDataset): Creates a Dataset instance from COCO formatted data. Args: - images_directory_path (str): The path to the directory containing the images. + images_directory_path (str): The path to the + directory containing the images. annotations_path (str): The path to the json annotation files. - force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present. + force_masks (bool, optional): If True, + forces masks to be loaded for all annotations, + regardless of whether they are present. Returns: - DetectionDataset: A DetectionDataset instance containing the loaded images and annotations. + DetectionDataset: A DetectionDataset instance containing + the loaded images and annotations. Example: ```python @@ -412,20 +449,26 @@ class DetectionDataset(BaseDataset): approximation_percentage: float = 0.0, ) -> None: """ - Exports the dataset to COCO format. This method saves the images and their corresponding - annotations in COCO format. + Exports the dataset to COCO format. This method saves the + images and their corresponding annotations in COCO format. Args: - images_directory_path (Optional[str]): The path to the directory where the images should be saved. + images_directory_path (Optional[str]): The path to the directory + where the images should be saved. If not provided, images will not be saved. annotations_path (Optional[str]): The path to COCO annotation file. - min_image_area_percentage (float): The minimum percentage of detection area relative to - the image area for a detection to be included. Argument is used only for segmentation datasets. - max_image_area_percentage (float): The maximum percentage of detection area relative to - the image area for a detection to be included. Argument is used only for segmentation datasets. - approximation_percentage (float): The percentage of polygon points to be removed from the input polygon, - in the range [0, 1). This is useful for simplifying the annotations. Argument is used only for - segmentation datasets. + min_image_area_percentage (float): The minimum percentage of + detection area relative to + the image area for a detection to be included. + Argument is used only for segmentation datasets. + max_image_area_percentage (float): The maximum percentage of + detection area relative to + the image area for a detection to be included. + Argument is used only for segmentation datasets. + approximation_percentage (float): The percentage of polygon points + to be removed from the input polygon, + in the range [0, 1). This is useful for simplifying the annotations. + Argument is used only for segmentation datasets. """ if images_directory_path is not None: save_dataset_images( @@ -445,16 +488,20 @@ class DetectionDataset(BaseDataset): @classmethod def merge(cls, dataset_list: List[DetectionDataset]) -> DetectionDataset: """ - Merge a list of `DetectionDataset` objects into a single `DetectionDataset` object. + Merge a list of `DetectionDataset` objects into a single + `DetectionDataset` object. - This method takes a list of `DetectionDataset` objects and combines their respective fields (`classes`, `images`, + This method takes a list of `DetectionDataset` objects and combines + their respective fields (`classes`, `images`, `annotations`) into a single `DetectionDataset` object. Args: - dataset_list (List[DetectionDataset]): A list of `DetectionDataset` objects to merge. + dataset_list (List[DetectionDataset]): A list of `DetectionDataset` + objects to merge. Returns: - (DetectionDataset): A single `DetectionDataset` object containing the merged data from the input list. + (DetectionDataset): A single `DetectionDataset` object containing + the merged data from the input list. Example: ```python @@ -512,7 +559,8 @@ class ClassificationDataset(BaseDataset): Attributes: classes (List[str]): List containing dataset class names. images (Dict[str, np.ndarray]): Dictionary mapping image name to image. - annotations (Dict[str, Detections]): Dictionary mapping image name to annotations. + annotations (Dict[str, Detections]): Dictionary mapping + image name to annotations. """ classes: List[str] @@ -526,22 +574,28 @@ class ClassificationDataset(BaseDataset): self, split_ratio=0.8, random_state=None, shuffle: bool = True ) -> Tuple[ClassificationDataset, ClassificationDataset]: """ - Splits the dataset into two parts (training and testing) using the provided split_ratio. + Splits the dataset into two parts (training and testing) + using the provided split_ratio. Args: - split_ratio (float, optional): The ratio of the training set to the entire dataset. - random_state (int, optional): The seed for the random number generator. This is used for reproducibility. + split_ratio (float, optional): The ratio of the training + set to the entire dataset. + random_state (int, optional): The seed for the + random number generator. + This is used for reproducibility. shuffle (bool, optional): Whether to shuffle the data before splitting. Returns: - Tuple[ClassificationDataset, ClassificationDataset]: A tuple containing the training and testing datasets. + Tuple[ClassificationDataset, ClassificationDataset]: A tuple containing + the training and testing datasets. Example: ```python >>> import supervision as sv >>> cd = sv.ClassificationDataset(...) - >>> train_cd, test_cd = cd.split(split_ratio=0.7, random_state=42, shuffle=True) + >>> train_cd,test_cd = cd.split(split_ratio=0.7, + ... random_state=42,shuffle=True) >>> len(train_cd), len(test_cd) (700, 300) ``` @@ -571,7 +625,8 @@ class ClassificationDataset(BaseDataset): Saves the dataset as a multi-class folder structure. Args: - root_directory_path (str): The path to the directory where the dataset will be saved. + root_directory_path (str): The path to the directory + where the dataset will be saved. """ os.makedirs(root_directory_path, exist_ok=True) diff --git a/supervision/dataset/formats/pascal_voc.py b/supervision/dataset/formats/pascal_voc.py index a3005eb9..8753821c 100644 --- a/supervision/dataset/formats/pascal_voc.py +++ b/supervision/dataset/formats/pascal_voc.py @@ -52,13 +52,19 @@ def detections_to_pascal_voc( Converts Detections object to Pascal VOC XML format. Args: - detections (Detections): A Detections object containing bounding boxes, class ids, and other relevant information. - classes (List[str]): A list of class names corresponding to the class ids in the Detections object. + detections (Detections): A Detections object containing bounding boxes, + class ids, and other relevant information. + classes (List[str]): A list of class names corresponding to the + class ids in the Detections object. filename (str): The name of the image file associated with the detections. - image_shape (Tuple[int, int, int]): The shape of the image file associated with the detections. - min_image_area_percentage (float): Minimum detection area relative to area of image associated with it. - max_image_area_percentage (float): Maximum detection area relative to area of image associated with it. - approximation_percentage (float): The percentage of polygon points to be removed from the input polygon, in the range [0, 1). + image_shape (Tuple[int, int, int]): The shape of the image + file associated with the detections. + min_image_area_percentage (float): Minimum detection area + relative to area of image associated with it. + max_image_area_percentage (float): Maximum detection area + relative to area of image associated with it. + approximation_percentage (float): The percentage of + polygon points to be removed from the input polygon, in the range [0, 1). Returns: str: An XML string in Pascal VOC format representing the detections. """ @@ -123,13 +129,16 @@ def load_pascal_voc_annotations( annotation_path: str, ) -> Tuple[str, Detections, List[str]]: """ - Loads PASCAL VOC XML annotations and returns the image name, a Detections instance, and a list of class names. + Loads PASCAL VOC XML annotations and returns the image name, + a Detections instance, and a list of class names. Args: annotation_path (str): The path to the PASCAL VOC XML annotations file. Returns: - Tuple[str, Detections, List[str]]: A tuple containing the image name, a Detections instance, and a list of class names of objects in the detections. + Tuple[str, Detections, List[str]]: A tuple containing the image name, + a Detections instance, and a list of class + names of objects in the detections. """ tree = parse(annotation_path) root = tree.getroot() diff --git a/supervision/dataset/formats/yolo.py b/supervision/dataset/formats/yolo.py index 2ccf3a30..dee8c4ee 100644 --- a/supervision/dataset/formats/yolo.py +++ b/supervision/dataset/formats/yolo.py @@ -112,16 +112,23 @@ def load_yolo_annotations( force_masks: bool = False, ) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]: """ - Loads YOLO annotations and returns class names, images, and their corresponding detections. + Loads YOLO annotations and returns class names, images, + and their corresponding detections. Args: images_directory_path (str): The path to the directory containing the images. - annotations_directory_path (str): The path to the directory containing the YOLO annotation files. - data_yaml_path (str): The path to the data YAML file containing class information. - force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present. + annotations_directory_path (str): The path to the directory + containing the YOLO annotation files. + data_yaml_path (str): The path to the data + YAML file containing class information. + force_masks (bool, optional): If True, forces masks to be loaded + for all annotations, regardless of whether they are present. Returns: - Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]: A tuple containing a list of class names, a dictionary with image names as keys and images as values, and a dictionary with image names as keys and corresponding Detections instances as values. + Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]: + A tuple containing a list of class names, a dictionary with + image names as keys and images as values, and a dictionary + with image names as keys and corresponding Detections instances as values. """ image_paths = list_files_with_extensions( directory=images_directory_path, extensions=["jpg", "jpeg", "png"] diff --git a/supervision/detection/annotate.py b/supervision/detection/annotate.py index 3b3e40da..88d065f9 100644 --- a/supervision/detection/annotate.py +++ b/supervision/detection/annotate.py @@ -12,12 +12,15 @@ class BoxAnnotator: A class for drawing bounding boxes on an image using detections provided. Attributes: - color (Union[Color, ColorPalette]): The color to draw the bounding box, can be a single color or a color palette + color (Union[Color, ColorPalette]): The color to draw the bounding box, + can be a single color or a color palette thickness (int): The thickness of the bounding box lines, default is 2 text_color (Color): The color of the text on the bounding box, default is white text_scale (float): The scale of the text on the bounding box, default is 0.5 - text_thickness (int): The thickness of the text on the bounding box, default is 1 - text_padding (int): The padding around the text on the bounding box, default is 5 + text_thickness (int): The thickness of the text on the bounding box, + default is 1 + text_padding (int): The padding around the text on the bounding box, + default is 5 """ @@ -49,8 +52,11 @@ class BoxAnnotator: Args: scene (np.ndarray): The image on which the bounding boxes will be drawn - detections (Detections): The detections for which the bounding boxes will be drawn - labels (Optional[List[str]]): An optional list of labels corresponding to each detection. If `labels` are not provided, corresponding `class_id` will be used as label. + detections (Detections): The detections for which the + bounding boxes will be drawn + labels (Optional[List[str]]): An optional list of labels + corresponding to each detection. If `labels` are not provided, + corresponding `class_id` will be used as label. skip_label (bool): Is set to `True`, skips bounding box label annotation. Returns: np.ndarray: The image with the bounding boxes drawn on it @@ -145,7 +151,8 @@ class MaskAnnotator: A class for overlaying masks on an image using detections provided. Attributes: - color (Union[Color, ColorPalette]): The color to fill the mask, can be a single color or a color palette + color (Union[Color, ColorPalette]): The color to fill the mask, + can be a single color or a color palette """ def __init__( @@ -158,11 +165,13 @@ class MaskAnnotator: self, scene: np.ndarray, detections: Detections, opacity: float = 0.5 ) -> np.ndarray: """ - Overlays the masks on the given image based on the provided detections, with a specified opacity. + Overlays the masks on the given image based on the provided detections, + with a specified opacity. Args: scene (np.ndarray): The image on which the masks will be overlaid - detections (Detections): The detections for which the masks will be overlaid + detections (Detections): The detections for which the + masks will be overlaid opacity (float): The opacity of the masks, between 0 and 1, default is 0.5 Returns: diff --git a/supervision/detection/core.py b/supervision/detection/core.py index a67908ab..69692e02 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -57,11 +57,16 @@ class Detections: """ Data class containing information about the detections in a video frame. Attributes: - xyxy (np.ndarray): An array of shape `(n, 4)` containing the bounding boxes coordinates in format `[x1, y1, x2, y2]` - mask: (Optional[np.ndarray]): An array of shape `(n, W, H)` containing the segmentation masks. - confidence (Optional[np.ndarray]): An array of shape `(n,)` containing the confidence scores of the detections. - class_id (Optional[np.ndarray]): An array of shape `(n,)` containing the class ids of the detections. - tracker_id (Optional[np.ndarray]): An array of shape `(n,)` containing the tracker ids of the detections. + xyxy (np.ndarray): An array of shape `(n, 4)` containing the + bounding boxes coordinates in format `[x1, y1, x2, y2]` + mask: (Optional[np.ndarray]): An array of shape `(n, W, H)` + containing the segmentation masks. + confidence (Optional[np.ndarray]): An array of shape `(n,)` + containing the confidence scores of the detections. + class_id (Optional[np.ndarray]): An array of shape `(n,)` + containing the class ids of the detections. + tracker_id (Optional[np.ndarray]): An array of shape `(n,)` + containing the tracker ids of the detections. """ xyxy: np.ndarray @@ -96,7 +101,8 @@ class Detections: ] ]: """ - Iterates over the Detections object and yield a tuple of `(xyxy, mask, confidence, class_id, tracker_id)` for each detection. + Iterates over the Detections object and yield a tuple of + `(xyxy, mask, confidence, class_id, tracker_id)` for each detection. """ for i in range(len(self.xyxy)): yield ( @@ -141,10 +147,12 @@ class Detections: @classmethod def from_yolov5(cls, yolov5_results) -> Detections: """ - Creates a Detections instance from a [YOLOv5](https://github.com/ultralytics/yolov5) inference result. + Creates a Detections instance from a + [YOLOv5](https://github.com/ultralytics/yolov5) inference result. Args: - yolov5_results (yolov5.models.common.Detections): The output Detections instance from YOLOv5 + yolov5_results (yolov5.models.common.Detections): + The output Detections instance from YOLOv5 Returns: Detections: A new Detections object. @@ -171,10 +179,12 @@ class Detections: @classmethod def from_yolov8(cls, yolov8_results) -> Detections: """ - Creates a Detections instance from a [YOLOv8](https://github.com/ultralytics/ultralytics) inference result. + Creates a Detections instance from a + [YOLOv8](https://github.com/ultralytics/ultralytics) inference result. Args: - yolov8_results (ultralytics.yolo.engine.results.Results): The output Results instance from YOLOv8 + yolov8_results (ultralytics.yolo.engine.results.Results): + The output Results instance from YOLOv8 Returns: Detections: A new Detections object. @@ -201,10 +211,15 @@ class Detections: @classmethod def from_yolo_nas(cls, yolo_nas_results) -> Detections: """ - Creates a Detections instance from a [YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md) inference result. + Creates a Detections instance from a + [YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md) + inference result. Args: - yolo_nas_results (super_gradients.training.models.prediction_results.ImageDetectionPrediction): The output Results instance from YOLO-NAS + yolo_nas_results (ImageDetectionPrediction): + The output Results instance from YOLO-NAS + ImageDetectionPrediction is coming from + 'super_gradients.training.models.prediction_results' Returns: Detections: A new Detections object. @@ -230,11 +245,13 @@ class Detections: @classmethod def from_mmdetection(cls, mmdet_results) -> Detections: """ - Creates a Detections instance from a [mmdetection](https://github.com/open-mmlab/mmdetection) inference result. + Creates a Detections instance from + a [mmdetection](https://github.com/open-mmlab/mmdetection) inference result. Also supported for [mmyolo](https://github.com/open-mmlab/mmyolo) Args: - mmdet_results (mmdet.structures.DetDataSample): The output Results instance from MMDetection + mmdet_results (mmdet.structures.DetDataSample): + The output Results instance from MMDetection Returns: Detections: A new Detections object. @@ -246,7 +263,8 @@ class Detections: >>> from mmdet.apis import DetInferencer >>> inferencer = DetInferencer(model_name, checkpoint, device) - >>> mmdet_result = inferencer(SOURCE_IMAGE_PATH, out_dir='./output', return_datasample=True)["predictions"][0] + >>> mmdet_result = inferencer(SOURCE_IMAGE_PATH, out_dir='./output', + ... return_datasample=True)["predictions"][0] >>> detections = sv.Detections.from_mmdet(mmdet_result) ``` """ @@ -259,7 +277,8 @@ class Detections: @classmethod def from_transformers(cls, transformers_results: dict) -> Detections: """ - Creates a Detections instance from object detection [transformer](https://github.com/huggingface/transformers) inference result. + Creates a Detections instance from object detection + [transformer](https://github.com/huggingface/transformers) inference result. Returns: Detections: A new Detections object. @@ -273,13 +292,16 @@ class Detections: @classmethod def from_detectron2(cls, detectron2_results) -> Detections: """ - Create a Detections object from the [Detectron2](https://github.com/facebookresearch/detectron2) inference result. + Create a Detections object from the + [Detectron2](https://github.com/facebookresearch/detectron2) inference result. Args: - detectron2_results: The output of a Detectron2 model containing instances with prediction data. + detectron2_results: The output of a + Detectron2 model containing instances with prediction data. Returns: - (Detections): A Detections object containing the bounding boxes, class IDs, and confidences of the predictions. + (Detections): A Detections object containing the bounding boxes, + class IDs, and confidences of the predictions. Example: ```python @@ -294,7 +316,6 @@ class Detections: >>> cfg.MODEL.WEIGHTS = "path/to/model_weights.pth" >>> predictor = DefaultPredictor(cfg) >>> result = predictor(image) - >>> detections = sv.Detections.from_detectron2(result) ``` """ @@ -310,14 +331,18 @@ class Detections: @classmethod def from_roboflow(cls, roboflow_result: dict, class_list: List[str]) -> Detections: """ - Create a Detections object from the [Roboflow](https://roboflow.com/) API inference result. + Create a Detections object from the [Roboflow](https://roboflow.com/) + API inference result. Args: - roboflow_result (dict): The result from the Roboflow API containing predictions. - class_list (List[str]): A list of class names corresponding to the class IDs in the API result. + roboflow_result (dict): The result from the + Roboflow API containing predictions. + class_list (List[str]): A list of class names + corresponding to the class IDs in the API result. Returns: - (Detections): A Detections object containing the bounding boxes, class IDs, and confidences of the predictions. + (Detections): A Detections object containing the bounding boxes, class IDs, + and confidences of the predictions. Example: ```python @@ -354,7 +379,9 @@ class Detections: @classmethod def from_sam(cls, sam_result: List[dict]) -> Detections: """ - Creates a Detections instance from [Segment Anything Model](https://github.com/facebookresearch/segment-anything) inference result. + Creates a Detections instance from + [Segment Anything Model](https://github.com/facebookresearch/segment-anything) + inference result. Args: sam_result (List[dict]): The output Results instance from SAM @@ -365,14 +392,19 @@ class Detections: Example: ```python >>> import supervision as sv - >>> from segment_anything import sam_model_registry, SamAutomaticMaskGenerator + >>> from segment_anything import ( + ... sam_model_registry, + ... SamAutomaticMaskGenerator + ... ) - >>> sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device=DEVICE) + >>> sam_model_reg = sam_model_registry[MODEL_TYPE] + >>> sam = sam_model_reg(checkpoint=CHECKPOINT_PATH).to(device=DEVICE) >>> mask_generator = SamAutomaticMaskGenerator(sam) >>> sam_result = mask_generator.generate(IMAGE) >>> detections = sv.Detections.from_sam(sam_result=sam_result) ``` """ + sorted_generated_masks = sorted( sam_result, key=lambda x: x["area"], reverse=True ) @@ -385,7 +417,8 @@ class Detections: @classmethod def empty(cls) -> Detections: """ - Create an empty Detections object with no bounding boxes, confidences, or class IDs. + Create an empty Detections object with no bounding boxes, + confidences, or class IDs. Returns: (Detections): An empty Detections object. @@ -408,15 +441,18 @@ class Detections: """ Merge a list of Detections objects into a single Detections object. - This method takes a list of Detections objects and combines their respective fields (`xyxy`, `mask`, - `confidence`, `class_id`, and `tracker_id`) into a single Detections object. If all elements in a field are not - `None`, the corresponding field will be stacked. Otherwise, the field will be set to `None`. + This method takes a list of Detections objects and combines their + respective fields (`xyxy`, `mask`, `confidence`, `class_id`, and `tracker_id`) + into a single Detections object. If all elements in a field are not + `None`, the corresponding field will be stacked. + Otherwise, the field will be set to `None`. Args: detections_list (List[Detections]): A list of Detections objects to merge. Returns: - (Detections): A single Detections object containing the merged data from the input list. + (Detections): A single Detections object containing + the merged data from the input list. Example: ```python @@ -436,14 +472,14 @@ class Detections: list(field) for field in zip(*detections_tuples_list) ] - def all_not_none(l): - return all(x is not None for x in l) + def __all_not_none(item_list: List[Any]): + return all(x is not None for x in item_list) xyxy = np.vstack(xyxy) - mask = np.vstack(mask) if all_not_none(mask) else None - confidence = np.hstack(confidence) if all_not_none(confidence) else None - class_id = np.hstack(class_id) if all_not_none(class_id) else None - tracker_id = np.hstack(tracker_id) if all_not_none(tracker_id) else None + mask = np.vstack(mask) if __all_not_none(mask) else None + confidence = np.hstack(confidence) if __all_not_none(confidence) else None + class_id = np.hstack(class_id) if __all_not_none(class_id) else None + tracker_id = np.hstack(tracker_id) if __all_not_none(tracker_id) else None return cls( xyxy=xyxy, @@ -458,10 +494,12 @@ class Detections: Returns the bounding box coordinates for a specific anchor. Args: - anchor (Position): Position of bounding box anchor for which to return the coordinates. + anchor (Position): Position of bounding box anchor + for which to return the coordinates. Returns: - np.ndarray: An array of shape `(n, 2)` containing the bounding box anchor coordinates in format `[x, y]`. + np.ndarray: An array of shape `(n, 2)` containing the bounding + box anchor coordinates in format `[x, y]`. """ if anchor == Position.CENTER: return np.array( @@ -484,7 +522,8 @@ class Detections: Get a subset of the Detections object. Args: - index (Union[int, slice, List[int], np.ndarray]): The index or indices of the subset of the Detections + index (Union[int, slice, List[int], np.ndarray]): + The index or indices of the subset of the Detections Returns: (Detections): A subset of the Detections object. @@ -519,11 +558,14 @@ class Detections: @property def area(self) -> np.ndarray: """ - Calculate the area of each detection in the set of object detections. If masks field is defined property - returns are of each mask. If only box is given property return area of each box. + Calculate the area of each detection in the set of object detections. + If masks field is defined property returns are of each mask. + If only box is given property return area of each box. Returns: - np.ndarray: An array of floats containing the area of each detection in the format of `(area_1, area_2, ..., area_n)`, where n is the number of detections. + np.ndarray: An array of floats containing the area of each detection + in the format of `(area_1, area_2, ..., area_n)`, + where n is the number of detections. """ if self.mask is not None: return np.array([np.sum(mask) for mask in self.mask]) @@ -536,7 +578,9 @@ class Detections: Calculate the area of each bounding box in the set of object detections. Returns: - np.ndarray: An array of floats containing the area of each bounding box in the format of `(area_1, area_2, ..., area_n)`, where n is the number of detections. + np.ndarray: An array of floats containing the area of each bounding + box in the format of `(area_1, area_2, ..., area_n)`, + where n is the number of detections. """ return (self.xyxy[:, 3] - self.xyxy[:, 1]) * (self.xyxy[:, 2] - self.xyxy[:, 0]) @@ -547,14 +591,19 @@ class Detections: Perform non-maximum suppression on the current set of object detections. Args: - threshold (float, optional): The intersection-over-union threshold to use for non-maximum suppression. Defaults to 0.5. - class_agnostic (bool, optional): Whether to perform class-agnostic non-maximum suppression. If True, the class_id of each detection will be ignored. Defaults to False. + threshold (float, optional): The intersection-over-union threshold + to use for non-maximum suppression. Defaults to 0.5. + class_agnostic (bool, optional): Whether to perform class-agnostic + non-maximum suppression. If True, the class_id of each detection + will be ignored. Defaults to False. Returns: - Detections: A new Detections object containing the subset of detections after non-maximum suppression. + Detections: A new Detections object containing the subset of detections + after non-maximum suppression. Raises: - AssertionError: If `confidence` is None and class_agnostic is False. If `class_id` is None and class_agnostic is False. + AssertionError: If `confidence` is None and class_agnostic is False. + If `class_id` is None and class_agnostic is False. """ if len(self) == 0: return self diff --git a/supervision/detection/line_counter.py b/supervision/detection/line_counter.py index 45a9f7ee..ea97fcb8 100644 --- a/supervision/detection/line_counter.py +++ b/supervision/detection/line_counter.py @@ -113,7 +113,8 @@ class LineZoneAnnotator: Attributes: frame (np.ndarray): The image on which the line will be drawn. - line_counter (LineCounter): The line counter that will be used to draw the line. + line_counter (LineCounter): The line counter + that will be used to draw the line. Returns: np.ndarray: The image with the line drawn on it. diff --git a/supervision/detection/tools/polygon_zone.py b/supervision/detection/tools/polygon_zone.py index ecd78c47..c4406d56 100644 --- a/supervision/detection/tools/polygon_zone.py +++ b/supervision/detection/tools/polygon_zone.py @@ -17,9 +17,11 @@ class PolygonZone: A class for defining a polygon-shaped zone within a frame for detecting objects. Attributes: - polygon (np.ndarray): A polygon represented by a numpy array of shape `(N, 2)`, containing the `x`, `y` coordinates of the points. + 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_position (Position): The position within the bounding box that triggers the zone (default: Position.BOTTOM_CENTER) + triggering_position (Position): The position within the bounding + box that triggers the zone (default: Position.BOTTOM_CENTER) current_count (int): The current count of detected objects within the zone mask (np.ndarray): The 2D bool mask for the polygon zone """ @@ -45,10 +47,12 @@ class PolygonZone: Determines if the detections are within the polygon zone. Parameters: - detections (Detections): The detections to be checked against the polygon zone + detections (Detections): The detections + to be checked against the polygon zone Returns: - np.ndarray: A boolean numpy array indicating if each detection is within the polygon zone + np.ndarray: A boolean numpy array indicating + if each detection is within the polygon zone """ clipped_xyxy = clip_boxes( @@ -65,7 +69,8 @@ class PolygonZone: class PolygonZoneAnnotator: """ - A class for annotating a polygon-shaped zone within a frame with a count of detected objects. + A class for annotating a polygon-shaped zone within a + frame with a count of detected objects. Attributes: zone (PolygonZone): The polygon zone to be annotated @@ -75,7 +80,8 @@ class PolygonZoneAnnotator: text_scale (float): The scale of the text on the polygon, default is 0.5 text_thickness (int): The thickness of the text on the polygon, default is 1 text_padding (int): The padding around the text on the polygon, default is 10 - font (int): The font type for the text on the polygon, default is cv2.FONT_HERSHEY_SIMPLEX + font (int): The font type for the text on the polygon, + default is cv2.FONT_HERSHEY_SIMPLEX center (Tuple[int, int]): The center of the polygon for text placement """ @@ -105,7 +111,8 @@ class PolygonZoneAnnotator: Parameters: scene (np.ndarray): The image on which the polygon zone will be annotated - label (Optional[str]): An optional label for the count of detected objects within the polygon zone (default: None) + label (Optional[str]): An optional label for the count of detected objects + within the polygon zone (default: None) Returns: np.ndarray: The image with the polygon zone and count of detected objects diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 13c5a389..5fa5ef0f 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -10,11 +10,13 @@ def polygon_to_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.n """Generate a mask from a polygon. Args: - polygon (np.ndarray): The polygon for which the mask should be generated, given as a list of vertices. + polygon (np.ndarray): The polygon for which the mask should be generated, + given as a list of vertices. resolution_wh (Tuple[int, int]): The width and height of the desired resolution. Returns: - np.ndarray: The generated 2D mask, where the polygon is marked with `1`'s and the rest is filled with `0`'s. + np.ndarray: The generated 2D mask, where the polygon is marked with + `1`'s and the rest is filled with `0`'s. """ width, height = resolution_wh mask = np.zeros((height, width), dtype=np.uint8) @@ -24,15 +26,20 @@ def polygon_to_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.n def box_iou_batch(boxes_true: np.ndarray, boxes_detection: np.ndarray) -> np.ndarray: """ - Compute Intersection over Union (IoU) of two sets of bounding boxes - `boxes_true` and `boxes_detection`. Both sets - of boxes are expected to be in `(x_min, y_min, x_max, y_max)` format. + Compute Intersection over Union (IoU) of two sets of bounding boxes - + `boxes_true` and `boxes_detection`. Both sets + of boxes are expected to be in `(x_min, y_min, x_max, y_max)` format. Args: - boxes_true (np.ndarray): 2D `np.ndarray` representing ground-truth boxes. `shape = (N, 4)` where `N` is number of true objects. - boxes_detection (np.ndarray): 2D `np.ndarray` representing detection boxes. `shape = (M, 4)` where `M` is number of detected objects. + boxes_true (np.ndarray): 2D `np.ndarray` representing ground-truth boxes. + `shape = (N, 4)` where `N` is number of true objects. + boxes_detection (np.ndarray): 2D `np.ndarray` representing detection boxes. + `shape = (M, 4)` where `M` is number of detected objects. Returns: - np.ndarray: Pairwise IoU of boxes from `boxes_true` and `boxes_detection`. `shape = (N, M)` where `N` is number of true objects and `M` is number of detected objects. + np.ndarray: Pairwise IoU of boxes from `boxes_true` and `boxes_detection`. + `shape = (N, M)` where `N` is number of true objects and + `M` is number of detected objects. """ def box_area(box): @@ -55,14 +62,19 @@ def non_max_suppression( Perform Non-Maximum Suppression (NMS) on object detection predictions. Args: - predictions (np.ndarray): An array of object detection predictions in the format of `(x_min, y_min, x_max, y_max, score)` or `(x_min, y_min, x_max, y_max, score, class)`. - iou_threshold (float, optional): The intersection-over-union threshold to use for non-maximum suppression. + predictions (np.ndarray): An array of object detection predictions in + the format of `(x_min, y_min, x_max, y_max, score)` + or `(x_min, y_min, x_max, y_max, score, class)`. + iou_threshold (float, optional): The intersection-over-union threshold + to use for non-maximum suppression. Returns: - np.ndarray: A boolean array indicating which predictions to keep after non-maximum suppression. + np.ndarray: A boolean array indicating which predictions to keep after n + on-maximum suppression. Raises: - AssertionError: If `iou_threshold` is not within the closed range from `0` to `1`. + AssertionError: If `iou_threshold` is not within the + closed range from `0` to `1`. """ assert 0 <= iou_threshold <= 1, ( "Value of `iou_threshold` must be in the closed range from 0 to 1, " @@ -89,7 +101,8 @@ def non_max_suppression( if not keep[index]: continue - # drop detections with iou > iou_threshold and same category as current detections + # drop detections with iou > iou_threshold and + # same category as current detections condition = (iou > iou_threshold) & (categories == category) keep = keep & ~condition @@ -103,14 +116,16 @@ def clip_boxes( Clips bounding boxes coordinates to fit within the frame resolution. Args: - boxes_xyxy (np.ndarray): A numpy array of shape `(N, 4)` where each row corresponds to a bounding box in + boxes_xyxy (np.ndarray): A numpy array of shape `(N, 4)` where each + row corresponds to a bounding box in the format `(x_min, y_min, x_max, y_max)`. - frame_resolution_wh (Tuple[int, int]): A tuple of the form `(width, height)` representing the resolution of the - frame. + frame_resolution_wh (Tuple[int, int]): A tuple of the form `(width, height)` + representing the resolution of the frame. Returns: - np.ndarray: A numpy array of shape `(N, 4)` where each row corresponds to a bounding box with coordinates - clipped to fit within the frame resolution. + np.ndarray: A numpy array of shape `(N, 4)` where each row + corresponds to a bounding box with coordinates clipped to fit + within the frame resolution. """ result = np.copy(boxes_xyxy) width, height = frame_resolution_wh @@ -131,10 +146,12 @@ def mask_to_xyxy(masks: np.ndarray) -> np.ndarray: Converts a 3D `np.array` of 2D bool masks into a 2D `np.array` of bounding boxes. Parameters: - masks (np.ndarray): A 3D `np.array` of shape `(N, W, H)` containing 2D bool masks + masks (np.ndarray): A 3D `np.array` of shape `(N, W, H)` + containing 2D bool masks Returns: - np.ndarray: A 2D `np.array` of shape `(N, 4)` containing the bounding boxes `(x_min, y_min, x_max, y_max)` for each mask + np.ndarray: A 2D `np.array` of shape `(N, 4)` containing the bounding boxes + `(x_min, y_min, x_max, y_max)` for each mask """ n = masks.shape[0] bboxes = np.zeros((n, 4), dtype=int) @@ -155,12 +172,14 @@ def mask_to_polygons(mask: np.ndarray) -> List[np.ndarray]: Converts a binary mask to a list of polygons. Parameters: - mask (np.ndarray): A binary mask represented as a 2D NumPy array of shape `(H, W)`, - where H and W are the height and width of the mask, respectively. + mask (np.ndarray): A binary mask represented as a 2D NumPy array of + shape `(H, W)`, where H and W are the height and width of + the mask, respectively. Returns: - List[np.ndarray]: A list of polygons, where each polygon is represented by a NumPy array of shape `(N, 2)`, - containing the `x`, `y` coordinates of the points. Polygons with fewer points than `MIN_POLYGON_POINT_COUNT = 3` + List[np.ndarray]: A list of polygons, where each polygon is represented by a + NumPy array of shape `(N, 2)`, containing the `x`, `y` coordinates + of the points. Polygons with fewer points than `MIN_POLYGON_POINT_COUNT = 3` are excluded from the output. """ @@ -183,15 +202,21 @@ def filter_polygons_by_area( Filters a list of polygons based on their area. Parameters: - polygons (List[np.ndarray]): A list of polygons, where each polygon is represented by a NumPy array of shape `(N, 2)`, + polygons (List[np.ndarray]): A list of polygons, where each polygon is + represented by a NumPy array of shape `(N, 2)`, containing the `x`, `y` coordinates of the points. - min_area (Optional[float]): The minimum area threshold. Only polygons with an area greater than or equal to this value - will be included in the output. If set to None, no minimum area constraint will be applied. - max_area (Optional[float]): The maximum area threshold. Only polygons with an area less than or equal to this value - will be included in the output. If set to None, no maximum area constraint will be applied. + min_area (Optional[float]): The minimum area threshold. + Only polygons with an area greater than or equal to this value + will be included in the output. If set to None, + no minimum area constraint will be applied. + max_area (Optional[float]): The maximum area threshold. + Only polygons with an area less than or equal to this value + will be included in the output. If set to None, + no maximum area constraint will be applied. Returns: - List[np.ndarray]: A new list of polygons containing only those with areas within the specified thresholds. + List[np.ndarray]: A new list of polygons containing only those with + areas within the specified thresholds. """ if min_area is None and max_area is None: return polygons @@ -213,7 +238,8 @@ def polygon_to_xyxy(polygon: np.ndarray) -> np.ndarray: containing the `x`, `y` coordinates of the points. Returns: - np.ndarray: A 1D NumPy array containing the bounding box `(x_min, y_min, x_max, y_max)` of the input polygon. + np.ndarray: A 1D NumPy array containing the bounding box + `(x_min, y_min, x_max, y_max)` of the input polygon. """ x_min, y_min = np.min(polygon, axis=0) x_max, y_max = np.max(polygon, axis=0) @@ -226,16 +252,23 @@ def approximate_polygon( """ Approximates a given polygon by reducing a certain percentage of points. - This function uses the Ramer-Douglas-Peucker algorithm to simplify the input polygon by reducing the number of points - while preserving the general shape. + This function uses the Ramer-Douglas-Peucker algorithm to simplify the input + polygon by reducing the number of points + while preserving the general shape. Parameters: - polygon (np.ndarray): A 2D NumPy array of shape `(N, 2)` containing the `x`, `y` coordinates of the input polygon's points. - percentage (float): The percentage of points to be removed from the input polygon, in the range `[0, 1)`. - epsilon_step (float): Approximation accuracy step. Epsilon is the maximum distance between the original curve and its approximation. + polygon (np.ndarray): A 2D NumPy array of shape `(N, 2)` containing + the `x`, `y` coordinates of the input polygon's points. + percentage (float): The percentage of points to be removed from the + input polygon, in the range `[0, 1)`. + epsilon_step (float): Approximation accuracy step. + Epsilon is the maximum distance between the original curve + and its approximation. Returns: - np.ndarray: A new 2D NumPy array of shape `(M, 2)`, where `M <= N * (1 - percentage)`, containing the `x`, `y` coordinates of the + np.ndarray: A new 2D NumPy array of shape `(M, 2)`, + where `M <= N * (1 - percentage)`, containing + the `x`, `y` coordinates of the approximated polygon's points. """ diff --git a/supervision/draw/color.py b/supervision/draw/color.py index b597831e..9997c8b8 100644 --- a/supervision/draw/color.py +++ b/supervision/draw/color.py @@ -42,7 +42,8 @@ class Color: """ 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" + :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: @@ -105,7 +106,8 @@ class ColorPalette: """ 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" + :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: diff --git a/supervision/draw/utils.py b/supervision/draw/utils.py index 55a329b5..9d75690d 100644 --- a/supervision/draw/utils.py +++ b/supervision/draw/utils.py @@ -115,15 +115,19 @@ def draw_text( Draw text with background on a scene. Parameters: - scene (np.ndarray): A 2-dimensional numpy ndarray representing an image or scene. + scene (np.ndarray): A 2-dimensional numpy ndarray representing an image or scene text (str): The text to be drawn. - text_anchor (Point): The anchor point for the text, represented as a Point object with x and y attributes. + text_anchor (Point): The anchor point for the text, represented as a + Point object with x and y attributes. text_color (Color, optional): The color of the text. Defaults to black. text_scale (float, optional): The scale of the text. Defaults to 0.5. text_thickness (int, optional): The thickness of the text. Defaults to 1. - text_padding (int, optional): The amount of padding to add around the text when drawing a rectangle in the background. Defaults to 10. - text_font (int, optional): The font to use for the text. Defaults to cv2.FONT_HERSHEY_SIMPLEX. - background_color (Color, optional): The color of the background rectangle, if one is to be drawn. Defaults to None. + text_padding (int, optional): The amount of padding to add around the text + when drawing a rectangle in the background. Defaults to 10. + text_font (int, optional): The font to use for the text. + Defaults to cv2.FONT_HERSHEY_SIMPLEX. + background_color (Color, optional): The color of the background rectangle, + if one is to be drawn. Defaults to None. Returns: np.ndarray: The input scene with the text drawn on it. @@ -132,7 +136,7 @@ def draw_text( ```python >>> scene = np.zeros((100, 100, 3), dtype=np.uint8) >>> text_anchor = Point(x=50, y=50) - >>> scene = draw_text(scene=scene, text="Hello, world!", text_anchor=text_anchor) + >>> scene = draw_text(scene=scene, text="Hello, world!",text_anchor=text_anchor) ``` """ text_width, text_height = cv2.getTextSize( diff --git a/supervision/geometry/utils.py b/supervision/geometry/utils.py index 400e4d75..f1b66db2 100644 --- a/supervision/geometry/utils.py +++ b/supervision/geometry/utils.py @@ -7,13 +7,18 @@ def get_polygon_center(polygon: np.ndarray) -> Point: """ Calculate the center of a polygon. - This function takes in a polygon as a 2-dimensional numpy ndarray and returns the center of the polygon as a Point object. The center is calculated as the mean of the polygon's vertices along each axis, and is rounded down to the nearest integer. + This function takes in a polygon as a 2-dimensional numpy ndarray and + returns the center of the polygon as a Point object. + The center is calculated as the mean of the polygon's vertices along each axis, + and is rounded down to the nearest integer. Parameters: - polygon (np.ndarray): A 2-dimensional numpy ndarray representing the vertices of the polygon. + polygon (np.ndarray): A 2-dimensional numpy ndarray representing the + vertices of the polygon. Returns: - Point: The center of the polygon, represented as a Point object with x and y attributes. + Point: The center of the polygon, represented as a + Point object with x and y attributes. Examples: ```python diff --git a/supervision/metrics/detection.py b/supervision/metrics/detection.py index 9683c5db..7fc345ba 100644 --- a/supervision/metrics/detection.py +++ b/supervision/metrics/detection.py @@ -18,10 +18,14 @@ class ConfusionMatrix: Confusion matrix for object detection tasks. Attributes: - matrix (np.ndarray): An 2D `np.ndarray` of shape `(len(classes) + 1, len(classes) + 1)` containing the number of `TP`, `FP`, `FN` and `TN` for each class. + matrix (np.ndarray): An 2D `np.ndarray` of shape + `(len(classes) + 1, len(classes) + 1)` + containing the number of `TP`, `FP`, `FN` and `TN` for each class. classes (List[str]): Model class names. - conf_threshold (float): Detection confidence threshold between `0` and `1`. Detections with lower confidence will be excluded from the matrix. - iou_threshold (float): Detection IoU threshold between `0` and `1`. Detections with lower IoU will be classified as `FP`. + conf_threshold (float): Detection confidence threshold between `0` and `1`. + Detections with lower confidence will be excluded from the matrix. + iou_threshold (float): Detection IoU threshold between `0` and `1`. + Detections with lower IoU will be classified as `FP`. """ matrix: np.ndarray @@ -45,8 +49,10 @@ class ConfusionMatrix: targets (List[Detections]): Detections objects from ground-truth. predictions (List[Detections]): Detections objects predicted by the model. classes (List[str]): Model class names. - conf_threshold (float): Detection confidence threshold between `0` and `1`. Detections with lower confidence will be excluded. - iou_threshold (float): Detection IoU threshold between `0` and `1`. Detections with lower IoU will be classified as `FP`. + conf_threshold (float): Detection confidence threshold between `0` and `1`. + Detections with lower confidence will be excluded. + iou_threshold (float): Detection IoU threshold between `0` and `1`. + Detections with lower IoU will be classified as `FP`. Returns: ConfusionMatrix: New instance of ConfusionMatrix. @@ -132,11 +138,19 @@ class ConfusionMatrix: Calculate confusion matrix based on predicted and ground-truth detections. Args: - predictions (List[np.ndarray]): Each element of the list describes a single image and has `shape = (M, 6)` where `M` is the number of detected objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class, conf)` format. - targets (List[np.ndarray]): Each element of the list describes a single image and has `shape = (N, 5)` where `N` is the number of ground-truth objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class)` format. + predictions (List[np.ndarray]): Each element of the list describes a single + image and has `shape = (M, 6)` where `M` is the number of detected + objects. Each row is expected to be in + `(x_min, y_min, x_max, y_max, class, conf)` format. + targets (List[np.ndarray]): Each element of the list describes a single + image and has `shape = (N, 5)` where `N` is the number of + ground-truth objects. Each row is expected to be in + `(x_min, y_min, x_max, y_max, class)` format. classes (List[str]): Model class names. - conf_threshold (float): Detection confidence threshold between `0` and `1`. Detections with lower confidence will be excluded. - iou_threshold (float): Detection iou threshold between `0` and `1`. Detections with lower iou will be classified as `FP`. + conf_threshold (float): Detection confidence threshold between `0` and `1`. + Detections with lower confidence will be excluded. + iou_threshold (float): Detection iou threshold between `0` and `1`. + Detections with lower iou will be classified as `FP`. Returns: ConfusionMatrix: New instance of ConfusionMatrix. @@ -246,11 +260,19 @@ class ConfusionMatrix: Calculate confusion matrix for a batch of detections for a single image. Args: - predictions (List[np.ndarray]): Each element of the list describes a single image and has `shape = (M, 6)` where `M` is the number of detected objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class, conf)` format. - targets (List[np.ndarray]): Each element of the list describes a single image and has `shape = (N, 5)` where `N` is the number of ground-truth objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class)` format. + predictions (List[np.ndarray]): Each element of the list describes a single + image and has `shape = (M, 6)` where `M` is the number of detected + objects. Each row is expected to be in + `(x_min, y_min, x_max, y_max, class, conf)` format. + targets (List[np.ndarray]): Each element of the list describes a single + image and has `shape = (N, 5)` where `N` is the number of ground-truth + objects. Each row is expected to be in + `(x_min, y_min, x_max, y_max, class)` format. num_classes (int): Number of classes. - conf_threshold (float): Detection confidence threshold between `0` and `1`. Detections with lower confidence will be excluded. - iou_threshold (float): Detection iou threshold between `0` and `1`. Detections with lower iou will be classified as `FP`. + conf_threshold (float): Detection confidence threshold between `0` and `1`. + Detections with lower confidence will be excluded. + iou_threshold (float): Detection iou threshold between `0` and `1`. + Detections with lower iou will be classified as `FP`. Returns: np.ndarray: Confusion matrix based on a single image. @@ -304,8 +326,8 @@ class ConfusionMatrix: @staticmethod def _drop_extra_matches(matches: np.ndarray) -> np.ndarray: """ - Deduplicate matches. If there are multiple matches for the same true or predicted box, - only the one with the highest IoU is kept. + Deduplicate matches. If there are multiple matches for the same true or + predicted box, only the one with the highest IoU is kept. """ if matches.shape[0] > 0: matches = matches[matches[:, 2].argsort()[::-1]] @@ -327,9 +349,12 @@ class ConfusionMatrix: Args: dataset (DetectionDataset): Object detection dataset used for evaluation. - callback (Callable[[np.ndarray], Detections]): Function that takes an image as input and returns Detections object. - conf_threshold (float): Detection confidence threshold between `0` and `1`. Detections with lower confidence will be excluded. - iou_threshold (float): Detection IoU threshold between `0` and `1`. Detections with lower IoU will be classified as `FP`. + callback (Callable[[np.ndarray], Detections]): Function that takes an image + as input and returns Detections object. + conf_threshold (float): Detection confidence threshold between `0` and `1`. + Detections with lower confidence will be excluded. + iou_threshold (float): Detection IoU threshold between `0` and `1`. + Detections with lower IoU will be classified as `FP`. Returns: ConfusionMatrix: New instance of ConfusionMatrix. @@ -386,9 +411,11 @@ class ConfusionMatrix: Create confusion matrix plot and save it at selected location. Args: - save_path (Optional[str]): Path to save the plot. If not provided, plot will be displayed. + save_path (Optional[str]): Path to save the plot. If not provided, + plot will be displayed. title (Optional[str]): Title of the plot. - classes (Optional[List[str]]): List of classes to be displayed on the plot. If not provided, all classes will be displayed. + classes (Optional[List[str]]): List of classes to be displayed on the plot. + If not provided, all classes will be displayed. normalize (bool): If True, normalize the confusion matrix. fig_size (Tuple[int, int]): Size of the plot. diff --git a/supervision/utils/file.py b/supervision/utils/file.py index aaaaf6fa..3954f9d5 100644 --- a/supervision/utils/file.py +++ b/supervision/utils/file.py @@ -21,11 +21,13 @@ def list_files_with_extensions( directory: Union[str, Path], extensions: Optional[List[str]] = None ) -> List[Path]: """ - List files in a directory with specified extensions or all files if no extensions are provided. + List files in a directory with specified extensions or + all files if no extensions are provided. Args: directory (Union[str, Path]): The directory path as a string or Path object. - extensions (Optional[List[str]]): A list of file extensions to filter. Default is None, which lists all files. + extensions (Optional[List[str]]): A list of file extensions to filter. + Default is None, which lists all files. Returns: (List[Path]): A list of Path objects for the matching files. @@ -38,9 +40,11 @@ def list_files_with_extensions( >>> files = sv.list_files_with_extensions(directory='my_directory') >>> # List only files with '.txt' and '.md' extensions - >>> files = sv.list_files_with_extensions(directory='my_directory', extensions=['txt', 'md']) + >>> files = sv.list_files_with_extensions( + ... directory='my_directory', extensions=['txt', 'md']) ``` """ + directory = Path(directory) files_with_extensions = [] diff --git a/supervision/utils/image.py b/supervision/utils/image.py index a0843854..30d4e6a8 100644 --- a/supervision/utils/image.py +++ b/supervision/utils/image.py @@ -12,7 +12,8 @@ def crop(image: np.ndarray, xyxy: np.ndarray) -> np.ndarray: 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). + xyxy (np.ndarray): A numpy array containing the bounding box coordinates + in the format (x1, y1, x2, y2). Returns: (np.ndarray): The cropped image as a numpy array. @@ -47,18 +48,23 @@ class ImageSink: Args: target_dir_path (str): The target directory where images will be saved. - overwrite (bool, optional): Whether to overwrite the existing directory. Defaults to False. - image_name_pattern (str, optional): The image file name pattern. Defaults to "image_{:05d}.png". + overwrite (bool, optional): Whether to overwrite the existing directory. + Defaults to False. + image_name_pattern (str, optional): The image file name pattern. + Defaults to "image_{:05d}.png". Examples: ```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): + >>> 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): ... sink.save_image(image=image) ``` """ + self.target_dir_path = target_dir_path self.overwrite = overwrite self.image_name_pattern = image_name_pattern @@ -80,7 +86,9 @@ class ImageSink: Args: image (np.ndarray): The image to be saved. - image_name (str, optional): The name to use for the saved image. If not provided, a name will be generated using the `image_name_pattern`. + image_name (str, optional): The name to use for the saved image. + If not provided, a name will be + generated using the `image_name_pattern`. """ if image_name is None: image_name = self.image_name_pattern.format(self.image_count) diff --git a/supervision/utils/notebook.py b/supervision/utils/notebook.py index b084aaea..97b6659c 100644 --- a/supervision/utils/notebook.py +++ b/supervision/utils/notebook.py @@ -50,9 +50,12 @@ def plot_images_grid( Args: images (List[np.ndarray]): A list of images as numpy arrays. - grid_size (Tuple[int, int]): A tuple specifying the number of rows and columns for the grid. - titles (Optional[List[str]]): A list of titles for each image. Defaults to None. - size (Tuple[int, int]): A tuple specifying the width and height of the entire plot in inches. + grid_size (Tuple[int, int]): A tuple specifying the number + of rows and columns for the grid. + titles (Optional[List[str]]): A list of titles for each image. + Defaults to None. + size (Tuple[int, int]): A tuple specifying the width and + height of the entire plot in inches. cmap (str): the colormap to use for single channel images. Raises: diff --git a/supervision/utils/video.py b/supervision/utils/video.py index 2aaefaa0..333890d3 100644 --- a/supervision/utils/video.py +++ b/supervision/utils/video.py @@ -10,13 +10,15 @@ import numpy as np @dataclass class VideoInfo: """ - A class to store video information, including width, height, fps and total number of frames. + A class to store video information, including width, height, fps and + total number of frames. Attributes: width (int): width of the video in pixels height (int): height of the video in pixels fps (int): frames per second of the video - total_frames (int, optional): total number of frames in the video, default is None + total_frames (int, optional): total number of frames in the video, + default is None Examples: ```python @@ -61,7 +63,8 @@ class VideoSink: Attributes: target_path (str): The path to the output file where the video will be saved. - video_info (VideoInfo): Information about the video resolution, fps, and total frame count. + video_info (VideoInfo): Information about the video resolution, fps, + and total frame count. Example: ```python @@ -69,8 +72,10 @@ class VideoSink: >>> video_info = sv.VideoInfo.from_video_path(video_path='source_video.mp4') - >>> with sv.VideoSink(target_path='target_video.mp4', video_info=video_info) as sink: - ... for frame in get_video_frames_generator(source_path='source_video.mp4', stride=2): + >>> with sv.VideoSink(target_path='target_video.mp4', + ... video_info=video_info) as sink: + ... for frame in get_video_frames_generator(source_path='source_video.mp4', + ... stride=2): ... sink.write_frame(frame=frame) ``` """ @@ -118,12 +123,16 @@ def get_video_frames_generator( Args: source_path (str): The path of the video file. - stride (int): Indicates the interval at which frames are returned, skipping stride - 1 frames between each. - start (int): Indicates the starting position from which video should generate frames - end (Optional[int]): Indicates the ending position at which video should stop generating frames. If None, video will be read to the end. + stride (int): Indicates the interval at which frames are returned, + skipping stride - 1 frames between each. + start (int): Indicates the starting position from which + video should generate frames + end (Optional[int]): Indicates the ending position at which video + should stop generating frames. If None, video will be read to the end. Returns: - (Generator[np.ndarray, None, None]): A generator that yields the frames of the video. + (Generator[np.ndarray, None, None]): A generator that yields the + frames of the video. Examples: ```python @@ -154,12 +163,16 @@ def process_video( callback: Callable[[np.ndarray, int], np.ndarray], ) -> None: """ - Process a video file by applying a callback function on each frame and saving the result to a target video file. + Process a video file by applying a callback function on each frame + and saving the result to a target video file. Args: source_path (str): The path to the source video file. target_path (str): The path to the target video file. - callback (Callable[[np.ndarray, int], np.ndarray]): A function that takes in a numpy ndarray representation of a video frame and an int index of the frame and returns a processed numpy ndarray representation of the frame. + callback (Callable[[np.ndarray, int], np.ndarray]): A function that takes in + a numpy ndarray representation of a video frame and an + int index of the frame and returns a processed numpy ndarray + representation of the frame. Examples: ```python diff --git a/test/dataset/formats/test_coco.py b/test/dataset/formats/test_coco.py index 34f1df57..62d1b75b 100644 --- a/test/dataset/formats/test_coco.py +++ b/test/dataset/formats/test_coco.py @@ -47,7 +47,8 @@ def mock_cock_coco_annotation( ], ["fashion-assistant", "baseball cap"], DoesNotRaise(), - ), # two coco categories; one with supercategory == "none" and one with supercategory != "none" + ), # two coco categories; one with supercategory == "none" and + # one with supercategory != "none" ( [ {"id": 0, "name": "fashion-assistant", "supercategory": "none"}, @@ -56,7 +57,8 @@ def mock_cock_coco_annotation( ], ["fashion-assistant", "baseball cap", "hoodie"], DoesNotRaise(), - ), # three coco categories; one with supercategory == "none" and two with supercategory != "none" + ), # three coco categories; one with supercategory == "none" and + # two with supercategory != "none" ( [ {"id": 0, "name": "fashion-assistant", "supercategory": "none"}, @@ -65,7 +67,8 @@ def mock_cock_coco_annotation( ], ["fashion-assistant", "baseball cap", "hoodie"], DoesNotRaise(), - ), # three coco categories; one with supercategory == "none" and two with supercategory != "none" (different order) + ), # three coco categories; one with supercategory == "none" and + # two with supercategory != "none" (different order) ], ) def test_coco_categories_to_classes( diff --git a/test/dataset/formats/test_yolo.py b/test/dataset/formats/test_yolo.py index 17714a06..f84e275b 100644 --- a/test/dataset/formats/test_yolo.py +++ b/test/dataset/formats/test_yolo.py @@ -154,7 +154,8 @@ def test_with_mask( ), ), DoesNotRaise(), - ), # yolo annotation file with two lines - one box and one polygon in with_masks mode + ), # yolo annotation file with two lines - + # one box and one polygon in with_masks mode ( ["0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6", "1 0.11 0.47 0.22 0.30"], (1000, 1000), diff --git a/test/detection/test_utils.py b/test/detection/test_utils.py index 1acd3f56..b096fd08 100644 --- a/test/detection/test_utils.py +++ b/test/detection/test_utils.py @@ -224,7 +224,8 @@ def test_clip_boxes( 100, [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])], DoesNotRaise(), - ), # two polygons with min_area and max_area equal to the area of the first polygon + ), # two polygons with min_area and + # max_area equal to the area of the first polygon ( [ np.array([[0, 0], [0, 10], [10, 10], [10, 0]]), @@ -234,7 +235,8 @@ def test_clip_boxes( 400, [np.array([[0, 0], [0, 20], [20, 20], [20, 0]])], DoesNotRaise(), - ), # two polygons with min_area and max_area equal to the area of the second polygon + ), # two polygons with min_area and + # max_area equal to the area of the second polygon ], ) def test_filter_polygons_by_area(