Spaces:
Running
Running
| import logging | |
| import glob | |
| from tqdm import tqdm | |
| import numpy as np | |
| import torch | |
| import cv2 | |
| class FaceDetector(object): | |
| """An abstract class representing a face detector. | |
| Any other face detection implementation must subclass it. All subclasses | |
| must implement ``detect_from_image``, that return a list of detected | |
| bounding boxes. Optionally, for speed considerations detect from path is | |
| recommended. | |
| """ | |
| def __init__(self, device, verbose): | |
| self.device = device | |
| self.verbose = verbose | |
| if verbose: | |
| if 'cpu' in device: | |
| logger = logging.getLogger(__name__) | |
| logger.warning("Detection running on CPU, this may be potentially slow.") | |
| if 'cpu' not in device and 'cuda' not in device: | |
| if verbose: | |
| logger.error("Expected values for device are: {cpu, cuda} but got: %s", device) | |
| raise ValueError | |
| def detect_from_image(self, tensor_or_path): | |
| """Detects faces in a given image. | |
| This function detects the faces present in a provided BGR(usually) | |
| image. The input can be either the image itself or the path to it. | |
| Arguments: | |
| tensor_or_path {numpy.ndarray, torch.tensor or string} -- the path | |
| to an image or the image itself. | |
| Example:: | |
| >>> path_to_image = 'data/image_01.jpg' | |
| ... detected_faces = detect_from_image(path_to_image) | |
| [A list of bounding boxes (x1, y1, x2, y2)] | |
| >>> image = cv2.imread(path_to_image) | |
| ... detected_faces = detect_from_image(image) | |
| [A list of bounding boxes (x1, y1, x2, y2)] | |
| """ | |
| raise NotImplementedError | |
| def detect_from_directory(self, path, extensions=['.jpg', '.png'], recursive=False, show_progress_bar=True): | |
| """Detects faces from all the images present in a given directory. | |
| Arguments: | |
| path {string} -- a string containing a path that points to the folder containing the images | |
| Keyword Arguments: | |
| extensions {list} -- list of string containing the extensions to be | |
| consider in the following format: ``.extension_name`` (default: | |
| {['.jpg', '.png']}) recursive {bool} -- option wherever to scan the | |
| folder recursively (default: {False}) show_progress_bar {bool} -- | |
| display a progressbar (default: {True}) | |
| Example: | |
| >>> directory = 'data' | |
| ... detected_faces = detect_from_directory(directory) | |
| {A dictionary of [lists containing bounding boxes(x1, y1, x2, y2)]} | |
| """ | |
| if self.verbose: | |
| logger = logging.getLogger(__name__) | |
| if len(extensions) == 0: | |
| if self.verbose: | |
| logger.error("Expected at list one extension, but none was received.") | |
| raise ValueError | |
| if self.verbose: | |
| logger.info("Constructing the list of images.") | |
| additional_pattern = '/**/*' if recursive else '/*' | |
| files = [] | |
| for extension in extensions: | |
| files.extend(glob.glob(path + additional_pattern + extension, recursive=recursive)) | |
| if self.verbose: | |
| logger.info("Finished searching for images. %s images found", len(files)) | |
| logger.info("Preparing to run the detection.") | |
| predictions = {} | |
| for image_path in tqdm(files, disable=not show_progress_bar): | |
| if self.verbose: | |
| logger.info("Running the face detector on image: %s", image_path) | |
| predictions[image_path] = self.detect_from_image(image_path) | |
| if self.verbose: | |
| logger.info("The detector was successfully run on all %s images", len(files)) | |
| return predictions | |
| def reference_scale(self): | |
| raise NotImplementedError | |
| def reference_x_shift(self): | |
| raise NotImplementedError | |
| def reference_y_shift(self): | |
| raise NotImplementedError | |
| def tensor_or_path_to_ndarray(tensor_or_path, rgb=True): | |
| """Convert path (represented as a string) or torch.tensor to a numpy.ndarray | |
| Arguments: | |
| tensor_or_path {numpy.ndarray, torch.tensor or string} -- path to the image, or the image itself | |
| """ | |
| if isinstance(tensor_or_path, str): | |
| return cv2.imread(tensor_or_path) if not rgb else cv2.imread(tensor_or_path)[..., ::-1] | |
| elif torch.is_tensor(tensor_or_path): | |
| # Call cpu in case its coming from cuda | |
| return tensor_or_path.cpu().numpy()[..., ::-1].copy() if not rgb else tensor_or_path.cpu().numpy() | |
| elif isinstance(tensor_or_path, np.ndarray): | |
| return tensor_or_path[..., ::-1].copy() if not rgb else tensor_or_path | |
| else: | |
| raise TypeError | |