| import os |
| import re |
| import cv2 |
| import yaml |
| import torch |
| import hashlib |
| import argparse |
|
|
| import albumentations as A |
| from albumentations.core.transforms_interface import ImageOnlyTransform |
|
|
| import numpy as np |
|
|
| from PIL import Image |
| from threading import Thread |
| from easydict import EasyDict |
|
|
| VID_EXTS = ('mp4', 'avi', 'h264', 'mkv', 'mov', 'flv', 'wmv', 'webm', 'ts', 'm4v', 'vob', '3gp', '3g2', 'rm', 'rmvb', 'ogv', 'ogg', 'drc', 'gif', 'gifv', 'mng', 'avi', 'mov', 'qt', 'wmv', 'yuv', 'rm', 'rmvb', 'asf', 'amv', 'mp4', 'm4p', 'm4v', 'mpg', 'mp2', 'mpeg', 'mpe', 'mpv', 'mpg', 'mpeg', 'm2v', 'm4v', 'svi', '3gp', '3g2', 'mxf', 'roq', 'nsv', 'flv', 'f4v', 'f4p', 'f4a', 'f4b') |
| IMG_EXTS = ('jpg', 'jpeg', 'bmp', 'png', 'ppm', 'pgm', 'pbm', 'pnm', 'webp', 'sr', 'ras', 'tiff', 'tif', 'exr', 'hdr', 'pic', 'dib', 'jpe', 'jp2', 'j2k', 'jpf', 'jpx', 'jpm', 'mj2', 'jxr', 'hdp', 'wdp', 'cur', 'ico', 'ani', 'icns', 'bpg', 'jp2', 'j2k', 'jpf', 'jpx', 'jpm', 'mj2', 'jxr', 'hdp', 'wdp', 'cur', 'ico', 'ani', 'icns', 'bpg') |
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--source', '-s', type=str, help="Path to the source. Single image, video, directory of images, directory of videos is supported.") |
| parser.add_argument('--dest', '-d', type=str, default=None, help="Path to destination. Results will be stored in current directory if not specified.") |
| parser.add_argument('--type', '-t', type=str, default='rgba', help="Specify output type. If not specified, output results will make the background transparent. Please refer to the documentation for other types.") |
| parser.add_argument('--reverse', '-R', action='store_true', help="Output will be reverse and foreground will be removed instead of background if specified.") |
| parser.add_argument('--format', '-f', type=str, default=None, help="Specify output format. If not specified, it will be saved with the format of input.") |
| parser.add_argument('--resize', '-r', type=str, default='static', help="Specify resizing method. If not specified, static resize will be used. Choose from (static|dynamic).") |
| parser.add_argument('--jit', '-j', action='store_true', help="Speed up inference speed by using torchscript, but decreases output quality.") |
| parser.add_argument('--device', '-D', type=str, default=None, help="Designate device. If not specified, it will find available device.") |
| parser.add_argument('--mode', '-m', type=str, default='base', help="choose between base and fast mode. Also, use base-nightly for nightly release checkpoint.") |
| parser.add_argument('--ckpt', '-c', type=str, default=None, help="Designate checkpoint. If not specified, it will download or load pre-downloaded default checkpoint.") |
| parser.add_argument('--threshold', '-th', type=str, default=None, help="Designate threshold. If specified, it will output hard prediction above threshold. If not specified, it will output soft prediction.") |
| return parser.parse_args() |
|
|
| def get_backend(): |
| if torch.cuda.is_available(): |
| return "cuda:0" |
| elif torch.backends.mps.is_available(): |
| return "mps:0" |
| else: |
| return "cpu" |
| |
| def load_config(config_dir, easy=True): |
| cfg = yaml.load(open(config_dir), yaml.FullLoader) |
| if easy is True: |
| cfg = EasyDict(cfg) |
| return cfg |
|
|
| def get_format(source): |
| img_count = len([i for i in source if i.lower().endswith(IMG_EXTS)]) |
| vid_count = len([i for i in source if i.lower().endswith(VID_EXTS)]) |
| |
| if img_count * vid_count != 0: |
| return '' |
| elif img_count != 0: |
| return 'Image' |
| elif vid_count != 0: |
| return 'Video' |
| else: |
| return '' |
|
|
| def sort(x): |
| convert = lambda text: int(text) if text.isdigit() else text.lower() |
| alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)] |
| return sorted(x, key=alphanum_key) |
|
|
| def download_and_unzip(filename, url, dest, unzip=True, **kwargs): |
| if not os.path.isdir(dest): |
| os.makedirs(dest, exist_ok=True) |
| |
| if os.path.isfile(os.path.join(dest, filename)) is False: |
| os.system("wget -O {} {}".format(os.path.join(dest, filename), url)) |
| elif 'md5' in kwargs.keys() and kwargs['md5'] != hashlib.md5(open(os.path.join(dest, filename), 'rb').read()).hexdigest(): |
| os.system("wget -O {} {}".format(os.path.join(dest, filename), url)) |
| |
| if unzip: |
| os.system("unzip -o {} -d {}".format(os.path.join(dest, filename), dest)) |
| os.system("rm {}".format(os.path.join(dest, filename))) |
| |
| class dynamic_resize: |
| def __init__(self, L=1280): |
| self.L = L |
| |
| def __call__(self, img): |
| size = list(img.size) |
| if (size[0] >= size[1]) and size[1] > self.L: |
| size[0] = size[0] / (size[1] / self.L) |
| size[1] = self.L |
| elif (size[1] > size[0]) and size[0] > self.L: |
| size[1] = size[1] / (size[0] / self.L) |
| size[0] = self.L |
| size = (int(round(size[0] / 32)) * 32, int(round(size[1] / 32)) * 32) |
| |
| return img.resize(size, Image.BILINEAR) |
|
|
| class dynamic_resize_a(ImageOnlyTransform): |
| def __init__(self, L=1280, always_apply=False, p=1.0): |
| super(dynamic_resize_a, self).__init__(always_apply, p) |
| self.L = L |
|
|
| def apply(self, img, **params): |
| size = list(img.shape[:2]) |
| if (size[0] >= size[1]) and size[1] > self.L: |
| size[0] = size[0] / (size[1] / self.L) |
| size[1] = self.L |
| elif (size[1] > size[0]) and size[0] > self.L: |
| size[1] = size[1] / (size[0] / self.L) |
| size[0] = self.L |
| size = (int(round(size[0] / 32)) * 32, int(round(size[1] / 32)) * 32) |
|
|
| return A.resize(img, height=size[0], width=size[1]) |
|
|
| def get_transform_init_args_names(self): |
| return ("L",) |
| |
| class static_resize: |
| def __init__(self, size=[1024, 1024]): |
| self.size = size |
| |
| def __call__(self, img): |
| return img.resize(self.size, Image.BILINEAR) |
|
|
| class normalize: |
| def __init__(self, mean=None, std=None, div=255): |
| self.mean = mean if mean is not None else 0.0 |
| self.std = std if std is not None else 1.0 |
| self.div = div |
| |
| def __call__(self, img): |
| img /= self.div |
| img -= self.mean |
| img /= self.std |
| |
| return img |
| |
| class tonumpy: |
| def __init__(self): |
| pass |
|
|
| def __call__(self, img): |
| img = np.array(img, dtype=np.float32) |
| return img |
| |
| class totensor: |
| def __init__(self): |
| pass |
|
|
| def __call__(self, img): |
| img = img.transpose((2, 0, 1)) |
| img = torch.from_numpy(img).float() |
| |
| return img |
|
|
| class ImageLoader: |
| def __init__(self, root): |
| if os.path.isdir(root): |
| self.images = [os.path.join(root, f) for f in os.listdir(root) if f.lower().endswith(('.jpg', '.png', '.jpeg'))] |
| self.images = sort(self.images) |
| elif os.path.isfile(root): |
| self.images = [root] |
| self.size = len(self.images) |
|
|
| def __iter__(self): |
| self.index = 0 |
| return self |
|
|
| def __next__(self): |
| if self.index == self.size: |
| raise StopIteration |
| |
| img = Image.open(self.images[self.index]).convert('RGB') |
| name = os.path.split(self.images[self.index])[-1] |
| |
| |
| self.index += 1 |
| return img, name |
|
|
| def __len__(self): |
| return self.size |
| |
| class VideoLoader: |
| def __init__(self, root): |
| if os.path.isdir(root): |
| self.videos = [os.path.join(root, f) for f in os.listdir(root) if f.lower().endswith(('.mp4', '.avi', 'mov'))] |
| elif os.path.isfile(root): |
| self.videos = [root] |
| self.size = len(self.videos) |
|
|
| def __iter__(self): |
| self.index = 0 |
| self.cap = None |
| self.fps = None |
| return self |
|
|
| def __next__(self): |
| if self.index == self.size: |
| raise StopIteration |
| |
| if self.cap is None: |
| self.cap = cv2.VideoCapture(self.videos[self.index]) |
| self.fps = self.cap.get(cv2.CAP_PROP_FPS) |
| ret, frame = self.cap.read() |
| name = os.path.split(self.videos[self.index])[-1] |
| |
| if ret is False: |
| self.cap.release() |
| self.cap = None |
| img = None |
| self.index += 1 |
| |
| else: |
| frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
| img = Image.fromarray(frame).convert('RGB') |
| |
| return img, name |
| |
| def __len__(self): |
| return self.size |
| |
| class WebcamLoader: |
| def __init__(self, ID): |
| self.ID = int(ID) |
| self.cap = cv2.VideoCapture(self.ID) |
| self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640) |
| self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480) |
| self.imgs = [] |
| self.imgs.append(self.cap.read()[1]) |
| self.thread = Thread(target=self.update, daemon=True) |
| self.thread.start() |
| |
| def update(self): |
| while self.cap.isOpened(): |
| ret, frame = self.cap.read() |
| if ret is True: |
| self.imgs.append(frame) |
| else: |
| break |
| |
| def __iter__(self): |
| return self |
|
|
| def __next__(self): |
| if len(self.imgs) > 0: |
| frame = self.imgs[-1] |
| else: |
| frame = Image.fromarray(np.zeros((480, 640, 3)).astype(np.uint8)) |
| |
| if self.thread.is_alive() is False or cv2.waitKey(1) == ord('q'): |
| cv2.destroyAllWindows() |
| raise StopIteration |
| |
| else: |
| frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
| frame = Image.fromarray(frame).convert('RGB') |
| |
| del self.imgs[:-1] |
| return frame, None |
|
|
| def __len__(self): |
| return 0 |
|
|
