diff --git a/supervision/draw/utils.py b/supervision/draw/utils.py index 9d75690d..eb61f0ac 100644 --- a/supervision/draw/utils.py +++ b/supervision/draw/utils.py @@ -1,4 +1,5 @@ -from typing import Optional +from typing import Optional, Union +import os import cv2 import numpy as np @@ -168,3 +169,62 @@ def draw_text( lineType=cv2.LINE_AA, ) return scene + +def draw_image( + scene: np.ndarray, image: Union[str, np.ndarray], opacity: float, rect: Rect +) -> np.ndarray: + """ + Draws an image onto a given scene with specified opacity and dimensions. + + Args: + scene (np.ndarray): The background image onto which the image will be drawn. + image (Union[str, np.ndarray]): The image to be drawn. Can be either a file path or a NumPy array. + opacity (float): The opacity level of the image to be drawn, ranging from 0.0 to 1.0. + rect (Rect): A Rect object specifying the dimensions and position where the image will be drawn. + + Returns: + np.ndarray: The scene with the image drawn onto it. + + Example: + >>> scene = np.zeros((400, 400, 3), dtype=np.uint8) + >>> image_path = "path/to/image.jpg" + >>> opacity = 0.5 + >>> rect = Rect(x=50, y=50, width=200, height=200) + >>> new_scene = draw_image(scene, image_path, opacity, rect) + """ + if isinstance(image, str): + assert os.path.exists(image), f'The specified path ("{image}") does not exist.' + image = cv2.imread(image, cv2.IMREAD_UNCHANGED) + + assert 0.0 <= opacity <= 1.0, "The opacity has to be between 0.0 and 1.0." + + assert ( + rect.x >= 0 and rect.y >= 0 + ), "The top left coordinates of the rectangle have to be positive." + assert ( + rect.x + rect.width <= scene.shape[1] and rect.y + rect.height <= scene.shape[0] + ), "The image you are trying to draw exceeds the bounds of the scene." + + image = cv2.resize(image, (rect.width, rect.height)) + + # watermark with transparent background + if image.shape[2] == 4: + b, g, r, a = cv2.split(image) + b = cv2.bitwise_and(b, b, mask=a) + g = cv2.bitwise_and(g, g, mask=a) + r = cv2.bitwise_and(r, r, mask=a) + image = cv2.merge([b, g, r, a]) + del b, g, r, a # immediately free up memory + + if scene.shape[2] == 3: + scene = np.dstack([scene, np.ones(scene.shape[:2], dtype=np.uint8) * 255]) + + scene_h, scene_w, channels = scene.shape[:3] + water_h, water_w = image.shape[:2] + + overlay = np.zeros((scene_h, scene_w, channels), dtype="uint8") + overlay[rect.y : rect.y + water_h, rect.x : rect.x + water_w] = image + + cv2.addWeighted(overlay, opacity, scene, 1.0, 0, scene) + + return scene[:, :, :3] \ No newline at end of file