# Copyright (c) OpenMMLab. All rights reserved. import os import os.path as osp import sys import time from contextlib import contextmanager from typing import Optional from urllib.parse import urlparse from urllib.request import urlopen import cv2 import numpy as np from torch.hub import HASH_REGEX, download_url_to_file @contextmanager def limit_max_fps(fps: Optional[float]): t_start = time.time() try: yield finally: t_end = time.time() if fps is not None: t_sleep = 1.0 / fps - t_end + t_start if t_sleep > 0: time.sleep(t_sleep) def _is_url(filename): """Check if the file is a url link. Args: filename (str): the file name or url link. Returns: bool: is url or not. """ prefixes = ['http://', 'https://'] for p in prefixes: if filename.startswith(p): return True return False def load_image_from_disk_or_url(filename, readFlag=cv2.IMREAD_COLOR): """Load an image file, from disk or url. Args: filename (str): file name on the disk or url link. readFlag (int): readFlag for imdecode. Returns: np.ndarray: A loaded image """ if _is_url(filename): # download the image, convert it to a NumPy array, and then read # it into OpenCV format resp = urlopen(filename) image = np.asarray(bytearray(resp.read()), dtype='uint8') image = cv2.imdecode(image, readFlag) return image else: image = cv2.imread(filename, readFlag) return image def mkdir_or_exist(dir_name, mode=0o777): if dir_name == '': return dir_name = osp.expanduser(dir_name) os.makedirs(dir_name, mode=mode, exist_ok=True) def get_cached_file_path(url, save_dir=None, progress=True, check_hash=False, file_name=None): r"""Loads the Torch serialized object at the given URL. If downloaded file is a zip file, it will be automatically decompressed If the object is already present in `model_dir`, it's deserialized and returned. The default value of ``model_dir`` is ``/checkpoints`` where ``hub_dir`` is the directory returned by :func:`~torch.hub.get_dir`. Args: url (str): URL of the object to download save_dir (str, optional): directory in which to save the object progress (bool, optional): whether or not to display a progress bar to stderr. Default: True check_hash(bool, optional): If True, the filename part of the URL should follow the naming convention ``filename-.ext`` where ```` is the first eight or more digits of the SHA256 hash of the contents of the file. The hash is used to ensure unique names and to verify the contents of the file. Default: False file_name (str, optional): name for the downloaded file. Filename from ``url`` will be used if not set. Default: None. """ if save_dir is None: save_dir = os.path.join('webcam_resources') mkdir_or_exist(save_dir) parts = urlparse(url) filename = os.path.basename(parts.path) if file_name is not None: filename = file_name cached_file = os.path.join(save_dir, filename) if not os.path.exists(cached_file): sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file)) hash_prefix = None if check_hash: r = HASH_REGEX.search(filename) # r is Optional[Match[str]] hash_prefix = r.group(1) if r else None download_url_to_file(url, cached_file, hash_prefix, progress=progress) return cached_file def screen_matting(img, color_low=None, color_high=None, color=None): """Screen Matting. Args: img (np.ndarray): Image data. color_low (tuple): Lower limit (b, g, r). color_high (tuple): Higher limit (b, g, r). color (str): Support colors include: - 'green' or 'g' - 'blue' or 'b' - 'black' or 'k' - 'white' or 'w' """ if color_high is None or color_low is None: if color is not None: if color.lower() == 'g' or color.lower() == 'green': color_low = (0, 200, 0) color_high = (60, 255, 60) elif color.lower() == 'b' or color.lower() == 'blue': color_low = (230, 0, 0) color_high = (255, 40, 40) elif color.lower() == 'k' or color.lower() == 'black': color_low = (0, 0, 0) color_high = (40, 40, 40) elif color.lower() == 'w' or color.lower() == 'white': color_low = (230, 230, 230) color_high = (255, 255, 255) else: NotImplementedError(f'Not supported color: {color}.') else: ValueError('color or color_high | color_low should be given.') mask = cv2.inRange(img, np.array(color_low), np.array(color_high)) == 0 return mask.astype(np.uint8) def expand_and_clamp(box, im_shape, s=1.25): """Expand the bbox and clip it to fit the image shape. Args: box (list): x1, y1, x2, y2 im_shape (ndarray): image shape (h, w, c) s (float): expand ratio Returns: list: x1, y1, x2, y2 """ x1, y1, x2, y2 = box[:4] w = x2 - x1 h = y2 - y1 deta_w = w * (s - 1) / 2 deta_h = h * (s - 1) / 2 x1, y1, x2, y2 = x1 - deta_w, y1 - deta_h, x2 + deta_w, y2 + deta_h img_h, img_w = im_shape[:2] x1 = min(max(0, int(x1)), img_w - 1) y1 = min(max(0, int(y1)), img_h - 1) x2 = min(max(0, int(x2)), img_w - 1) y2 = min(max(0, int(y2)), img_h - 1) return [x1, y1, x2, y2] def _find_connected_components(mask): """Find connected components and sort with areas. Args: mask (ndarray): instance segmentation result. Returns: ndarray (N, 5): Each item contains (x, y, w, h, area). """ num, labels, stats, centroids = cv2.connectedComponentsWithStats(mask) stats = stats[stats[:, 4].argsort()] return stats def _find_bbox(mask): """Find the bounding box for the mask. Args: mask (ndarray): Mask. Returns: list(4, ): Returned box (x1, y1, x2, y2). """ mask_shape = mask.shape if len(mask_shape) == 3: assert mask_shape[-1] == 1, 'the channel of the mask should be 1.' elif len(mask_shape) == 2: pass else: NotImplementedError() h, w = mask_shape[:2] mask_w = mask.sum(0) mask_h = mask.sum(1) left = 0 right = w - 1 up = 0 down = h - 1 for i in range(w): if mask_w[i] > 0: break left += 1 for i in range(w - 1, left, -1): if mask_w[i] > 0: break right -= 1 for i in range(h): if mask_h[i] > 0: break up += 1 for i in range(h - 1, up, -1): if mask_h[i] > 0: break down -= 1 return [left, up, right, down] def copy_and_paste(img, background_img, mask, bbox=None, effect_region=(0.2, 0.2, 0.8, 0.8), min_size=(20, 20)): """Copy the image region and paste to the background. Args: img (np.ndarray): Image data. background_img (np.ndarray): Background image data. mask (ndarray): instance segmentation result. bbox (ndarray): instance bbox, (x1, y1, x2, y2). effect_region (tuple(4, )): The region to apply mask, the coordinates are normalized (x1, y1, x2, y2). """ background_img = background_img.copy() background_h, background_w = background_img.shape[:2] region_h = (effect_region[3] - effect_region[1]) * background_h region_w = (effect_region[2] - effect_region[0]) * background_w region_aspect_ratio = region_w / region_h if bbox is None: bbox = _find_bbox(mask) instance_w = bbox[2] - bbox[0] instance_h = bbox[3] - bbox[1] if instance_w > min_size[0] and instance_h > min_size[1]: aspect_ratio = instance_w / instance_h if region_aspect_ratio > aspect_ratio: resize_rate = region_h / instance_h else: resize_rate = region_w / instance_w mask_inst = mask[int(bbox[1]):int(bbox[3]), int(bbox[0]):int(bbox[2])] img_inst = img[int(bbox[1]):int(bbox[3]), int(bbox[0]):int(bbox[2])] img_inst = cv2.resize(img_inst, (int( resize_rate * instance_w), int(resize_rate * instance_h))) mask_inst = cv2.resize( mask_inst, (int(resize_rate * instance_w), int(resize_rate * instance_h)), interpolation=cv2.INTER_NEAREST) mask_ids = list(np.where(mask_inst == 1)) mask_ids[1] += int(effect_region[0] * background_w) mask_ids[0] += int(effect_region[1] * background_h) background_img[tuple(mask_ids)] = img_inst[np.where(mask_inst == 1)] return background_img def is_image_file(path): if isinstance(path, str): if path.lower().endswith(('.png', '.jpg', '.jpeg', '.tiff', '.bmp')): return True return False class ImageCapture: """A mock-up version of cv2.VideoCapture that always return a const image. Args: image (str | ndarray): The image or image path """ def __init__(self, image): if isinstance(image, str): self.image = load_image_from_disk_or_url(image) else: self.image = image def isOpened(self): return (self.image is not None) def read(self): return True, self.image.copy() def release(self): pass def get(self, propId): if propId == cv2.CAP_PROP_FRAME_WIDTH: return self.image.shape[1] elif propId == cv2.CAP_PROP_FRAME_HEIGHT: return self.image.shape[0] elif propId == cv2.CAP_PROP_FPS: return np.nan else: raise NotImplementedError()