diff --git a/supervision/annotators/utils.py b/supervision/annotators/utils.py index ed0070f0..6f9cd9db 100644 --- a/supervision/annotators/utils.py +++ b/supervision/annotators/utils.py @@ -125,6 +125,12 @@ class Trace: return self.xy[self.tracker_id == tracker_id] +def pillow_to_cv2(image: Image.Image) -> np.ndarray: + scene = np.array(image) + scene = cv2.cvtColor(scene, cv2.COLOR_RGB2BGR) + return scene + + def scene_to_annotator_img_type(annotate_func): """ Decorates `BaseAnnotator.annotate` implementations, converts scene to @@ -138,8 +144,7 @@ def scene_to_annotator_img_type(annotate_func): return annotate_func(self, scene, *args, **kwargs) if isinstance(scene, Image.Image): - scene = np.array(scene) - scene = cv2.cvtColor(scene, cv2.COLOR_RGB2BGR) + scene = pillow_to_cv2(scene) annotated = annotate_func(self, scene, *args, **kwargs) annotated = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB) annotated = Image.fromarray(annotated) diff --git a/supervision/utils/notebook.py b/supervision/utils/notebook.py index e8a0c7aa..159f3c17 100644 --- a/supervision/utils/notebook.py +++ b/supervision/utils/notebook.py @@ -2,17 +2,21 @@ from typing import List, Optional, Tuple import cv2 import matplotlib.pyplot as plt -import numpy as np +from PIL import Image + +from supervision.annotators.base import ImageType +from supervision.annotators.utils import pillow_to_cv2 def plot_image( - image: np.ndarray, size: Tuple[int, int] = (12, 12), cmap: Optional[str] = "gray" + image: ImageType, size: Tuple[int, int] = (12, 12), cmap: Optional[str] = "gray" ) -> None: """ Plots image using matplotlib. Args: - image (np.ndarray): The frame to be displayed. + image (ImageType): The frame to be displayed ImageType + is a flexible type, accepting either `numpy.ndarray` or `PIL.Image.Image`. size (Tuple[int, int]): The size of the plot. cmap (str): the colormap to use for single channel images. @@ -27,6 +31,9 @@ def plot_image( sv.plot_image(image=image, size=(16, 16)) ``` """ + if isinstance(image, Image.Image): + image = pillow_to_cv2(image) + plt.figure(figsize=size) if image.ndim == 2: @@ -39,7 +46,7 @@ def plot_image( def plot_images_grid( - images: List[np.ndarray], + images: List[ImageType], grid_size: Tuple[int, int], titles: Optional[List[str]] = None, size: Tuple[int, int] = (12, 12), @@ -49,7 +56,8 @@ def plot_images_grid( Plots images in a grid using matplotlib. Args: - images (List[np.ndarray]): A list of images as numpy arrays. + images (List[ImageType]): A list of images as ImageType + is a flexible type, accepting either `numpy.ndarray` or `PIL.Image.Image`. 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. @@ -65,9 +73,10 @@ def plot_images_grid( ```python import cv2 import supervision as sv + from PIL import Image image1 = cv2.imread("path/to/image1.jpg") - image2 = cv2.imread("path/to/image2.jpg") + image2 = Image.open("path/to/image2.jpg") image3 = cv2.imread("path/to/image3.jpg") images = [image1, image2, image3] @@ -79,6 +88,10 @@ def plot_images_grid( """ nrows, ncols = grid_size + for idx, img in enumerate(images): + if isinstance(img, Image.Image): + images[idx] = pillow_to_cv2(img) + if len(images) > nrows * ncols: raise ValueError( "The number of images exceeds the grid size. Please increase the grid size"