93 lines
3.9 KiB
Python
93 lines
3.9 KiB
Python
|
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||
|
|
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import cv2
|
||
|
|
|
||
|
|
from ultralytics.solutions.solutions import BaseSolution, SolutionAnnotator, SolutionResults
|
||
|
|
from ultralytics.utils import LOGGER
|
||
|
|
from ultralytics.utils.plotting import colors
|
||
|
|
|
||
|
|
|
||
|
|
class ObjectBlurrer(BaseSolution):
|
||
|
|
"""
|
||
|
|
A class to manage the blurring of detected objects in a real-time video stream.
|
||
|
|
|
||
|
|
This class extends the BaseSolution class and provides functionality for blurring objects based on detected bounding
|
||
|
|
boxes. The blurred areas are updated directly in the input image, allowing for privacy preservation or other effects.
|
||
|
|
|
||
|
|
Attributes:
|
||
|
|
blur_ratio (int): The intensity of the blur effect applied to detected objects (higher values create more blur).
|
||
|
|
iou (float): Intersection over Union threshold for object detection.
|
||
|
|
conf (float): Confidence threshold for object detection.
|
||
|
|
|
||
|
|
Methods:
|
||
|
|
process: Apply a blurring effect to detected objects in the input image.
|
||
|
|
extract_tracks: Extract tracking information from detected objects.
|
||
|
|
display_output: Display the processed output image.
|
||
|
|
|
||
|
|
Examples:
|
||
|
|
>>> blurrer = ObjectBlurrer()
|
||
|
|
>>> frame = cv2.imread("frame.jpg")
|
||
|
|
>>> processed_results = blurrer.process(frame)
|
||
|
|
>>> print(f"Total blurred objects: {processed_results.total_tracks}")
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, **kwargs: Any) -> None:
|
||
|
|
"""
|
||
|
|
Initialize the ObjectBlurrer class for applying a blur effect to objects detected in video streams or images.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
**kwargs (Any): Keyword arguments passed to the parent class and for configuration.
|
||
|
|
blur_ratio (float): Intensity of the blur effect (0.1-1.0, default=0.5).
|
||
|
|
"""
|
||
|
|
super().__init__(**kwargs)
|
||
|
|
blur_ratio = self.CFG["blur_ratio"]
|
||
|
|
if blur_ratio < 0.1:
|
||
|
|
LOGGER.warning("blur ratio cannot be less than 0.1, updating it to default value 0.5")
|
||
|
|
blur_ratio = 0.5
|
||
|
|
self.blur_ratio = int(blur_ratio * 100)
|
||
|
|
|
||
|
|
def process(self, im0) -> SolutionResults:
|
||
|
|
"""
|
||
|
|
Apply a blurring effect to detected objects in the input image.
|
||
|
|
|
||
|
|
This method extracts tracking information, applies blur to regions corresponding to detected objects,
|
||
|
|
and annotates the image with bounding boxes.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
im0 (numpy.ndarray): The input image containing detected objects.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
(SolutionResults): Object containing the processed image and number of tracked objects.
|
||
|
|
- plot_im (numpy.ndarray): The annotated output image with blurred objects.
|
||
|
|
- total_tracks (int): The total number of tracked objects in the frame.
|
||
|
|
|
||
|
|
Examples:
|
||
|
|
>>> blurrer = ObjectBlurrer()
|
||
|
|
>>> frame = cv2.imread("image.jpg")
|
||
|
|
>>> results = blurrer.process(frame)
|
||
|
|
>>> print(f"Blurred {results.total_tracks} objects")
|
||
|
|
"""
|
||
|
|
self.extract_tracks(im0) # Extract tracks
|
||
|
|
annotator = SolutionAnnotator(im0, self.line_width)
|
||
|
|
|
||
|
|
# Iterate over bounding boxes and classes
|
||
|
|
for box, cls, conf in zip(self.boxes, self.clss, self.confs):
|
||
|
|
# Crop and blur the detected object
|
||
|
|
blur_obj = cv2.blur(
|
||
|
|
im0[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])],
|
||
|
|
(self.blur_ratio, self.blur_ratio),
|
||
|
|
)
|
||
|
|
# Update the blurred area in the original image
|
||
|
|
im0[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] = blur_obj
|
||
|
|
annotator.box_label(
|
||
|
|
box, label=self.adjust_box_label(cls, conf), color=colors(cls, True)
|
||
|
|
) # Annotate bounding box
|
||
|
|
|
||
|
|
plot_im = annotator.result()
|
||
|
|
self.display_output(plot_im) # Display the output using the base class function
|
||
|
|
|
||
|
|
# Return a SolutionResults
|
||
|
|
return SolutionResults(plot_im=plot_im, total_tracks=len(self.track_ids))
|