diff --git a/TODO.md b/TODO.md
index c754f50ed..436f42de8 100644
--- a/TODO.md
+++ b/TODO.md
@@ -22,8 +22,6 @@
- Refactor: Move `nunchaku` models to refernce instead of internal decision, owner @CalamitousFelicitousness
- Refactor: [GGUF](https://huggingface.co/docs/diffusers/main/en/quantization/gguf)
- Refactor: move sampler options to settings to config
-- Refactor: remove `CodeFormer`, owner @CalamitousFelicitousness
-- Refactor: remove `GFPGAN`, owner @CalamitousFelicitousness
- Reimplement `llama` remover for Kanvas, pending end-to-end review of `Kanvas`
## Modular
diff --git a/cli/generate.json b/cli/generate.json
index d2bfddb37..4af6bad9b 100644
--- a/cli/generate.json
+++ b/cli/generate.json
@@ -24,10 +24,7 @@
{
"upscaler_1": "SwinIR_4x",
"upscaler_2": "None",
- "upscaling_resize": 0,
- "gfpgan_visibility": 0,
- "codeformer_visibility": 0,
- "codeformer_weight": 0.5
+ "upscaling_resize": 0
},
"options":
{
diff --git a/cli/generate.py b/cli/generate.py
index 1e533adc3..f804739ca 100755
--- a/cli/generate.py
+++ b/cli/generate.py
@@ -65,8 +65,6 @@ def exif(info, i = None, op = 'generate'):
seed = ', '.join([str(x) for x in seed]) # int list to str list to single str
template = '{prompt} | negative {negative_prompt} | seed {s} | steps {steps} | cfgscale {cfg_scale} | sampler {sampler_name} | batch {batch_size} | timestamp {job_timestamp} | model {model} | vae {vae}'.format(s = seed, model = sd.options['sd_model_checkpoint'], vae = sd.options['sd_vae'], **info) # pylint: disable=consider-using-f-string
if op == 'upscale':
- template += ' | faces gfpgan' if sd.upscale.gfpgan_visibility > 0 else ''
- template += ' | faces codeformer' if sd.upscale.codeformer_visibility > 0 else ''
template += ' | upscale {resize}x {upscaler}'.format(resize = sd.upscale.upscaling_resize, upscaler = sd.upscale.upscaler_1) if sd.upscale.upscaler_1 != 'None' else '' # pylint: disable=consider-using-f-string
template += ' | upscale {resize}x {upscaler}'.format(resize = sd.upscale.upscaling_resize, upscaler = sd.upscale.upscaler_2) if sd.upscale.upscaler_2 != 'None' else '' # pylint: disable=consider-using-f-string
if op == 'grid':
@@ -309,7 +307,6 @@ def args(): # parse cmd arguments
sd.generate.height = params.height if params.height > 0 else sd.generate.height
sd.generate.steps = params.steps if params.steps > 0 else sd.generate.steps
sd.upscale.upscaling_resize = params.upscale if params.upscale > 0 else sd.upscale.upscaling_resize
- sd.upscale.codeformer_visibility = 1 if params.detailer else sd.upscale.codeformer_visibility
sd.options.sd_vae = params.vae if params.vae != '' else sd.options.sd_vae
sd.options.sd_model_checkpoint = params.model if params.model != '' else sd.options.sd_model_checkpoint
sd.upscale.upscaler_1 = 'SwinIR_4x' if params.upscale > 1 else sd.upscale.upscaler_1
diff --git a/cli/process.py b/cli/process.py
index 61890f442..591b7c694 100644
--- a/cli/process.py
+++ b/cli/process.py
@@ -126,11 +126,9 @@ def reset():
all_images = []
-def upscale_restore_image(res: Result, upscale: bool = False, restore: bool = False):
+def upscale_restore_image(res: Result, upscale: bool = False):
kwargs = util.Map({
'image': encode(res.image),
- 'codeformer_visibility': 0.0,
- 'codeformer_weight': 0.0,
})
if res.image.width >= options.process.target_size and res.image.height >= options.process.target_size:
upscale = False
@@ -138,14 +136,10 @@ def upscale_restore_image(res: Result, upscale: bool = False, restore: bool = Fa
kwargs.upscaler_1 = 'SwinIR_4x'
kwargs.upscaling_resize = 2
res.ops.append('upscale')
- if restore:
- kwargs.codeformer_visibility = 1.0
- kwargs.codeformer_weight = 0.2
- res.ops.append('restore')
- if upscale or restore:
+ if upscale:
result = sdapi.postsync('/sdapi/v1/extra-single-image', kwargs)
if 'image' not in result:
- res.message = 'failed to upscale/restore image'
+ res.message = 'failed to upscale image'
else:
res.image = Image.open(io.BytesIO(base64.b64decode(result['image'])))
return res
@@ -309,7 +303,7 @@ def file(filename: str, folder: str, tag = None, requested = []):
if res.image is None:
return res
# post processing steps
- res = upscale_restore_image(res, 'upscale' in requested, 'restore' in requested)
+ res = upscale_restore_image(res, 'upscale' in requested)
if res.image.width < options.process.target_size or res.image.height < options.process.target_size:
res.message = f'low resolution: [{res.image.width}, {res.image.height}]'
res.image = None
diff --git a/html/licenses.html b/html/licenses.html
index dc0e1fdbe..5c282685e 100644
--- a/html/licenses.html
+++ b/html/licenses.html
@@ -4,46 +4,6 @@
#licenses pre { margin: 1em 0 2em 0;}
-
-Parts of CodeFormer code had to be copied to be compatible with GFPGAN.
-
-S-Lab License 1.0
-
-Copyright 2022 S-Lab
-
-Redistribution and use for non-commercial purpose in source and
-binary forms, with or without modification, are permitted provided
-that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright
- notice, this list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright
- notice, this list of conditions and the following disclaimer in
- the documentation and/or other materials provided with the
- distribution.
-
-3. Neither the name of the copyright holder nor the names of its
- contributors may be used to endorse or promote products derived
- from this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
-"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
-LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
-A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
-HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
-SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
-LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
-DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
-THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
-(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
-In the event that redistribution and/or use for commercial purpose in
-source or binary forms, with or without modification is required,
-please contact the contributor(s) of the work.
-
-
Code for architecture and reading models copied.
diff --git a/html/locale_en.json b/html/locale_en.json
index f5ad4b18c..7c2da4777 100644
--- a/html/locale_en.json
+++ b/html/locale_en.json
@@ -390,7 +390,7 @@
{"id":"","label":"Image Paths","localized":"","reload":"","hint":"Settings related to image filenames, and output directories"},
{"id":"","label":"Live Previews","localized":"","reload":"","hint":"Settings related to live previews, audio notification"},
{"id":"","label":"Sampler Settings","localized":"","reload":"","hint":"Settings related to sampler selection and configuration, and diffuser specific sampler configuration"},
- {"id":"","label":"Postprocessing","localized":"","reload":"","hint":"Settings related to post image generation processing, face restoration, and upscaling"},
+ {"id":"","label":"Postprocessing","localized":"","reload":"","hint":"Settings related to post image generation processing and upscaling"},
{"id":"","label":"Control Options","localized":"","reload":"","hint":"Settings related the Control tab"},
{"id":"","label":"Huggingface","localized":"","reload":"","hint":"Settings related huggingface access"},
{"id":"","label":"Show all pages","localized":"","reload":"","hint":"Show all settings pages"},
@@ -436,9 +436,6 @@
{"id":"","label":"Sigma negative guidance minimum","localized":"","reload":"","hint":"Skip negative prompt for some steps when the image is almost ready, 0=disable"},
{"id":"","label":"Upscaler tile size","localized":"","reload":"","hint":"0 = no tiling"},
{"id":"","label":"Upscaler tile overlap","localized":"","reload":"","hint":"Low values = visible seam"},
- {"id":"","label":"GFPGAN","localized":"","reload":"","hint":"Restore low quality faces using GFPGAN neural network"},
- {"id":"","label":"CodeFormer","localized":"","reload":"","hint":"Restore low quality faces using Codeformer neural network"},
- {"id":"","label":"CodeFormer weight parameter","localized":"","reload":"","hint":"0 = maximum effect; 1 = minimum effect"},
{"id":"","label":"ToMe token merging ratio","localized":"","reload":"","hint":"Enable redundant token merging via tomesd for speed and memory improvements, 0=disabled"},
{"id":"","label":"Todo token merging ratio","localized":"","reload":"","hint":"Enable redundant token merging via todo for speed and memory improvements, 0=disabled"},
{"id":"","label":"Model pipeline","localized":"","reload":"","hint":"If autodetect does not detect model automatically, select model type before loading a model"},
@@ -795,10 +792,8 @@
{"id":"","label":"folder with bsrgan models","localized":"","reload":"","hint":"folder with bsrgan models"},
{"id":"","label":"folder with chainner models","localized":"","reload":"","hint":"folder with chainner models"},
{"id":"","label":"folder with clip models","localized":"","reload":"","hint":"folder with clip models"},
- {"id":"","label":"folder with codeformer models","localized":"","reload":"","hint":"folder with codeformer models"},
{"id":"","label":"folder with control models","localized":"","reload":"","hint":"folder with control models"},
{"id":"","label":"folder with esrgan models","localized":"","reload":"","hint":"folder with esrgan models"},
- {"id":"","label":"folder with gfpgan models","localized":"","reload":"","hint":"folder with gfpgan models"},
{"id":"","label":"folder with huggingface models","localized":"","reload":"","hint":"folder with huggingface models"},
{"id":"","label":"folder with hypernetwork models","localized":"","reload":"","hint":"folder with hypernetwork models"},
{"id":"","label":"folder with ldsr models","localized":"","reload":"","hint":"folder with ldsr models"},
@@ -1124,8 +1119,6 @@
{"id":"","label":"resize method","localized":"","reload":"","hint":"resize method"},
{"id":"","label":"resize scale","localized":"","reload":"","hint":"resize scale"},
{"id":"","label":"restart step","localized":"","reload":"","hint":"restart step"},
- {"id":"","label":"restore faces: codeformer","localized":"","reload":"","hint":"restore faces: codeformer"},
- {"id":"","label":"restore faces: gfpgan","localized":"","reload":"","hint":"restore faces: gfpgan"},
{"id":"","label":"restore pipe on end","localized":"","reload":"","hint":"restore pipe on end"},
{"id":"","label":"restore unparsed prompt","localized":"","reload":"","hint":"restore unparsed prompt"},
{"id":"","label":"reswapper model","localized":"","reload":"","hint":"reswapper model"},
diff --git a/installer.py b/installer.py
index 0e38b1daf..84bce28b6 100644
--- a/installer.py
+++ b/installer.py
@@ -1379,7 +1379,6 @@ def install_optional():
t_start = time.time()
log.info('Installing optional requirements...')
install('--no-build-isolation git+https://github.com/Disty0/BasicSR@23c1fb6f5c559ef5ce7ad657f2fa56e41b121754', 'basicsr', ignore=True, quiet=True)
- install('--no-build-isolation git+https://github.com/Disty0/GFPGAN@ae0f7e44fafe0ef4716f3c10067f8f379b74c21c', 'gfpgan', ignore=True, quiet=True)
install('av', ignore=True, quiet=True)
install('beautifulsoup4', ignore=True, quiet=True)
install('clean-fid', ignore=True, quiet=True)
diff --git a/launch.py b/launch.py
index e0fbc61ed..bfe5e9e33 100755
--- a/launch.py
+++ b/launch.py
@@ -166,7 +166,7 @@ def clean_server():
modules_to_remove = ['webui', 'modules', 'scripts', 'gradio',
'onnx', 'torch', 'pytorch', 'lightning', 'tensor', 'diffusers', 'transformers', 'tokenize', 'safetensors', 'gguf', 'accelerate', 'peft', 'triton', 'huggingface',
'PIL', 'cv2', 'timm', 'numpy', 'scipy', 'sympy', 'sklearn', 'skimage', 'sqlalchemy', 'flash_attn', 'bitsandbytes', 'xformers', 'matplotlib', 'optimum', 'pandas', 'pi', 'git', 're', 'altair',
- 'framepack', 'nudenet', 'agent_scheduler', 'basicsr', 'gfpgan', 'war',
+ 'framepack', 'nudenet', 'agent_scheduler', 'basicsr', 'war',
'fastapi', 'urllib', 'uvicorn', 'web', 'http', 'google', 'starlette', 'socket']
removed_removed = []
for module_loaded in modules_loaded:
diff --git a/modules/api/api.py b/modules/api/api.py
index 9fbbcadc5..9ba641b23 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -80,7 +80,6 @@ class Api:
self.add_api_route("/sdapi/v1/upscalers", endpoints.get_upscalers, methods=["GET"], response_model=List[models.ItemUpscaler])
self.add_api_route("/sdapi/v1/sd-models", endpoints.get_sd_models, methods=["GET"], response_model=List[models.ItemModel])
self.add_api_route("/sdapi/v1/controlnets", endpoints.get_controlnets, methods=["GET"], response_model=List[str])
- self.add_api_route("/sdapi/v1/face-restorers", endpoints.get_restorers, methods=["GET"], response_model=List[models.ItemDetailer])
self.add_api_route("/sdapi/v1/detailers", endpoints.get_detailers, methods=["GET"], response_model=List[models.ItemDetailer])
self.add_api_route("/sdapi/v1/prompt-styles", endpoints.get_prompt_styles, methods=["GET"], response_model=List[models.ItemStyle])
self.add_api_route("/sdapi/v1/embeddings", endpoints.get_embeddings, methods=["GET"], response_model=models.ResEmbeddings)
diff --git a/modules/api/endpoints.py b/modules/api/endpoints.py
index e31fdbb24..0ad141b48 100644
--- a/modules/api/endpoints.py
+++ b/modules/api/endpoints.py
@@ -47,9 +47,6 @@ def get_controlnets(model_type: Optional[str] = None):
from modules.control.units.controlnet import api_list_models
return api_list_models(model_type)
-def get_restorers():
- return [{"name":x.name(), "path": getattr(x, "cmd_dir", None)} for x in shared.face_restorers]
-
def get_detailers():
shared.yolo.enumerate()
return [{"name": k, "path": v} for k, v in shared.yolo.list.items()]
diff --git a/modules/api/models.py b/modules/api/models.py
index 3f9063ee6..915ad15ec 100644
--- a/modules/api/models.py
+++ b/modules/api/models.py
@@ -268,9 +268,6 @@ class FileData(BaseModel):
class ReqProcess(BaseModel):
resize_mode: float = Field(default=0, title="Resize Mode", description="Sets the resize mode: 0 to upscale by upscaling_resize amount, 1 to upscale up to upscaling_resize_h x upscaling_resize_w.")
show_extras_results: bool = Field(default=True, title="Show results", description="Should the backend return the generated image?")
- gfpgan_visibility: float = Field(default=0, title="GFPGAN Visibility", ge=0, le=1, allow_inf_nan=False, description="Sets the visibility of GFPGAN, values should be between 0 and 1.")
- codeformer_visibility: float = Field(default=0, title="CodeFormer Visibility", ge=0, le=1, allow_inf_nan=False, description="Sets the visibility of CodeFormer, values should be between 0 and 1.")
- codeformer_weight: float = Field(default=0, title="CodeFormer Weight", ge=0, le=1, allow_inf_nan=False, description="Sets the weight of CodeFormer, values should be between 0 and 1.")
upscaling_resize: float = Field(default=2, title="Upscaling Factor", ge=1, le=8, description="By how much to upscale the image, only used when resize_mode=0.")
upscaling_resize_w: int = Field(default=512, title="Target Width", ge=1, description="Target width for the upscaler to hit. Only used when resize_mode=1.")
upscaling_resize_h: int = Field(default=512, title="Target Height", ge=1, description="Target height for the upscaler to hit. Only used when resize_mode=1.")
diff --git a/modules/cmd_args.py b/modules/cmd_args.py
index 2d0195f91..522742147 100644
--- a/modules/cmd_args.py
+++ b/modules/cmd_args.py
@@ -74,8 +74,6 @@ def settings_args(opts, args):
group_compat.add_argument("--vae-dir", type=str, help=argparse.SUPPRESS, default=opts.vae_dir)
group_compat.add_argument("--embeddings-dir", type=str, help=argparse.SUPPRESS, default=opts.embeddings_dir)
group_compat.add_argument("--embeddings-templates-dir", type=str, help=argparse.SUPPRESS, default=opts.embeddings_templates_dir)
- group_compat.add_argument("--codeformer-models-path", type=str, help=argparse.SUPPRESS, default=opts.codeformer_models_path)
- group_compat.add_argument("--gfpgan-models-path", type=str, help=argparse.SUPPRESS, default=opts.gfpgan_models_path)
group_compat.add_argument("--esrgan-models-path", type=str, help=argparse.SUPPRESS, default=opts.esrgan_models_path)
group_compat.add_argument("--bsrgan-models-path", type=str, help=argparse.SUPPRESS, default=opts.bsrgan_models_path)
group_compat.add_argument("--realesrgan-models-path", type=str, help=argparse.SUPPRESS, default=opts.realesrgan_models_path)
diff --git a/modules/face_restoration.py b/modules/face_restoration.py
deleted file mode 100644
index d17191fdf..000000000
--- a/modules/face_restoration.py
+++ /dev/null
@@ -1,17 +0,0 @@
-from modules import shared
-
-
-class FaceRestoration:
- def name(self):
- return "None"
-
- def restore(self, np_image):
- return np_image
-
-
-def restore_faces(np_image, p=None):
- face_restorers = [x for x in shared.face_restorers if x.name() == shared.opts.face_restoration_model or shared.opts.face_restoration_model is None]
- if len(face_restorers) == 0:
- return np_image
- face_restorer = face_restorers[0]
- return face_restorer.restore(np_image, p)
diff --git a/modules/facelib/__init__.py b/modules/facelib/__init__.py
deleted file mode 100644
index e69de29bb..000000000
diff --git a/modules/facelib/detection/__init__.py b/modules/facelib/detection/__init__.py
deleted file mode 100644
index 14ea22a6f..000000000
--- a/modules/facelib/detection/__init__.py
+++ /dev/null
@@ -1,73 +0,0 @@
-import os
-from copy import deepcopy
-import torch
-from torch import nn
-from ..utils import load_file_from_url
-from ..utils import download_pretrained_models
-from ..detection.yolov5face.models.common import Conv
-from .retinaface.retinaface import RetinaFace
-from .yolov5face.face_detector import YoloDetector
-from modules import paths
-
-
-model_dir = os.path.join(paths.models_path, 'Codeformer')
-
-
-def init_detection_model(model_name, half=False, device='cuda'):
- if 'retinaface' in model_name:
- model = init_retinaface_model(model_name, half, device)
- elif 'YOLOv5' in model_name:
- model = init_yolov5face_model(model_name, device)
- else:
- raise NotImplementedError(f'{model_name} is not implemented.')
-
- return model
-
-
-def init_retinaface_model(model_name, half=False, device='cuda'):
- if model_name == 'retinaface_resnet50':
- model = RetinaFace(network_name='resnet50', half=half)
- model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth'
- elif model_name == 'retinaface_mobile0.25':
- model = RetinaFace(network_name='mobile0.25', half=half)
- model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_mobilenet0.25_Final.pth'
- else:
- raise NotImplementedError(f'{model_name} is not implemented.')
-
- model_path = load_file_from_url(url=model_url, model_dir=model_dir, progress=True, file_name=None)
- load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
- # remove unnecessary 'module.'
- for k, v in deepcopy(load_net).items():
- if k.startswith('module.'):
- load_net[k[7:]] = v
- load_net.pop(k)
- model.load_state_dict(load_net, strict=True)
- model.eval()
- model = model.to(device)
-
- return model
-
-
-def init_yolov5face_model(model_name, device='cuda'):
- if model_name == 'YOLOv5l':
- model = YoloDetector(config_name='facelib/detection/yolov5face/models/yolov5l.yaml', device=device)
- model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5l-face.pth'
- elif model_name == 'YOLOv5n':
- model = YoloDetector(config_name='facelib/detection/yolov5face/models/yolov5n.yaml', device=device)
- model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5n-face.pth'
- else:
- raise NotImplementedError(f'{model_name} is not implemented.')
-
- model_path = load_file_from_url(url=model_url, model_dir=model_dir, progress=True, file_name=None)
- load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
- model.detector.load_state_dict(load_net, strict=True)
- model.detector.eval()
- model.detector = model.detector.to(device).float()
-
- for m in model.detector.modules():
- if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
- m.inplace = True # pytorch 1.7.0 compatibility
- elif isinstance(m, Conv):
- m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
-
- return model
diff --git a/modules/facelib/detection/align_trans.py b/modules/facelib/detection/align_trans.py
deleted file mode 100644
index 07f1eb365..000000000
--- a/modules/facelib/detection/align_trans.py
+++ /dev/null
@@ -1,219 +0,0 @@
-import cv2
-import numpy as np
-
-from .matlab_cp2tform import get_similarity_transform_for_cv2
-
-# reference facial points, a list of coordinates (x,y)
-REFERENCE_FACIAL_POINTS = [[30.29459953, 51.69630051], [65.53179932, 51.50139999], [48.02519989, 71.73660278],
- [33.54930115, 92.3655014], [62.72990036, 92.20410156]]
-
-DEFAULT_CROP_SIZE = (96, 112)
-
-
-class FaceWarpException(Exception):
-
- def __str__(self):
- return 'In File {}:{}'.format(__file__, super.__str__(self))
-
-
-def get_reference_facial_points(output_size=None, inner_padding_factor=0.0, outer_padding=(0, 0), default_square=False):
- """
- Function:
- ----------
- get reference 5 key points according to crop settings:
- 0. Set default crop_size:
- if default_square:
- crop_size = (112, 112)
- else:
- crop_size = (96, 112)
- 1. Pad the crop_size by inner_padding_factor in each side;
- 2. Resize crop_size into (output_size - outer_padding*2),
- pad into output_size with outer_padding;
- 3. Output reference_5point;
- Parameters:
- ----------
- @output_size: (w, h) or None
- size of aligned face image
- @inner_padding_factor: (w_factor, h_factor)
- padding factor for inner (w, h)
- @outer_padding: (w_pad, h_pad)
- each row is a pair of coordinates (x, y)
- @default_square: True or False
- if True:
- default crop_size = (112, 112)
- else:
- default crop_size = (96, 112);
- !!! make sure, if output_size is not None:
- (output_size - outer_padding)
- = some_scale * (default crop_size * (1.0 +
- inner_padding_factor))
- Returns:
- ----------
- @reference_5point: 5x2 np.array
- each row is a pair of transformed coordinates (x, y)
- """
-
- tmp_5pts = np.array(REFERENCE_FACIAL_POINTS)
- tmp_crop_size = np.array(DEFAULT_CROP_SIZE)
-
- # 0) make the inner region a square
- if default_square:
- size_diff = max(tmp_crop_size) - tmp_crop_size
- tmp_5pts += size_diff / 2
- tmp_crop_size += size_diff
-
- if (output_size and output_size[0] == tmp_crop_size[0] and output_size[1] == tmp_crop_size[1]):
-
- return tmp_5pts
-
- if (inner_padding_factor == 0 and outer_padding == (0, 0)):
- if output_size is None:
- return tmp_5pts
- else:
- raise FaceWarpException('No paddings to do, output_size must be None or {}'.format(tmp_crop_size))
-
- # check output size
- if not (0 <= inner_padding_factor <= 1.0):
- raise FaceWarpException('Not (0 <= inner_padding_factor <= 1.0)')
-
- if ((inner_padding_factor > 0 or outer_padding[0] > 0 or outer_padding[1] > 0) and output_size is None):
- output_size = tmp_crop_size * \
- (1 + inner_padding_factor * 2).astype(np.int32)
- output_size += np.array(outer_padding)
- if not (outer_padding[0] < output_size[0] and outer_padding[1] < output_size[1]):
- raise FaceWarpException('Not (outer_padding[0] < output_size[0] and outer_padding[1] < output_size[1])')
-
- # 1) pad the inner region according inner_padding_factor
- if inner_padding_factor > 0:
- size_diff = tmp_crop_size * inner_padding_factor * 2
- tmp_5pts += size_diff / 2
- tmp_crop_size += np.round(size_diff).astype(np.int32)
-
- # 2) resize the padded inner region
- size_bf_outer_pad = np.array(output_size) - np.array(outer_padding) * 2
-
- if size_bf_outer_pad[0] * tmp_crop_size[1] != size_bf_outer_pad[1] * tmp_crop_size[0]:
- raise FaceWarpException('Must have (output_size - outer_padding)'
- '= some_scale * (crop_size * (1.0 + inner_padding_factor)')
-
- scale_factor = size_bf_outer_pad[0].astype(np.float32) / tmp_crop_size[0]
- tmp_5pts = tmp_5pts * scale_factor
- # size_diff = tmp_crop_size * (scale_factor - min(scale_factor))
- # tmp_5pts = tmp_5pts + size_diff / 2
- tmp_crop_size = size_bf_outer_pad
-
- # 3) add outer_padding to make output_size
- reference_5point = tmp_5pts + np.array(outer_padding)
- tmp_crop_size = output_size
-
- return reference_5point
-
-
-def get_affine_transform_matrix(src_pts, dst_pts):
- """
- Function:
- ----------
- get affine transform matrix 'tfm' from src_pts to dst_pts
- Parameters:
- ----------
- @src_pts: Kx2 np.array
- source points matrix, each row is a pair of coordinates (x, y)
- @dst_pts: Kx2 np.array
- destination points matrix, each row is a pair of coordinates (x, y)
- Returns:
- ----------
- @tfm: 2x3 np.array
- transform matrix from src_pts to dst_pts
- """
-
- tfm = np.float32([[1, 0, 0], [0, 1, 0]])
- n_pts = src_pts.shape[0]
- ones = np.ones((n_pts, 1), src_pts.dtype)
- src_pts_ = np.hstack([src_pts, ones])
- dst_pts_ = np.hstack([dst_pts, ones])
-
- A, res, rank, s = np.linalg.lstsq(src_pts_, dst_pts_)
-
- if rank == 3:
- tfm = np.float32([[A[0, 0], A[1, 0], A[2, 0]], [A[0, 1], A[1, 1], A[2, 1]]])
- elif rank == 2:
- tfm = np.float32([[A[0, 0], A[1, 0], 0], [A[0, 1], A[1, 1], 0]])
-
- return tfm
-
-
-def warp_and_crop_face(src_img, facial_pts, reference_pts=None, crop_size=(96, 112), align_type='smilarity'):
- """
- Function:
- ----------
- apply affine transform 'trans' to uv
- Parameters:
- ----------
- @src_img: 3x3 np.array
- input image
- @facial_pts: could be
- 1)a list of K coordinates (x,y)
- or
- 2) Kx2 or 2xK np.array
- each row or col is a pair of coordinates (x, y)
- @reference_pts: could be
- 1) a list of K coordinates (x,y)
- or
- 2) Kx2 or 2xK np.array
- each row or col is a pair of coordinates (x, y)
- or
- 3) None
- if None, use default reference facial points
- @crop_size: (w, h)
- output face image size
- @align_type: transform type, could be one of
- 1) 'similarity': use similarity transform
- 2) 'cv2_affine': use the first 3 points to do affine transform,
- by calling cv2.getAffineTransform()
- 3) 'affine': use all points to do affine transform
- Returns:
- ----------
- @face_img: output face image with size (w, h) = @crop_size
- """
-
- if reference_pts is None:
- if crop_size[0] == 96 and crop_size[1] == 112:
- reference_pts = REFERENCE_FACIAL_POINTS
- else:
- default_square = False
- inner_padding_factor = 0
- outer_padding = (0, 0)
- output_size = crop_size
-
- reference_pts = get_reference_facial_points(output_size, inner_padding_factor, outer_padding,
- default_square)
-
- ref_pts = np.float32(reference_pts)
- ref_pts_shp = ref_pts.shape
- if max(ref_pts_shp) < 3 or min(ref_pts_shp) != 2:
- raise FaceWarpException('reference_pts.shape must be (K,2) or (2,K) and K>2')
-
- if ref_pts_shp[0] == 2:
- ref_pts = ref_pts.T
-
- src_pts = np.float32(facial_pts)
- src_pts_shp = src_pts.shape
- if max(src_pts_shp) < 3 or min(src_pts_shp) != 2:
- raise FaceWarpException('facial_pts.shape must be (K,2) or (2,K) and K>2')
-
- if src_pts_shp[0] == 2:
- src_pts = src_pts.T
-
- if src_pts.shape != ref_pts.shape:
- raise FaceWarpException('facial_pts and reference_pts must have the same shape')
-
- if align_type == 'cv2_affine':
- tfm = cv2.getAffineTransform(src_pts[0:3], ref_pts[0:3])
- elif align_type == 'affine':
- tfm = get_affine_transform_matrix(src_pts, ref_pts)
- else:
- tfm = get_similarity_transform_for_cv2(src_pts, ref_pts)
-
- face_img = cv2.warpAffine(src_img, tfm, (crop_size[0], crop_size[1]))
-
- return face_img
diff --git a/modules/facelib/detection/matlab_cp2tform.py b/modules/facelib/detection/matlab_cp2tform.py
deleted file mode 100644
index b2a8b54a9..000000000
--- a/modules/facelib/detection/matlab_cp2tform.py
+++ /dev/null
@@ -1,317 +0,0 @@
-import numpy as np
-from numpy.linalg import inv, lstsq
-from numpy.linalg import matrix_rank as rank
-from numpy.linalg import norm
-
-
-class MatlabCp2tormException(Exception):
-
- def __str__(self):
- return 'In File {}:{}'.format(__file__, super.__str__(self))
-
-
-def tformfwd(trans, uv):
- """
- Function:
- ----------
- apply affine transform 'trans' to uv
-
- Parameters:
- ----------
- @trans: 3x3 np.array
- transform matrix
- @uv: Kx2 np.array
- each row is a pair of coordinates (x, y)
-
- Returns:
- ----------
- @xy: Kx2 np.array
- each row is a pair of transformed coordinates (x, y)
- """
- uv = np.hstack((uv, np.ones((uv.shape[0], 1))))
- xy = np.dot(uv, trans)
- xy = xy[:, 0:-1]
- return xy
-
-
-def tforminv(trans, uv):
- """
- Function:
- ----------
- apply the inverse of affine transform 'trans' to uv
-
- Parameters:
- ----------
- @trans: 3x3 np.array
- transform matrix
- @uv: Kx2 np.array
- each row is a pair of coordinates (x, y)
-
- Returns:
- ----------
- @xy: Kx2 np.array
- each row is a pair of inverse-transformed coordinates (x, y)
- """
- Tinv = inv(trans)
- xy = tformfwd(Tinv, uv)
- return xy
-
-
-def findNonreflectiveSimilarity(uv, xy, options=None):
- options = {'K': 2}
-
- K = options['K']
- M = xy.shape[0]
- x = xy[:, 0].reshape((-1, 1)) # use reshape to keep a column vector
- y = xy[:, 1].reshape((-1, 1)) # use reshape to keep a column vector
-
- tmp1 = np.hstack((x, y, np.ones((M, 1)), np.zeros((M, 1))))
- tmp2 = np.hstack((y, -x, np.zeros((M, 1)), np.ones((M, 1))))
- X = np.vstack((tmp1, tmp2))
-
- u = uv[:, 0].reshape((-1, 1)) # use reshape to keep a column vector
- v = uv[:, 1].reshape((-1, 1)) # use reshape to keep a column vector
- U = np.vstack((u, v))
-
- # We know that X * r = U
- if rank(X) >= 2 * K:
- r, _, _, _ = lstsq(X, U, rcond=-1)
- r = np.squeeze(r)
- else:
- raise Exception('cp2tform:twoUniquePointsReq')
- sc = r[0]
- ss = r[1]
- tx = r[2]
- ty = r[3]
-
- Tinv = np.array([[sc, -ss, 0], [ss, sc, 0], [tx, ty, 1]])
- T = inv(Tinv)
- T[:, 2] = np.array([0, 0, 1])
-
- return T, Tinv
-
-
-def findSimilarity(uv, xy, options=None):
- options = {'K': 2}
-
- # uv = np.array(uv)
- # xy = np.array(xy)
-
- # Solve for trans1
- trans1, trans1_inv = findNonreflectiveSimilarity(uv, xy, options)
-
- # Solve for trans2
-
- # manually reflect the xy data across the Y-axis
- xyR = xy
- xyR[:, 0] = -1 * xyR[:, 0]
-
- trans2r, trans2r_inv = findNonreflectiveSimilarity(uv, xyR, options)
-
- # manually reflect the tform to undo the reflection done on xyR
- TreflectY = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, 1]])
-
- trans2 = np.dot(trans2r, TreflectY)
-
- # Figure out if trans1 or trans2 is better
- xy1 = tformfwd(trans1, uv)
- norm1 = norm(xy1 - xy)
-
- xy2 = tformfwd(trans2, uv)
- norm2 = norm(xy2 - xy)
-
- if norm1 <= norm2:
- return trans1, trans1_inv
- else:
- trans2_inv = inv(trans2)
- return trans2, trans2_inv
-
-
-def get_similarity_transform(src_pts, dst_pts, reflective=True):
- """
- Function:
- ----------
- Find Similarity Transform Matrix 'trans':
- u = src_pts[:, 0]
- v = src_pts[:, 1]
- x = dst_pts[:, 0]
- y = dst_pts[:, 1]
- [x, y, 1] = [u, v, 1] * trans
-
- Parameters:
- ----------
- @src_pts: Kx2 np.array
- source points, each row is a pair of coordinates (x, y)
- @dst_pts: Kx2 np.array
- destination points, each row is a pair of transformed
- coordinates (x, y)
- @reflective: True or False
- if True:
- use reflective similarity transform
- else:
- use non-reflective similarity transform
-
- Returns:
- ----------
- @trans: 3x3 np.array
- transform matrix from uv to xy
- trans_inv: 3x3 np.array
- inverse of trans, transform matrix from xy to uv
- """
-
- if reflective:
- trans, trans_inv = findSimilarity(src_pts, dst_pts)
- else:
- trans, trans_inv = findNonreflectiveSimilarity(src_pts, dst_pts)
-
- return trans, trans_inv
-
-
-def cvt_tform_mat_for_cv2(trans):
- """
- Function:
- ----------
- Convert Transform Matrix 'trans' into 'cv2_trans' which could be
- directly used by cv2.warpAffine():
- u = src_pts[:, 0]
- v = src_pts[:, 1]
- x = dst_pts[:, 0]
- y = dst_pts[:, 1]
- [x, y].T = cv_trans * [u, v, 1].T
-
- Parameters:
- ----------
- @trans: 3x3 np.array
- transform matrix from uv to xy
-
- Returns:
- ----------
- @cv2_trans: 2x3 np.array
- transform matrix from src_pts to dst_pts, could be directly used
- for cv2.warpAffine()
- """
- cv2_trans = trans[:, 0:2].T
-
- return cv2_trans
-
-
-def get_similarity_transform_for_cv2(src_pts, dst_pts, reflective=True):
- """
- Function:
- ----------
- Find Similarity Transform Matrix 'cv2_trans' which could be
- directly used by cv2.warpAffine():
- u = src_pts[:, 0]
- v = src_pts[:, 1]
- x = dst_pts[:, 0]
- y = dst_pts[:, 1]
- [x, y].T = cv_trans * [u, v, 1].T
-
- Parameters:
- ----------
- @src_pts: Kx2 np.array
- source points, each row is a pair of coordinates (x, y)
- @dst_pts: Kx2 np.array
- destination points, each row is a pair of transformed
- coordinates (x, y)
- reflective: True or False
- if True:
- use reflective similarity transform
- else:
- use non-reflective similarity transform
-
- Returns:
- ----------
- @cv2_trans: 2x3 np.array
- transform matrix from src_pts to dst_pts, could be directly used
- for cv2.warpAffine()
- """
- trans, trans_inv = get_similarity_transform(src_pts, dst_pts, reflective)
- cv2_trans = cvt_tform_mat_for_cv2(trans)
-
- return cv2_trans
-
-
-if __name__ == '__main__':
- """
- u = [0, 6, -2]
- v = [0, 3, 5]
- x = [-1, 0, 4]
- y = [-1, -10, 4]
-
- # In Matlab, run:
- #
- # uv = [u'; v'];
- # xy = [x'; y'];
- # tform_sim=cp2tform(uv,xy,'similarity');
- #
- # trans = tform_sim.tdata.T
- # ans =
- # -0.0764 -1.6190 0
- # 1.6190 -0.0764 0
- # -3.2156 0.0290 1.0000
- # trans_inv = tform_sim.tdata.Tinv
- # ans =
- #
- # -0.0291 0.6163 0
- # -0.6163 -0.0291 0
- # -0.0756 1.9826 1.0000
- # xy_m=tformfwd(tform_sim, u,v)
- #
- # xy_m =
- #
- # -3.2156 0.0290
- # 1.1833 -9.9143
- # 5.0323 2.8853
- # uv_m=tforminv(tform_sim, x,y)
- #
- # uv_m =
- #
- # 0.5698 1.3953
- # 6.0872 2.2733
- # -2.6570 4.3314
- """
- u = [0, 6, -2]
- v = [0, 3, 5]
- x = [-1, 0, 4]
- y = [-1, -10, 4]
-
- uv = np.array((u, v)).T
- xy = np.array((x, y)).T
-
- print('\n--->uv:')
- print(uv)
- print('\n--->xy:')
- print(xy)
-
- trans, trans_inv = get_similarity_transform(uv, xy)
-
- print('\n--->trans matrix:')
- print(trans)
-
- print('\n--->trans_inv matrix:')
- print(trans_inv)
-
- print('\n---> apply transform to uv')
- print('\nxy_m = uv_augmented * trans')
- uv_aug = np.hstack((uv, np.ones((uv.shape[0], 1))))
- xy_m = np.dot(uv_aug, trans)
- print(xy_m)
-
- print('\nxy_m = tformfwd(trans, uv)')
- xy_m = tformfwd(trans, uv)
- print(xy_m)
-
- print('\n---> apply inverse transform to xy')
- print('\nuv_m = xy_augmented * trans_inv')
- xy_aug = np.hstack((xy, np.ones((xy.shape[0], 1))))
- uv_m = np.dot(xy_aug, trans_inv)
- print(uv_m)
-
- print('\nuv_m = tformfwd(trans_inv, xy)')
- uv_m = tformfwd(trans_inv, xy)
- print(uv_m)
-
- uv_m = tforminv(trans, xy)
- print('\nuv_m = tforminv(trans, xy)')
- print(uv_m)
diff --git a/modules/facelib/detection/retinaface/retinaface.py b/modules/facelib/detection/retinaface/retinaface.py
deleted file mode 100644
index 92a226ca3..000000000
--- a/modules/facelib/detection/retinaface/retinaface.py
+++ /dev/null
@@ -1,370 +0,0 @@
-import cv2
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from PIL import Image
-from torchvision.models._utils import IntermediateLayerGetter as IntermediateLayerGetter
-
-from ...detection.align_trans import get_reference_facial_points, warp_and_crop_face
-from ...detection.retinaface.retinaface_net import FPN, SSH, MobileNetV1, make_bbox_head, make_class_head, make_landmark_head
-from ...detection.retinaface.retinaface_utils import (PriorBox, batched_decode, batched_decode_landm, decode, decode_landm,
- py_cpu_nms)
-
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-
-def generate_config(network_name):
-
- cfg_mnet = {
- 'name': 'mobilenet0.25',
- 'min_sizes': [[16, 32], [64, 128], [256, 512]],
- 'steps': [8, 16, 32],
- 'variance': [0.1, 0.2],
- 'clip': False,
- 'loc_weight': 2.0,
- 'gpu_train': True,
- 'batch_size': 32,
- 'ngpu': 1,
- 'epoch': 250,
- 'decay1': 190,
- 'decay2': 220,
- 'image_size': 640,
- 'return_layers': {
- 'stage1': 1,
- 'stage2': 2,
- 'stage3': 3
- },
- 'in_channel': 32,
- 'out_channel': 64
- }
-
- cfg_re50 = {
- 'name': 'Resnet50',
- 'min_sizes': [[16, 32], [64, 128], [256, 512]],
- 'steps': [8, 16, 32],
- 'variance': [0.1, 0.2],
- 'clip': False,
- 'loc_weight': 2.0,
- 'gpu_train': True,
- 'batch_size': 24,
- 'ngpu': 4,
- 'epoch': 100,
- 'decay1': 70,
- 'decay2': 90,
- 'image_size': 840,
- 'return_layers': {
- 'layer2': 1,
- 'layer3': 2,
- 'layer4': 3
- },
- 'in_channel': 256,
- 'out_channel': 256
- }
-
- if network_name == 'mobile0.25':
- return cfg_mnet
- elif network_name == 'resnet50':
- return cfg_re50
- else:
- raise NotImplementedError(f'network_name={network_name}')
-
-
-class RetinaFace(nn.Module):
-
- def __init__(self, network_name='resnet50', half=False, phase='test'):
- super(RetinaFace, self).__init__()
- self.half_inference = half
- cfg = generate_config(network_name)
- self.backbone = cfg['name']
-
- self.model_name = f'retinaface_{network_name}'
- self.cfg = cfg
- self.phase = phase
- self.target_size, self.max_size = 1600, 2150
- self.resize, self.scale, self.scale1 = 1., None, None
- self.mean_tensor = torch.tensor([[[[104.]], [[117.]], [[123.]]]]).to(device)
- self.reference = get_reference_facial_points(default_square=True)
- # Build network.
- backbone = None
- if cfg['name'] == 'mobilenet0.25':
- backbone = MobileNetV1()
- self.body = IntermediateLayerGetter(backbone, cfg['return_layers'])
- elif cfg['name'] == 'Resnet50':
- import torchvision.models as models
- backbone = models.resnet50(pretrained=False)
- self.body = IntermediateLayerGetter(backbone, cfg['return_layers'])
-
- in_channels_stage2 = cfg['in_channel']
- in_channels_list = [
- in_channels_stage2 * 2,
- in_channels_stage2 * 4,
- in_channels_stage2 * 8,
- ]
-
- out_channels = cfg['out_channel']
- self.fpn = FPN(in_channels_list, out_channels)
- self.ssh1 = SSH(out_channels, out_channels)
- self.ssh2 = SSH(out_channels, out_channels)
- self.ssh3 = SSH(out_channels, out_channels)
-
- self.ClassHead = make_class_head(fpn_num=3, inchannels=cfg['out_channel'])
- self.BboxHead = make_bbox_head(fpn_num=3, inchannels=cfg['out_channel'])
- self.LandmarkHead = make_landmark_head(fpn_num=3, inchannels=cfg['out_channel'])
-
- self.to(device)
- self.eval()
- if self.half_inference:
- self.half()
-
- def forward(self, inputs):
- out = self.body(inputs)
-
- if self.backbone == 'mobilenet0.25' or self.backbone == 'Resnet50':
- out = list(out.values())
- # FPN
- fpn = self.fpn(out)
-
- # SSH
- feature1 = self.ssh1(fpn[0])
- feature2 = self.ssh2(fpn[1])
- feature3 = self.ssh3(fpn[2])
- features = [feature1, feature2, feature3]
-
- bbox_regressions = torch.cat([self.BboxHead[i](feature) for i, feature in enumerate(features)], dim=1)
- classifications = torch.cat([self.ClassHead[i](feature) for i, feature in enumerate(features)], dim=1)
- tmp = [self.LandmarkHead[i](feature) for i, feature in enumerate(features)]
- ldm_regressions = (torch.cat(tmp, dim=1))
-
- if self.phase == 'train':
- output = (bbox_regressions, classifications, ldm_regressions)
- else:
- output = (bbox_regressions, F.softmax(classifications, dim=-1), ldm_regressions)
- return output
-
- def __detect_faces(self, inputs):
- # get scale
- height, width = inputs.shape[2:]
- self.scale = torch.tensor([width, height, width, height], dtype=torch.float32).to(device)
- tmp = [width, height, width, height, width, height, width, height, width, height]
- self.scale1 = torch.tensor(tmp, dtype=torch.float32).to(device)
-
- # forawrd
- inputs = inputs.to(device)
- if self.half_inference:
- inputs = inputs.half()
- loc, conf, landmarks = self(inputs)
-
- # get priorbox
- priorbox = PriorBox(self.cfg, image_size=inputs.shape[2:])
- priors = priorbox.forward().to(device)
-
- return loc, conf, landmarks, priors
-
- # single image detection
- def transform(self, image, use_origin_size):
- # convert to opencv format
- if isinstance(image, Image.Image):
- image = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)
- image = image.astype(np.float32)
-
- # testing scale
- im_size_min = np.min(image.shape[0:2])
- im_size_max = np.max(image.shape[0:2])
- resize = float(self.target_size) / float(im_size_min)
-
- # prevent bigger axis from being more than max_size
- if np.round(resize * im_size_max) > self.max_size:
- resize = float(self.max_size) / float(im_size_max)
- resize = 1 if use_origin_size else resize
-
- # resize
- if resize != 1:
- image = cv2.resize(image, None, None, fx=resize, fy=resize, interpolation=cv2.INTER_LINEAR)
-
- # convert to torch.tensor format
- # image -= (104, 117, 123)
- image = image.transpose(2, 0, 1)
- image = torch.from_numpy(image).unsqueeze(0)
-
- return image, resize
-
- def detect_faces(
- self,
- image,
- conf_threshold=0.8,
- nms_threshold=0.4,
- use_origin_size=True,
- ):
- """
- Params:
- imgs: BGR image
- """
- image, self.resize = self.transform(image, use_origin_size)
- image = image.to(device)
- if self.half_inference:
- image = image.half()
- image = image - self.mean_tensor
-
- loc, conf, landmarks, priors = self.__detect_faces(image)
-
- boxes = decode(loc.data.squeeze(0), priors.data, self.cfg['variance'])
- boxes = boxes * self.scale / self.resize
- boxes = boxes.cpu().numpy()
-
- scores = conf.squeeze(0).data.cpu().numpy()[:, 1]
-
- landmarks = decode_landm(landmarks.squeeze(0), priors, self.cfg['variance'])
- landmarks = landmarks * self.scale1 / self.resize
- landmarks = landmarks.cpu().numpy()
-
- # ignore low scores
- inds = np.where(scores > conf_threshold)[0]
- boxes, landmarks, scores = boxes[inds], landmarks[inds], scores[inds]
-
- # sort
- order = scores.argsort()[::-1]
- boxes, landmarks, scores = boxes[order], landmarks[order], scores[order]
-
- # do NMS
- bounding_boxes = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
- keep = py_cpu_nms(bounding_boxes, nms_threshold)
- bounding_boxes, landmarks = bounding_boxes[keep, :], landmarks[keep]
- # self.t['forward_pass'].toc()
- # print(self.t['forward_pass'].average_time)
- # import sys
- # sys.stdout.flush()
- return np.concatenate((bounding_boxes, landmarks), axis=1)
-
- def __align_multi(self, image, boxes, landmarks, limit=None):
-
- if len(boxes) < 1:
- return [], []
-
- if limit:
- boxes = boxes[:limit]
- landmarks = landmarks[:limit]
-
- faces = []
- for landmark in landmarks:
- facial5points = [[landmark[2 * j], landmark[2 * j + 1]] for j in range(5)]
-
- warped_face = warp_and_crop_face(np.array(image), facial5points, self.reference, crop_size=(112, 112))
- faces.append(warped_face)
-
- return np.concatenate((boxes, landmarks), axis=1), faces
-
- def align_multi(self, img, conf_threshold=0.8, limit=None):
-
- rlt = self.detect_faces(img, conf_threshold=conf_threshold)
- boxes, landmarks = rlt[:, 0:5], rlt[:, 5:]
-
- return self.__align_multi(img, boxes, landmarks, limit)
-
- # batched detection
- def batched_transform(self, frames, use_origin_size):
- """
- Arguments:
- frames: a list of PIL.Image, or torch.Tensor(shape=[n, h, w, c],
- type=np.float32, BGR format).
- use_origin_size: whether to use origin size.
- """
- from_PIL = True if isinstance(frames[0], Image.Image) else False
-
- # convert to opencv format
- if from_PIL:
- frames = [cv2.cvtColor(np.asarray(frame), cv2.COLOR_RGB2BGR) for frame in frames]
- frames = np.asarray(frames, dtype=np.float32)
-
- # testing scale
- im_size_min = np.min(frames[0].shape[0:2])
- im_size_max = np.max(frames[0].shape[0:2])
- resize = float(self.target_size) / float(im_size_min)
-
- # prevent bigger axis from being more than max_size
- if np.round(resize * im_size_max) > self.max_size:
- resize = float(self.max_size) / float(im_size_max)
- resize = 1 if use_origin_size else resize
-
- # resize
- if resize != 1:
- if not from_PIL:
- frames = F.interpolate(frames, scale_factor=resize)
- else:
- frames = [
- cv2.resize(frame, None, None, fx=resize, fy=resize, interpolation=cv2.INTER_LINEAR)
- for frame in frames
- ]
-
- # convert to torch.tensor format
- if not from_PIL:
- frames = frames.transpose(1, 2).transpose(1, 3).contiguous()
- else:
- frames = frames.transpose((0, 3, 1, 2))
- frames = torch.from_numpy(frames)
-
- return frames, resize
-
- def batched_detect_faces(self, frames, conf_threshold=0.8, nms_threshold=0.4, use_origin_size=True):
- """
- Arguments:
- frames: a list of PIL.Image, or np.array(shape=[n, h, w, c],
- type=np.uint8, BGR format).
- conf_threshold: confidence threshold.
- nms_threshold: nms threshold.
- use_origin_size: whether to use origin size.
- Returns:
- final_bounding_boxes: list of np.array ([n_boxes, 5],
- type=np.float32).
- final_landmarks: list of np.array ([n_boxes, 10], type=np.float32).
- """
- # self.t['forward_pass'].tic()
- frames, self.resize = self.batched_transform(frames, use_origin_size)
- frames = frames.to(device)
- frames = frames - self.mean_tensor
-
- b_loc, b_conf, b_landmarks, priors = self.__detect_faces(frames)
-
- final_bounding_boxes, final_landmarks = [], []
-
- # decode
- priors = priors.unsqueeze(0)
- b_loc = batched_decode(b_loc, priors, self.cfg['variance']) * self.scale / self.resize
- b_landmarks = batched_decode_landm(b_landmarks, priors, self.cfg['variance']) * self.scale1 / self.resize
- b_conf = b_conf[:, :, 1]
-
- # index for selection
- b_indice = b_conf > conf_threshold
-
- # concat
- b_loc_and_conf = torch.cat((b_loc, b_conf.unsqueeze(-1)), dim=2).float()
-
- for pred, landm, inds in zip(b_loc_and_conf, b_landmarks, b_indice):
-
- # ignore low scores
- pred, landm = pred[inds, :], landm[inds, :]
- if pred.shape[0] == 0:
- final_bounding_boxes.append(np.array([], dtype=np.float32))
- final_landmarks.append(np.array([], dtype=np.float32))
- continue
-
- # sort
- # order = score.argsort(descending=True)
- # box, landm, score = box[order], landm[order], score[order]
-
- # to CPU
- bounding_boxes, landm = pred.cpu().numpy(), landm.cpu().numpy()
-
- # NMS
- keep = py_cpu_nms(bounding_boxes, nms_threshold)
- bounding_boxes, landmarks = bounding_boxes[keep, :], landm[keep]
-
- # append
- final_bounding_boxes.append(bounding_boxes)
- final_landmarks.append(landmarks)
- # self.t['forward_pass'].toc(average=True)
- # self.batch_time += self.t['forward_pass'].diff
- # self.total_frame += len(frames)
- # print(self.batch_time / self.total_frame)
-
- return final_bounding_boxes, final_landmarks
diff --git a/modules/facelib/detection/retinaface/retinaface_net.py b/modules/facelib/detection/retinaface/retinaface_net.py
deleted file mode 100644
index ab6aa82d3..000000000
--- a/modules/facelib/detection/retinaface/retinaface_net.py
+++ /dev/null
@@ -1,196 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-def conv_bn(inp, oup, stride=1, leaky=0):
- return nn.Sequential(
- nn.Conv2d(inp, oup, 3, stride, 1, bias=False), nn.BatchNorm2d(oup),
- nn.LeakyReLU(negative_slope=leaky, inplace=True))
-
-
-def conv_bn_no_relu(inp, oup, stride):
- return nn.Sequential(
- nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
- nn.BatchNorm2d(oup),
- )
-
-
-def conv_bn1X1(inp, oup, stride, leaky=0):
- return nn.Sequential(
- nn.Conv2d(inp, oup, 1, stride, padding=0, bias=False), nn.BatchNorm2d(oup),
- nn.LeakyReLU(negative_slope=leaky, inplace=True))
-
-
-def conv_dw(inp, oup, stride, leaky=0.1):
- return nn.Sequential(
- nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
- nn.BatchNorm2d(inp),
- nn.LeakyReLU(negative_slope=leaky, inplace=True),
- nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
- nn.BatchNorm2d(oup),
- nn.LeakyReLU(negative_slope=leaky, inplace=True),
- )
-
-
-class SSH(nn.Module):
-
- def __init__(self, in_channel, out_channel):
- super(SSH, self).__init__()
- assert out_channel % 4 == 0
- leaky = 0
- if (out_channel <= 64):
- leaky = 0.1
- self.conv3X3 = conv_bn_no_relu(in_channel, out_channel // 2, stride=1)
-
- self.conv5X5_1 = conv_bn(in_channel, out_channel // 4, stride=1, leaky=leaky)
- self.conv5X5_2 = conv_bn_no_relu(out_channel // 4, out_channel // 4, stride=1)
-
- self.conv7X7_2 = conv_bn(out_channel // 4, out_channel // 4, stride=1, leaky=leaky)
- self.conv7x7_3 = conv_bn_no_relu(out_channel // 4, out_channel // 4, stride=1)
-
- def forward(self, input):
- conv3X3 = self.conv3X3(input)
-
- conv5X5_1 = self.conv5X5_1(input)
- conv5X5 = self.conv5X5_2(conv5X5_1)
-
- conv7X7_2 = self.conv7X7_2(conv5X5_1)
- conv7X7 = self.conv7x7_3(conv7X7_2)
-
- out = torch.cat([conv3X3, conv5X5, conv7X7], dim=1)
- out = F.relu(out)
- return out
-
-
-class FPN(nn.Module):
-
- def __init__(self, in_channels_list, out_channels):
- super(FPN, self).__init__()
- leaky = 0
- if (out_channels <= 64):
- leaky = 0.1
- self.output1 = conv_bn1X1(in_channels_list[0], out_channels, stride=1, leaky=leaky)
- self.output2 = conv_bn1X1(in_channels_list[1], out_channels, stride=1, leaky=leaky)
- self.output3 = conv_bn1X1(in_channels_list[2], out_channels, stride=1, leaky=leaky)
-
- self.merge1 = conv_bn(out_channels, out_channels, leaky=leaky)
- self.merge2 = conv_bn(out_channels, out_channels, leaky=leaky)
-
- def forward(self, input):
- # names = list(input.keys())
- # input = list(input.values())
-
- output1 = self.output1(input[0])
- output2 = self.output2(input[1])
- output3 = self.output3(input[2])
-
- up3 = F.interpolate(output3, size=[output2.size(2), output2.size(3)], mode='nearest')
- output2 = output2 + up3
- output2 = self.merge2(output2)
-
- up2 = F.interpolate(output2, size=[output1.size(2), output1.size(3)], mode='nearest')
- output1 = output1 + up2
- output1 = self.merge1(output1)
-
- out = [output1, output2, output3]
- return out
-
-
-class MobileNetV1(nn.Module):
-
- def __init__(self):
- super(MobileNetV1, self).__init__()
- self.stage1 = nn.Sequential(
- conv_bn(3, 8, 2, leaky=0.1), # 3
- conv_dw(8, 16, 1), # 7
- conv_dw(16, 32, 2), # 11
- conv_dw(32, 32, 1), # 19
- conv_dw(32, 64, 2), # 27
- conv_dw(64, 64, 1), # 43
- )
- self.stage2 = nn.Sequential(
- conv_dw(64, 128, 2), # 43 + 16 = 59
- conv_dw(128, 128, 1), # 59 + 32 = 91
- conv_dw(128, 128, 1), # 91 + 32 = 123
- conv_dw(128, 128, 1), # 123 + 32 = 155
- conv_dw(128, 128, 1), # 155 + 32 = 187
- conv_dw(128, 128, 1), # 187 + 32 = 219
- )
- self.stage3 = nn.Sequential(
- conv_dw(128, 256, 2), # 219 +3 2 = 241
- conv_dw(256, 256, 1), # 241 + 64 = 301
- )
- self.avg = nn.AdaptiveAvgPool2d((1, 1))
- self.fc = nn.Linear(256, 1000)
-
- def forward(self, x):
- x = self.stage1(x)
- x = self.stage2(x)
- x = self.stage3(x)
- x = self.avg(x)
- # x = self.model(x)
- x = x.view(-1, 256)
- x = self.fc(x)
- return x
-
-
-class ClassHead(nn.Module):
-
- def __init__(self, inchannels=512, num_anchors=3):
- super(ClassHead, self).__init__()
- self.num_anchors = num_anchors
- self.conv1x1 = nn.Conv2d(inchannels, self.num_anchors * 2, kernel_size=(1, 1), stride=1, padding=0)
-
- def forward(self, x):
- out = self.conv1x1(x)
- out = out.permute(0, 2, 3, 1).contiguous()
-
- return out.view(out.shape[0], -1, 2)
-
-
-class BboxHead(nn.Module):
-
- def __init__(self, inchannels=512, num_anchors=3):
- super(BboxHead, self).__init__()
- self.conv1x1 = nn.Conv2d(inchannels, num_anchors * 4, kernel_size=(1, 1), stride=1, padding=0)
-
- def forward(self, x):
- out = self.conv1x1(x)
- out = out.permute(0, 2, 3, 1).contiguous()
-
- return out.view(out.shape[0], -1, 4)
-
-
-class LandmarkHead(nn.Module):
-
- def __init__(self, inchannels=512, num_anchors=3):
- super(LandmarkHead, self).__init__()
- self.conv1x1 = nn.Conv2d(inchannels, num_anchors * 10, kernel_size=(1, 1), stride=1, padding=0)
-
- def forward(self, x):
- out = self.conv1x1(x)
- out = out.permute(0, 2, 3, 1).contiguous()
-
- return out.view(out.shape[0], -1, 10)
-
-
-def make_class_head(fpn_num=3, inchannels=64, anchor_num=2):
- classhead = nn.ModuleList()
- for i in range(fpn_num):
- classhead.append(ClassHead(inchannels, anchor_num))
- return classhead
-
-
-def make_bbox_head(fpn_num=3, inchannels=64, anchor_num=2):
- bboxhead = nn.ModuleList()
- for i in range(fpn_num):
- bboxhead.append(BboxHead(inchannels, anchor_num))
- return bboxhead
-
-
-def make_landmark_head(fpn_num=3, inchannels=64, anchor_num=2):
- landmarkhead = nn.ModuleList()
- for i in range(fpn_num):
- landmarkhead.append(LandmarkHead(inchannels, anchor_num))
- return landmarkhead
diff --git a/modules/facelib/detection/retinaface/retinaface_utils.py b/modules/facelib/detection/retinaface/retinaface_utils.py
deleted file mode 100644
index da38d6d13..000000000
--- a/modules/facelib/detection/retinaface/retinaface_utils.py
+++ /dev/null
@@ -1,420 +0,0 @@
-import numpy as np
-import torch
-import torchvision
-from itertools import product as product
-from math import ceil
-
-
-class PriorBox(object):
-
- def __init__(self, cfg, image_size=None, phase='train'):
- super(PriorBox, self).__init__()
- self.min_sizes = cfg['min_sizes']
- self.steps = cfg['steps']
- self.clip = cfg['clip']
- self.image_size = image_size
- self.feature_maps = [[ceil(self.image_size[0] / step), ceil(self.image_size[1] / step)] for step in self.steps]
- self.name = 's'
-
- def forward(self):
- anchors = []
- for k, f in enumerate(self.feature_maps):
- min_sizes = self.min_sizes[k]
- for i, j in product(range(f[0]), range(f[1])):
- for min_size in min_sizes:
- s_kx = min_size / self.image_size[1]
- s_ky = min_size / self.image_size[0]
- dense_cx = [x * self.steps[k] / self.image_size[1] for x in [j + 0.5]]
- dense_cy = [y * self.steps[k] / self.image_size[0] for y in [i + 0.5]]
- for cy, cx in product(dense_cy, dense_cx):
- anchors += [cx, cy, s_kx, s_ky]
-
- # back to torch land
- output = torch.Tensor(anchors).view(-1, 4)
- if self.clip:
- output.clamp_(max=1, min=0)
- return output
-
-
-def py_cpu_nms(dets, thresh):
- """Pure Python NMS baseline."""
- keep = torchvision.ops.nms(
- boxes=torch.Tensor(dets[:, :4]),
- scores=torch.Tensor(dets[:, 4]),
- iou_threshold=thresh,
- )
-
- return list(keep)
-
-
-def point_form(boxes):
- """ Convert prior_boxes to (xmin, ymin, xmax, ymax)
- representation for comparison to point form ground truth data.
- Args:
- boxes: (tensor) center-size default boxes from priorbox layers.
- Return:
- boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.
- """
- return torch.cat(
- (
- boxes[:, :2] - boxes[:, 2:] / 2, # xmin, ymin
- boxes[:, :2] + boxes[:, 2:] / 2),
- 1) # xmax, ymax
-
-
-def center_size(boxes):
- """ Convert prior_boxes to (cx, cy, w, h)
- representation for comparison to center-size form ground truth data.
- Args:
- boxes: (tensor) point_form boxes
- Return:
- boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.
- """
- return torch.cat(
- (boxes[:, 2:] + boxes[:, :2]) / 2, # cx, cy
- boxes[:, 2:] - boxes[:, :2],
- 1) # w, h
-
-
-def intersect(box_a, box_b):
- """ We resize both tensors to [A,B,2] without new malloc:
- [A,2] -> [A,1,2] -> [A,B,2]
- [B,2] -> [1,B,2] -> [A,B,2]
- Then we compute the area of intersect between box_a and box_b.
- Args:
- box_a: (tensor) bounding boxes, Shape: [A,4].
- box_b: (tensor) bounding boxes, Shape: [B,4].
- Return:
- (tensor) intersection area, Shape: [A,B].
- """
- A = box_a.size(0)
- B = box_b.size(0)
- max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2), box_b[:, 2:].unsqueeze(0).expand(A, B, 2))
- min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2), box_b[:, :2].unsqueeze(0).expand(A, B, 2))
- inter = torch.clamp((max_xy - min_xy), min=0)
- return inter[:, :, 0] * inter[:, :, 1]
-
-
-def jaccard(box_a, box_b):
- """Compute the jaccard overlap of two sets of boxes. The jaccard overlap
- is simply the intersection over union of two boxes. Here we operate on
- ground truth boxes and default boxes.
- E.g.:
- A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)
- Args:
- box_a: (tensor) Ground truth bounding boxes, Shape: [num_objects,4]
- box_b: (tensor) Prior boxes from priorbox layers, Shape: [num_priors,4]
- Return:
- jaccard overlap: (tensor) Shape: [box_a.size(0), box_b.size(0)]
- """
- inter = intersect(box_a, box_b)
- area_a = ((box_a[:, 2] - box_a[:, 0]) * (box_a[:, 3] - box_a[:, 1])).unsqueeze(1).expand_as(inter) # [A,B]
- area_b = ((box_b[:, 2] - box_b[:, 0]) * (box_b[:, 3] - box_b[:, 1])).unsqueeze(0).expand_as(inter) # [A,B]
- union = area_a + area_b - inter
- return inter / union # [A,B]
-
-
-def matrix_iou(a, b):
- """
- return iou of a and b, numpy version for data augenmentation
- """
- lt = np.maximum(a[:, np.newaxis, :2], b[:, :2])
- rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])
-
- area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2)
- area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)
- area_b = np.prod(b[:, 2:] - b[:, :2], axis=1)
- return area_i / (area_a[:, np.newaxis] + area_b - area_i)
-
-
-def matrix_iof(a, b):
- """
- return iof of a and b, numpy version for data augenmentation
- """
- lt = np.maximum(a[:, np.newaxis, :2], b[:, :2])
- rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])
-
- area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2)
- area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)
- return area_i / np.maximum(area_a[:, np.newaxis], 1)
-
-
-def match(threshold, truths, priors, variances, labels, landms, loc_t, conf_t, landm_t, idx):
- """Match each prior box with the ground truth box of the highest jaccard
- overlap, encode the bounding boxes, then return the matched indices
- corresponding to both confidence and location preds.
- Args:
- threshold: (float) The overlap threshold used when matching boxes.
- truths: (tensor) Ground truth boxes, Shape: [num_obj, 4].
- priors: (tensor) Prior boxes from priorbox layers, Shape: [n_priors,4].
- variances: (tensor) Variances corresponding to each prior coord,
- Shape: [num_priors, 4].
- labels: (tensor) All the class labels for the image, Shape: [num_obj].
- landms: (tensor) Ground truth landms, Shape [num_obj, 10].
- loc_t: (tensor) Tensor to be filled w/ encoded location targets.
- conf_t: (tensor) Tensor to be filled w/ matched indices for conf preds.
- landm_t: (tensor) Tensor to be filled w/ encoded landm targets.
- idx: (int) current batch index
- Return:
- The matched indices corresponding to 1)location 2)confidence
- 3)landm preds.
- """
- # jaccard index
- overlaps = jaccard(truths, point_form(priors))
- # (Bipartite Matching)
- # [1,num_objects] best prior for each ground truth
- best_prior_overlap, best_prior_idx = overlaps.max(1, keepdim=True)
-
- # ignore hard gt
- valid_gt_idx = best_prior_overlap[:, 0] >= 0.2
- best_prior_idx_filter = best_prior_idx[valid_gt_idx, :]
- if best_prior_idx_filter.shape[0] <= 0:
- loc_t[idx] = 0
- conf_t[idx] = 0
- return
-
- # [1,num_priors] best ground truth for each prior
- best_truth_overlap, best_truth_idx = overlaps.max(0, keepdim=True)
- best_truth_idx.squeeze_(0)
- best_truth_overlap.squeeze_(0)
- best_prior_idx.squeeze_(1)
- best_prior_idx_filter.squeeze_(1)
- best_prior_overlap.squeeze_(1)
- best_truth_overlap.index_fill_(0, best_prior_idx_filter, 2) # ensure best prior
- # ensure every gt matches with its prior of max overlap
- for j in range(best_prior_idx.size(0)): # 判别此anchor是预测哪一个boxes
- best_truth_idx[best_prior_idx[j]] = j
- matches = truths[best_truth_idx] # Shape: [num_priors,4] 此处为每一个anchor对应的bbox取出来
- conf = labels[best_truth_idx] # Shape: [num_priors] 此处为每一个anchor对应的label取出来
- conf[best_truth_overlap < threshold] = 0 # label as background overlap<0.35的全部作为负样本
- loc = encode(matches, priors, variances)
-
- matches_landm = landms[best_truth_idx]
- landm = encode_landm(matches_landm, priors, variances)
- loc_t[idx] = loc # [num_priors,4] encoded offsets to learn
- conf_t[idx] = conf # [num_priors] top class label for each prior
- landm_t[idx] = landm
-
-
-def encode(matched, priors, variances):
- """Encode the variances from the priorbox layers into the ground truth boxes
- we have matched (based on jaccard overlap) with the prior boxes.
- Args:
- matched: (tensor) Coords of ground truth for each prior in point-form
- Shape: [num_priors, 4].
- priors: (tensor) Prior boxes in center-offset form
- Shape: [num_priors,4].
- variances: (list[float]) Variances of priorboxes
- Return:
- encoded boxes (tensor), Shape: [num_priors, 4]
- """
-
- # dist b/t match center and prior's center
- g_cxcy = (matched[:, :2] + matched[:, 2:]) / 2 - priors[:, :2]
- # encode variance
- g_cxcy /= (variances[0] * priors[:, 2:])
- # match wh / prior wh
- g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:]
- g_wh = torch.log(g_wh) / variances[1]
- # return target for smooth_l1_loss
- return torch.cat([g_cxcy, g_wh], 1) # [num_priors,4]
-
-
-def encode_landm(matched, priors, variances):
- """Encode the variances from the priorbox layers into the ground truth boxes
- we have matched (based on jaccard overlap) with the prior boxes.
- Args:
- matched: (tensor) Coords of ground truth for each prior in point-form
- Shape: [num_priors, 10].
- priors: (tensor) Prior boxes in center-offset form
- Shape: [num_priors,4].
- variances: (list[float]) Variances of priorboxes
- Return:
- encoded landm (tensor), Shape: [num_priors, 10]
- """
-
- # dist b/t match center and prior's center
- matched = torch.reshape(matched, (matched.size(0), 5, 2))
- priors_cx = priors[:, 0].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
- priors_cy = priors[:, 1].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
- priors_w = priors[:, 2].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
- priors_h = priors[:, 3].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
- priors = torch.cat([priors_cx, priors_cy, priors_w, priors_h], dim=2)
- g_cxcy = matched[:, :, :2] - priors[:, :, :2]
- # encode variance
- g_cxcy /= (variances[0] * priors[:, :, 2:])
- # g_cxcy /= priors[:, :, 2:]
- g_cxcy = g_cxcy.reshape(g_cxcy.size(0), -1)
- # return target for smooth_l1_loss
- return g_cxcy
-
-
-# Adapted from https://github.com/Hakuyume/chainer-ssd
-def decode(loc, priors, variances):
- """Decode locations from predictions using priors to undo
- the encoding we did for offset regression at train time.
- Args:
- loc (tensor): location predictions for loc layers,
- Shape: [num_priors,4]
- priors (tensor): Prior boxes in center-offset form.
- Shape: [num_priors,4].
- variances: (list[float]) Variances of priorboxes
- Return:
- decoded bounding box predictions
- """
-
- boxes = torch.cat((priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
- priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)
- boxes[:, :2] -= boxes[:, 2:] / 2
- boxes[:, 2:] += boxes[:, :2]
- return boxes
-
-
-def decode_landm(pre, priors, variances):
- """Decode landm from predictions using priors to undo
- the encoding we did for offset regression at train time.
- Args:
- pre (tensor): landm predictions for loc layers,
- Shape: [num_priors,10]
- priors (tensor): Prior boxes in center-offset form.
- Shape: [num_priors,4].
- variances: (list[float]) Variances of priorboxes
- Return:
- decoded landm predictions
- """
- tmp = (
- priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:],
- priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:],
- priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:],
- priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:],
- priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:],
- )
- landms = torch.cat(tmp, dim=1)
- return landms
-
-
-def batched_decode(b_loc, priors, variances):
- """Decode locations from predictions using priors to undo
- the encoding we did for offset regression at train time.
- Args:
- b_loc (tensor): location predictions for loc layers,
- Shape: [num_batches,num_priors,4]
- priors (tensor): Prior boxes in center-offset form.
- Shape: [1,num_priors,4].
- variances: (list[float]) Variances of priorboxes
- Return:
- decoded bounding box predictions
- """
- boxes = (
- priors[:, :, :2] + b_loc[:, :, :2] * variances[0] * priors[:, :, 2:],
- priors[:, :, 2:] * torch.exp(b_loc[:, :, 2:] * variances[1]),
- )
- boxes = torch.cat(boxes, dim=2)
-
- boxes[:, :, :2] -= boxes[:, :, 2:] / 2
- boxes[:, :, 2:] += boxes[:, :, :2]
- return boxes
-
-
-def batched_decode_landm(pre, priors, variances):
- """Decode landm from predictions using priors to undo
- the encoding we did for offset regression at train time.
- Args:
- pre (tensor): landm predictions for loc layers,
- Shape: [num_batches,num_priors,10]
- priors (tensor): Prior boxes in center-offset form.
- Shape: [1,num_priors,4].
- variances: (list[float]) Variances of priorboxes
- Return:
- decoded landm predictions
- """
- landms = (
- priors[:, :, :2] + pre[:, :, :2] * variances[0] * priors[:, :, 2:],
- priors[:, :, :2] + pre[:, :, 2:4] * variances[0] * priors[:, :, 2:],
- priors[:, :, :2] + pre[:, :, 4:6] * variances[0] * priors[:, :, 2:],
- priors[:, :, :2] + pre[:, :, 6:8] * variances[0] * priors[:, :, 2:],
- priors[:, :, :2] + pre[:, :, 8:10] * variances[0] * priors[:, :, 2:],
- )
- landms = torch.cat(landms, dim=2)
- return landms
-
-
-def log_sum_exp(x):
- """Utility function for computing log_sum_exp while determining
- This will be used to determine unaveraged confidence loss across
- all examples in a batch.
- Args:
- x (Variable(tensor)): conf_preds from conf layers
- """
- x_max = x.data.max()
- return torch.log(torch.sum(torch.exp(x - x_max), 1, keepdim=True)) + x_max
-
-
-# Original author: Francisco Massa:
-# https://github.com/fmassa/object-detection.torch
-# Ported to PyTorch by Max deGroot (02/01/2017)
-def nms(boxes, scores, overlap=0.5, top_k=200):
- """Apply non-maximum suppression at test time to avoid detecting too many
- overlapping bounding boxes for a given object.
- Args:
- boxes: (tensor) The location preds for the img, Shape: [num_priors,4].
- scores: (tensor) The class predscores for the img, Shape:[num_priors].
- overlap: (float) The overlap thresh for suppressing unnecessary boxes.
- top_k: (int) The Maximum number of box preds to consider.
- Return:
- The indices of the kept boxes with respect to num_priors.
- """
-
- keep = torch.Tensor(scores.size(0)).fill_(0).long()
- if boxes.numel() == 0:
- return keep
- x1 = boxes[:, 0]
- y1 = boxes[:, 1]
- x2 = boxes[:, 2]
- y2 = boxes[:, 3]
- area = torch.mul(x2 - x1, y2 - y1)
- v, idx = scores.sort(0) # sort in ascending order
- # I = I[v >= 0.01]
- idx = idx[-top_k:] # indices of the top-k largest vals
- xx1 = boxes.new()
- yy1 = boxes.new()
- xx2 = boxes.new()
- yy2 = boxes.new()
- w = boxes.new()
- h = boxes.new()
-
- # keep = torch.Tensor()
- count = 0
- while idx.numel() > 0:
- i = idx[-1] # index of current largest val
- # keep.append(i)
- keep[count] = i
- count += 1
- if idx.size(0) == 1:
- break
- idx = idx[:-1] # remove kept element from view
- # load bboxes of next highest vals
- torch.index_select(x1, 0, idx, out=xx1)
- torch.index_select(y1, 0, idx, out=yy1)
- torch.index_select(x2, 0, idx, out=xx2)
- torch.index_select(y2, 0, idx, out=yy2)
- # store element-wise max with next highest score
- xx1 = torch.clamp(xx1, min=x1[i])
- yy1 = torch.clamp(yy1, min=y1[i])
- xx2 = torch.clamp(xx2, max=x2[i])
- yy2 = torch.clamp(yy2, max=y2[i])
- w.resize_as_(xx2)
- h.resize_as_(yy2)
- w = xx2 - xx1
- h = yy2 - yy1
- # check sizes of xx1 and xx2.. after each iteration
- w = torch.clamp(w, min=0.0)
- h = torch.clamp(h, min=0.0)
- inter = w * h
- # IoU = i / (area(a) + area(b) - i)
- rem_areas = torch.index_select(area, 0, idx) # load remaining areas)
- union = (rem_areas - inter) + area[i]
- IoU = inter / union # store result in iou
- # keep only elements with an IoU <= overlap
- idx = idx[IoU.le(overlap)]
- return keep, count
diff --git a/modules/facelib/detection/yolov5face/__init__.py b/modules/facelib/detection/yolov5face/__init__.py
deleted file mode 100644
index e69de29bb..000000000
diff --git a/modules/facelib/detection/yolov5face/face_detector.py b/modules/facelib/detection/yolov5face/face_detector.py
deleted file mode 100644
index 7c4aa52f2..000000000
--- a/modules/facelib/detection/yolov5face/face_detector.py
+++ /dev/null
@@ -1,142 +0,0 @@
-import copy
-import os
-from pathlib import Path
-
-import cv2
-import numpy as np
-import torch
-from torch import nn
-
-from ....facelib.detection.yolov5face.models.common import Conv
-from ....facelib.detection.yolov5face.models.yolo import Model
-from ....facelib.detection.yolov5face.utils.datasets import letterbox
-from ....facelib.detection.yolov5face.utils.general import (
- check_img_size,
- non_max_suppression_face,
- scale_coords,
- scale_coords_landmarks,
-)
-
-IS_HIGH_VERSION = tuple(map(int, torch.__version__.split('+')[0].split('.')[:2])) >= (1, 9, 0)
-
-
-def isListempty(inList):
- if isinstance(inList, list): # Is a list
- return all(map(isListempty, inList))
- return False # Not a list
-
-class YoloDetector:
- def __init__(
- self,
- config_name,
- min_face=10,
- target_size=None,
- device='cuda',
- ):
- """
- config_name: name of .yaml config with network configuration from models/ folder.
- min_face : minimal face size in pixels.
- target_size : target size of smaller image axis (choose lower for faster work). e.g. 480, 720, 1080.
- None for original resolution.
- """
- self._class_path = Path(__file__).parent.absolute()
- self.target_size = target_size
- self.min_face = min_face
- self.detector = Model(cfg=config_name)
- self.device = device
-
-
- def _preprocess(self, imgs):
- """
- Preprocessing image before passing through the network. Resize and conversion to torch tensor.
- """
- pp_imgs = []
- for img in imgs:
- h0, w0 = img.shape[:2] # orig hw
- if self.target_size:
- r = self.target_size / min(h0, w0) # resize image to img_size
- if r < 1:
- img = cv2.resize(img, (int(w0 * r), int(h0 * r)), interpolation=cv2.INTER_LINEAR)
-
- imgsz = check_img_size(max(img.shape[:2]), s=self.detector.stride.max()) # check img_size
- img = letterbox(img, new_shape=imgsz)[0]
- pp_imgs.append(img)
- pp_imgs = np.array(pp_imgs)
- pp_imgs = pp_imgs.transpose(0, 3, 1, 2)
- pp_imgs = torch.from_numpy(pp_imgs).to(self.device)
- pp_imgs = pp_imgs.float() # uint8 to fp16/32
- return pp_imgs / 255.0 # 0 - 255 to 0.0 - 1.0
-
- def _postprocess(self, imgs, origimgs, pred, conf_thres, iou_thres):
- """
- Postprocessing of raw pytorch model output.
- Returns:
- bboxes: list of arrays with 4 coordinates of bounding boxes with format x1,y1,x2,y2.
- points: list of arrays with coordinates of 5 facial keypoints (eyes, nose, lips corners).
- """
- bboxes = [[] for _ in range(len(origimgs))]
- landmarks = [[] for _ in range(len(origimgs))]
-
- pred = non_max_suppression_face(pred, conf_thres, iou_thres)
-
- for image_id, origimg in enumerate(origimgs):
- img_shape = origimg.shape
- image_height, image_width = img_shape[:2]
- gn = torch.tensor(img_shape)[[1, 0, 1, 0]] # normalization gain whwh
- gn_lks = torch.tensor(img_shape)[[1, 0, 1, 0, 1, 0, 1, 0, 1, 0]] # normalization gain landmarks
- det = pred[image_id].cpu()
- scale_coords(imgs[image_id].shape[1:], det[:, :4], img_shape).round()
- scale_coords_landmarks(imgs[image_id].shape[1:], det[:, 5:15], img_shape).round()
-
- for j in range(det.size()[0]):
- box = (det[j, :4].view(1, 4) / gn).view(-1).tolist()
- box = list(
- map(int, [box[0] * image_width, box[1] * image_height, box[2] * image_width, box[3] * image_height])
- )
- if box[3] - box[1] < self.min_face:
- continue
- lm = (det[j, 5:15].view(1, 10) / gn_lks).view(-1).tolist()
- lm = list(map(int, [i * image_width if j % 2 == 0 else i * image_height for j, i in enumerate(lm)]))
- lm = [lm[i : i + 2] for i in range(0, len(lm), 2)]
- bboxes[image_id].append(box)
- landmarks[image_id].append(lm)
- return bboxes, landmarks
-
- def detect_faces(self, imgs, conf_thres=0.7, iou_thres=0.5):
- """
- Get bbox coordinates and keypoints of faces on original image.
- Params:
- imgs: image or list of images to detect faces on with BGR order (convert to RGB order for inference)
- conf_thres: confidence threshold for each prediction
- iou_thres: threshold for NMS (filter of intersecting bboxes)
- Returns:
- bboxes: list of arrays with 4 coordinates of bounding boxes with format x1,y1,x2,y2.
- points: list of arrays with coordinates of 5 facial keypoints (eyes, nose, lips corners).
- """
- # Pass input images through face detector
- images = imgs if isinstance(imgs, list) else [imgs]
- images = [cv2.cvtColor(img, cv2.COLOR_BGR2RGB) for img in images]
- origimgs = copy.deepcopy(images)
-
- images = self._preprocess(images)
-
- if IS_HIGH_VERSION:
- with torch.inference_mode(): # for pytorch>=1.9
- pred = self.detector(images)[0]
- else:
- with torch.no_grad(): # for pytorch<1.9
- pred = self.detector(images)[0]
-
- bboxes, points = self._postprocess(images, origimgs, pred, conf_thres, iou_thres)
-
- # return bboxes, points
- if not isListempty(points):
- bboxes = np.array(bboxes).reshape(-1,4)
- points = np.array(points).reshape(-1,10)
- padding = bboxes[:,0].reshape(-1,1)
- return np.concatenate((bboxes, padding, points), axis=1)
- else:
- return None
-
- def __call__(self, *args):
- return self.predict(*args)
diff --git a/modules/facelib/detection/yolov5face/models/__init__.py b/modules/facelib/detection/yolov5face/models/__init__.py
deleted file mode 100644
index e69de29bb..000000000
diff --git a/modules/facelib/detection/yolov5face/models/common.py b/modules/facelib/detection/yolov5face/models/common.py
deleted file mode 100644
index c649f8185..000000000
--- a/modules/facelib/detection/yolov5face/models/common.py
+++ /dev/null
@@ -1,299 +0,0 @@
-# This file contains modules common to various models
-
-import math
-
-import numpy as np
-import torch
-from torch import nn
-
-from ....detection.yolov5face.utils.datasets import letterbox
-from ....detection.yolov5face.utils.general import (
- make_divisible,
- non_max_suppression,
- scale_coords,
- xyxy2xywh,
-)
-
-
-def autopad(k, p=None): # kernel, padding
- # Pad to 'same'
- if p is None:
- p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
- return p
-
-
-def channel_shuffle(x, groups):
- batchsize, num_channels, height, width = x.data.size()
- channels_per_group = torch.div(num_channels, groups, rounding_mode="trunc")
-
- # reshape
- x = x.view(batchsize, groups, channels_per_group, height, width)
- x = torch.transpose(x, 1, 2).contiguous()
-
- # flatten
- return x.view(batchsize, -1, height, width)
-
-
-def DWConv(c1, c2, k=1, s=1, act=True):
- # Depthwise convolution
- return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act)
-
-
-class Conv(nn.Module):
- # Standard convolution
- def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
- super().__init__()
- self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
- self.bn = nn.BatchNorm2d(c2)
- self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
-
- def forward(self, x):
- return self.act(self.bn(self.conv(x)))
-
- def fuseforward(self, x):
- return self.act(self.conv(x))
-
-
-class StemBlock(nn.Module):
- def __init__(self, c1, c2, k=3, s=2, p=None, g=1, act=True):
- super().__init__()
- self.stem_1 = Conv(c1, c2, k, s, p, g, act)
- self.stem_2a = Conv(c2, c2 // 2, 1, 1, 0)
- self.stem_2b = Conv(c2 // 2, c2, 3, 2, 1)
- self.stem_2p = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)
- self.stem_3 = Conv(c2 * 2, c2, 1, 1, 0)
-
- def forward(self, x):
- stem_1_out = self.stem_1(x)
- stem_2a_out = self.stem_2a(stem_1_out)
- stem_2b_out = self.stem_2b(stem_2a_out)
- stem_2p_out = self.stem_2p(stem_1_out)
- return self.stem_3(torch.cat((stem_2b_out, stem_2p_out), 1))
-
-
-class Bottleneck(nn.Module):
- # Standard bottleneck
- def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
- super().__init__()
- c_ = int(c2 * e) # hidden channels
- self.cv1 = Conv(c1, c_, 1, 1)
- self.cv2 = Conv(c_, c2, 3, 1, g=g)
- self.add = shortcut and c1 == c2
-
- def forward(self, x):
- return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
-
-
-class BottleneckCSP(nn.Module):
- # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
- def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
- super().__init__()
- c_ = int(c2 * e) # hidden channels
- self.cv1 = Conv(c1, c_, 1, 1)
- self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
- self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
- self.cv4 = Conv(2 * c_, c2, 1, 1)
- self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3)
- self.act = nn.LeakyReLU(0.1, inplace=True)
- self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
-
- def forward(self, x):
- y1 = self.cv3(self.m(self.cv1(x)))
- y2 = self.cv2(x)
- return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1))))
-
-
-class C3(nn.Module):
- # CSP Bottleneck with 3 convolutions
- def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
- super().__init__()
- c_ = int(c2 * e) # hidden channels
- self.cv1 = Conv(c1, c_, 1, 1)
- self.cv2 = Conv(c1, c_, 1, 1)
- self.cv3 = Conv(2 * c_, c2, 1) # act=FReLU(c2)
- self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
-
- def forward(self, x):
- return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1))
-
-
-class ShuffleV2Block(nn.Module):
- def __init__(self, inp, oup, stride):
- super().__init__()
-
- if not 1 <= stride <= 3:
- raise ValueError("illegal stride value")
- self.stride = stride
-
- branch_features = oup // 2
-
- if self.stride > 1:
- self.branch1 = nn.Sequential(
- self.depthwise_conv(inp, inp, kernel_size=3, stride=self.stride, padding=1),
- nn.BatchNorm2d(inp),
- nn.Conv2d(inp, branch_features, kernel_size=1, stride=1, padding=0, bias=False),
- nn.BatchNorm2d(branch_features),
- nn.SiLU(),
- )
- else:
- self.branch1 = nn.Sequential()
-
- self.branch2 = nn.Sequential(
- nn.Conv2d(
- inp if (self.stride > 1) else branch_features,
- branch_features,
- kernel_size=1,
- stride=1,
- padding=0,
- bias=False,
- ),
- nn.BatchNorm2d(branch_features),
- nn.SiLU(),
- self.depthwise_conv(branch_features, branch_features, kernel_size=3, stride=self.stride, padding=1),
- nn.BatchNorm2d(branch_features),
- nn.Conv2d(branch_features, branch_features, kernel_size=1, stride=1, padding=0, bias=False),
- nn.BatchNorm2d(branch_features),
- nn.SiLU(),
- )
-
- @staticmethod
- def depthwise_conv(i, o, kernel_size, stride=1, padding=0, bias=False):
- return nn.Conv2d(i, o, kernel_size, stride, padding, bias=bias, groups=i)
-
- def forward(self, x):
- if self.stride == 1:
- x1, x2 = x.chunk(2, dim=1)
- out = torch.cat((x1, self.branch2(x2)), dim=1)
- else:
- out = torch.cat((self.branch1(x), self.branch2(x)), dim=1)
- out = channel_shuffle(out, 2)
- return out
-
-
-class SPP(nn.Module):
- # Spatial pyramid pooling layer used in YOLOv3-SPP
- def __init__(self, c1, c2, k=(5, 9, 13)):
- super().__init__()
- c_ = c1 // 2 # hidden channels
- self.cv1 = Conv(c1, c_, 1, 1)
- self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
- self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
-
- def forward(self, x):
- x = self.cv1(x)
- return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
-
-
-class Focus(nn.Module):
- # Focus wh information into c-space
- def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
- super().__init__()
- self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
-
- def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2)
- return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
-
-
-class Concat(nn.Module):
- # Concatenate a list of tensors along dimension
- def __init__(self, dimension=1):
- super().__init__()
- self.d = dimension
-
- def forward(self, x):
- return torch.cat(x, self.d)
-
-
-class NMS(nn.Module):
- # Non-Maximum Suppression (NMS) module
- conf = 0.25 # confidence threshold
- iou = 0.45 # IoU threshold
- classes = None # (optional list) filter by class
-
- def forward(self, x):
- return non_max_suppression(x[0], conf_thres=self.conf, iou_thres=self.iou, classes=self.classes)
-
-
-class AutoShape(nn.Module):
- # input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
- img_size = 640 # inference size (pixels)
- conf = 0.25 # NMS confidence threshold
- iou = 0.45 # NMS IoU threshold
- classes = None # (optional list) filter by class
-
- def __init__(self, model):
- super().__init__()
- self.model = model.eval()
-
- def autoshape(self):
- print("autoShape already enabled, skipping... ") # model already converted to model.autoshape()
- return self
-
- def forward(self, imgs, size=640, augment=False, profile=False):
- # Inference from various sources. For height=720, width=1280, RGB images example inputs are:
- # OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(720,1280,3)
- # PIL: = Image.open('image.jpg') # HWC x(720,1280,3)
- # numpy: = np.zeros((720,1280,3)) # HWC
- # torch: = torch.zeros(16,3,720,1280) # BCHW
- # multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images
-
- p = next(self.model.parameters()) # for device and type
- if isinstance(imgs, torch.Tensor): # torch
- return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference
-
- # Pre-process
- n, imgs = (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs]) # number of images, list of images
- shape0, shape1 = [], [] # image and inference shapes
- for i, im in enumerate(imgs):
- im = np.array(im) # to numpy
- if im.shape[0] < 5: # image in CHW
- im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1)
- im = im[:, :, :3] if im.ndim == 3 else np.tile(im[:, :, None], 3) # enforce 3ch input
- s = im.shape[:2] # HWC
- shape0.append(s) # image shape
- g = size / max(s) # gain
- shape1.append([y * g for y in s])
- imgs[i] = im # update
- shape1 = [make_divisible(x, int(self.stride.max())) for x in np.stack(shape1, 0).max(0)] # inference shape
- x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad
- x = np.stack(x, 0) if n > 1 else x[0][None] # stack
- x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW
- x = torch.from_numpy(x).to(p.device).type_as(p) / 255.0 # uint8 to fp16/32
-
- # Inference
- with torch.no_grad():
- y = self.model(x, augment, profile)[0] # forward
- y = non_max_suppression(y, conf_thres=self.conf, iou_thres=self.iou, classes=self.classes) # NMS
-
- # Post-process
- for i in range(n):
- scale_coords(shape1, y[i][:, :4], shape0[i])
-
- return Detections(imgs, y, self.names)
-
-
-class Detections:
- # detections class for YOLOv5 inference results
- def __init__(self, imgs, pred, names=None):
- super().__init__()
- d = pred[0].device # device
- gn = [torch.tensor([*(im.shape[i] for i in [1, 0, 1, 0]), 1.0, 1.0], device=d) for im in imgs] # normalizations
- self.imgs = imgs # list of images as numpy arrays
- self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls)
- self.names = names # class names
- self.xyxy = pred # xyxy pixels
- self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels
- self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized
- self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized
- self.n = len(self.pred)
-
- def __len__(self):
- return self.n
-
- def tolist(self):
- # return a list of Detections objects, i.e. 'for result in results.tolist():'
- x = [Detections([self.imgs[i]], [self.pred[i]], self.names) for i in range(self.n)]
- for d in x:
- for k in ["imgs", "pred", "xyxy", "xyxyn", "xywh", "xywhn"]:
- setattr(d, k, getattr(d, k)[0]) # pop out of list
- return x
diff --git a/modules/facelib/detection/yolov5face/models/experimental.py b/modules/facelib/detection/yolov5face/models/experimental.py
deleted file mode 100644
index cd63bc206..000000000
--- a/modules/facelib/detection/yolov5face/models/experimental.py
+++ /dev/null
@@ -1,45 +0,0 @@
-# # This file contains experimental modules
-
-import numpy as np
-import torch
-from torch import nn
-
-from ....detection.yolov5face.models.common import Conv
-
-
-class CrossConv(nn.Module):
- # Cross Convolution Downsample
- def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False):
- # ch_in, ch_out, kernel, stride, groups, expansion, shortcut
- super().__init__()
- c_ = int(c2 * e) # hidden channels
- self.cv1 = Conv(c1, c_, (1, k), (1, s))
- self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g)
- self.add = shortcut and c1 == c2
-
- def forward(self, x):
- return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
-
-
-class MixConv2d(nn.Module):
- # Mixed Depthwise Conv https://arxiv.org/abs/1907.09595
- def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True):
- super().__init__()
- groups = len(k)
- if equal_ch: # equal c_ per group
- i = torch.linspace(0, groups - 1e-6, c2).floor() # c2 indices
- c_ = [(i == g).sum() for g in range(groups)] # intermediate channels
- else: # equal weight.numel() per group
- b = [c2] + [0] * groups
- a = np.eye(groups + 1, groups, k=-1)
- a -= np.roll(a, 1, axis=1)
- a *= np.array(k) ** 2
- a[0] = 1
- c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b
-
- self.m = nn.ModuleList([nn.Conv2d(c1, int(c_[g]), k[g], s, k[g] // 2, bias=False) for g in range(groups)])
- self.bn = nn.BatchNorm2d(c2)
- self.act = nn.LeakyReLU(0.1, inplace=True)
-
- def forward(self, x):
- return x + self.act(self.bn(torch.cat([m(x) for m in self.m], 1)))
diff --git a/modules/facelib/detection/yolov5face/models/yolo.py b/modules/facelib/detection/yolov5face/models/yolo.py
deleted file mode 100644
index 19fd70663..000000000
--- a/modules/facelib/detection/yolov5face/models/yolo.py
+++ /dev/null
@@ -1,235 +0,0 @@
-import math
-from copy import deepcopy
-from pathlib import Path
-
-import torch
-import yaml # for torch hub
-from torch import nn
-
-from ....detection.yolov5face.models.common import (
- C3,
- NMS,
- SPP,
- AutoShape,
- Bottleneck,
- BottleneckCSP,
- Concat,
- Conv,
- DWConv,
- Focus,
- ShuffleV2Block,
- StemBlock,
-)
-from ....detection.yolov5face.models.experimental import CrossConv, MixConv2d
-from ....detection.yolov5face.utils.autoanchor import check_anchor_order
-from ....detection.yolov5face.utils.general import make_divisible
-from ....detection.yolov5face.utils.torch_utils import copy_attr, fuse_conv_and_bn
-
-
-class Detect(nn.Module):
- stride = None # strides computed during build
- export = False # onnx export
-
- def __init__(self, nc=80, anchors=(), ch=()): # detection layer
- super().__init__()
- self.nc = nc # number of classes
- self.no = nc + 5 + 10 # number of outputs per anchor
-
- self.nl = len(anchors) # number of detection layers
- self.na = len(anchors[0]) // 2 # number of anchors
- self.grid = [torch.zeros(1)] * self.nl # init grid
- a = torch.tensor(anchors).float().view(self.nl, -1, 2)
- self.register_buffer("anchors", a) # shape(nl,na,2)
- self.register_buffer("anchor_grid", a.clone().view(self.nl, 1, -1, 1, 1, 2)) # shape(nl,1,na,1,1,2)
- self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv
-
- def forward(self, x):
- z = [] # inference output
- if self.export:
- for i in range(self.nl):
- x[i] = self.m[i](x[i])
- return x
- for i in range(self.nl):
- x[i] = self.m[i](x[i]) # conv
- bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
- x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
-
- if not self.training: # inference
- if self.grid[i].shape[2:4] != x[i].shape[2:4]:
- self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
-
- y = torch.full_like(x[i], 0)
- y[..., [0, 1, 2, 3, 4, 15]] = x[i][..., [0, 1, 2, 3, 4, 15]].sigmoid()
- y[..., 5:15] = x[i][..., 5:15]
-
- y[..., 0:2] = (y[..., 0:2] * 2.0 - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy
- y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
-
- y[..., 5:7] = (
- y[..., 5:7] * self.anchor_grid[i] + self.grid[i].to(x[i].device) * self.stride[i]
- ) # landmark x1 y1
- y[..., 7:9] = (
- y[..., 7:9] * self.anchor_grid[i] + self.grid[i].to(x[i].device) * self.stride[i]
- ) # landmark x2 y2
- y[..., 9:11] = (
- y[..., 9:11] * self.anchor_grid[i] + self.grid[i].to(x[i].device) * self.stride[i]
- ) # landmark x3 y3
- y[..., 11:13] = (
- y[..., 11:13] * self.anchor_grid[i] + self.grid[i].to(x[i].device) * self.stride[i]
- ) # landmark x4 y4
- y[..., 13:15] = (
- y[..., 13:15] * self.anchor_grid[i] + self.grid[i].to(x[i].device) * self.stride[i]
- ) # landmark x5 y5
-
- z.append(y.view(bs, -1, self.no))
-
- return x if self.training else (torch.cat(z, 1), x)
-
- @staticmethod
- def _make_grid(nx=20, ny=20):
- # yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)], indexing="ij") # for pytorch>=1.10
- yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
- return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()
-
-
-class Model(nn.Module):
- def __init__(self, cfg="yolov5s.yaml", ch=3, nc=None): # model, input channels, number of classes
- super().__init__()
- self.yaml_file = Path(cfg).name
- with Path(cfg).open(encoding="utf8") as f:
- self.yaml = yaml.safe_load(f) # model dict
-
- # Define model
- ch = self.yaml["ch"] = self.yaml.get("ch", ch) # input channels
- if nc and nc != self.yaml["nc"]:
- self.yaml["nc"] = nc # override yaml value
-
- self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist
- self.names = [str(i) for i in range(self.yaml["nc"])] # default names
-
- # Build strides, anchors
- m = self.model[-1] # Detect()
- if isinstance(m, Detect):
- s = 128 # 2x min stride
- m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))]) # forward
- m.anchors /= m.stride.view(-1, 1, 1)
- check_anchor_order(m)
- self.stride = m.stride
- self._initialize_biases() # only run once
-
- def forward(self, x):
- return self.forward_once(x) # single-scale inference, train
-
- def forward_once(self, x):
- y = [] # outputs
- for m in self.model:
- if m.f != -1: # if not from previous layer
- x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
-
- x = m(x) # run
- y.append(x if m.i in self.save else None) # save output
-
- return x
-
- def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency
- # https://arxiv.org/abs/1708.02002 section 3.3
- m = self.model[-1] # Detect() module
- for mi, s in zip(m.m, m.stride): # from
- b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85)
- b.data[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image)
- b.data[:, 5:] += math.log(0.6 / (m.nc - 0.99)) if cf is None else torch.log(cf / cf.sum()) # cls
- mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
-
- def _print_biases(self):
- m = self.model[-1] # Detect() module
- for mi in m.m: # from
- b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85)
- print(("%6g Conv2d.bias:" + "%10.3g" * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean()))
-
- def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers
- print("Fusing layers... ")
- for m in self.model.modules():
- if isinstance(m, Conv) and hasattr(m, "bn"):
- m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
- delattr(m, "bn") # remove batchnorm
- m.forward = m.fuseforward # update forward
- elif type(m) is nn.Upsample:
- m.recompute_scale_factor = None # torch 1.11.0 compatibility
- return self
-
- def nms(self, mode=True): # add or remove NMS module
- present = isinstance(self.model[-1], NMS) # last layer is NMS
- if mode and not present:
- print("Adding NMS... ")
- m = NMS() # module
- m.f = -1 # from
- m.i = self.model[-1].i + 1 # index
- self.model.add_module(name=str(m.i), module=m) # add
- self.eval()
- elif not mode and present:
- print("Removing NMS... ")
- self.model = self.model[:-1] # remove
- return self
-
- def autoshape(self): # add autoShape module
- print("Adding autoShape... ")
- m = AutoShape(self) # wrap model
- copy_attr(m, self, include=("yaml", "nc", "hyp", "names", "stride"), exclude=()) # copy attributes
- return m
-
-
-def parse_model(d, ch): # model_dict, input_channels(3)
- anchors, nc, gd, gw = d["anchors"], d["nc"], d["depth_multiple"], d["width_multiple"]
- na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors
- no = na * (nc + 5) # number of outputs = anchors * (classes + 5)
-
- layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out
- for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]): # from, number, module, args
- m = eval(m) if isinstance(m, str) else m # eval strings
- for j, a in enumerate(args):
- try:
- args[j] = eval(a) if isinstance(a, str) else a # eval strings
- except Exception:
- pass
-
- n = max(round(n * gd), 1) if n > 1 else n # depth gain
- if m in [
- Conv,
- Bottleneck,
- SPP,
- DWConv,
- MixConv2d,
- Focus,
- CrossConv,
- BottleneckCSP,
- C3,
- ShuffleV2Block,
- StemBlock,
- ]:
- c1, c2 = ch[f], args[0]
-
- c2 = make_divisible(c2 * gw, 8) if c2 != no else c2
-
- args = [c1, c2, *args[1:]]
- if m in [BottleneckCSP, C3]:
- args.insert(2, n)
- n = 1
- elif m is nn.BatchNorm2d:
- args = [ch[f]]
- elif m is Concat:
- c2 = sum(ch[-1 if x == -1 else x + 1] for x in f)
- elif m is Detect:
- args.append([ch[x + 1] for x in f])
- if isinstance(args[1], int): # number of anchors
- args[1] = [list(range(args[1] * 2))] * len(f)
- else:
- c2 = ch[f]
-
- m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args) # module
- t = str(m)[8:-2].replace("__main__.", "") # module type
- np = sum(x.numel() for x in m_.parameters()) # number params
- m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params
- save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
- layers.append(m_)
- ch.append(c2)
- return nn.Sequential(*layers), sorted(save)
diff --git a/modules/facelib/detection/yolov5face/models/yolov5l.yaml b/modules/facelib/detection/yolov5face/models/yolov5l.yaml
deleted file mode 100644
index 5c8302517..000000000
--- a/modules/facelib/detection/yolov5face/models/yolov5l.yaml
+++ /dev/null
@@ -1,47 +0,0 @@
-# parameters
-nc: 1 # number of classes
-depth_multiple: 1.0 # model depth multiple
-width_multiple: 1.0 # layer channel multiple
-
-# anchors
-anchors:
- - [4,5, 8,10, 13,16] # P3/8
- - [23,29, 43,55, 73,105] # P4/16
- - [146,217, 231,300, 335,433] # P5/32
-
-# YOLOv5 backbone
-backbone:
- # [from, number, module, args]
- [[-1, 1, StemBlock, [64, 3, 2]], # 0-P1/2
- [-1, 3, C3, [128]],
- [-1, 1, Conv, [256, 3, 2]], # 2-P3/8
- [-1, 9, C3, [256]],
- [-1, 1, Conv, [512, 3, 2]], # 4-P4/16
- [-1, 9, C3, [512]],
- [-1, 1, Conv, [1024, 3, 2]], # 6-P5/32
- [-1, 1, SPP, [1024, [3,5,7]]],
- [-1, 3, C3, [1024, False]], # 8
- ]
-
-# YOLOv5 head
-head:
- [[-1, 1, Conv, [512, 1, 1]],
- [-1, 1, nn.Upsample, [None, 2, 'nearest']],
- [[-1, 5], 1, Concat, [1]], # cat backbone P4
- [-1, 3, C3, [512, False]], # 12
-
- [-1, 1, Conv, [256, 1, 1]],
- [-1, 1, nn.Upsample, [None, 2, 'nearest']],
- [[-1, 3], 1, Concat, [1]], # cat backbone P3
- [-1, 3, C3, [256, False]], # 16 (P3/8-small)
-
- [-1, 1, Conv, [256, 3, 2]],
- [[-1, 13], 1, Concat, [1]], # cat head P4
- [-1, 3, C3, [512, False]], # 19 (P4/16-medium)
-
- [-1, 1, Conv, [512, 3, 2]],
- [[-1, 9], 1, Concat, [1]], # cat head P5
- [-1, 3, C3, [1024, False]], # 22 (P5/32-large)
-
- [[16, 19, 22], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
- ]
diff --git a/modules/facelib/detection/yolov5face/models/yolov5n.yaml b/modules/facelib/detection/yolov5face/models/yolov5n.yaml
deleted file mode 100644
index caba6bed6..000000000
--- a/modules/facelib/detection/yolov5face/models/yolov5n.yaml
+++ /dev/null
@@ -1,45 +0,0 @@
-# parameters
-nc: 1 # number of classes
-depth_multiple: 1.0 # model depth multiple
-width_multiple: 1.0 # layer channel multiple
-
-# anchors
-anchors:
- - [4,5, 8,10, 13,16] # P3/8
- - [23,29, 43,55, 73,105] # P4/16
- - [146,217, 231,300, 335,433] # P5/32
-
-# YOLOv5 backbone
-backbone:
- # [from, number, module, args]
- [[-1, 1, StemBlock, [32, 3, 2]], # 0-P2/4
- [-1, 1, ShuffleV2Block, [128, 2]], # 1-P3/8
- [-1, 3, ShuffleV2Block, [128, 1]], # 2
- [-1, 1, ShuffleV2Block, [256, 2]], # 3-P4/16
- [-1, 7, ShuffleV2Block, [256, 1]], # 4
- [-1, 1, ShuffleV2Block, [512, 2]], # 5-P5/32
- [-1, 3, ShuffleV2Block, [512, 1]], # 6
- ]
-
-# YOLOv5 head
-head:
- [[-1, 1, Conv, [128, 1, 1]],
- [-1, 1, nn.Upsample, [None, 2, 'nearest']],
- [[-1, 4], 1, Concat, [1]], # cat backbone P4
- [-1, 1, C3, [128, False]], # 10
-
- [-1, 1, Conv, [128, 1, 1]],
- [-1, 1, nn.Upsample, [None, 2, 'nearest']],
- [[-1, 2], 1, Concat, [1]], # cat backbone P3
- [-1, 1, C3, [128, False]], # 14 (P3/8-small)
-
- [-1, 1, Conv, [128, 3, 2]],
- [[-1, 11], 1, Concat, [1]], # cat head P4
- [-1, 1, C3, [128, False]], # 17 (P4/16-medium)
-
- [-1, 1, Conv, [128, 3, 2]],
- [[-1, 7], 1, Concat, [1]], # cat head P5
- [-1, 1, C3, [128, False]], # 20 (P5/32-large)
-
- [[14, 17, 20], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
- ]
diff --git a/modules/facelib/detection/yolov5face/utils/__init__.py b/modules/facelib/detection/yolov5face/utils/__init__.py
deleted file mode 100644
index e69de29bb..000000000
diff --git a/modules/facelib/detection/yolov5face/utils/autoanchor.py b/modules/facelib/detection/yolov5face/utils/autoanchor.py
deleted file mode 100644
index a4eba3e94..000000000
--- a/modules/facelib/detection/yolov5face/utils/autoanchor.py
+++ /dev/null
@@ -1,12 +0,0 @@
-# Auto-anchor utils
-
-
-def check_anchor_order(m):
- # Check anchor order against stride order for YOLOv5 Detect() module m, and correct if necessary
- a = m.anchor_grid.prod(-1).view(-1) # anchor area
- da = a[-1] - a[0] # delta a
- ds = m.stride[-1] - m.stride[0] # delta s
- if da.sign() != ds.sign(): # same order
- print("Reversing anchor order")
- m.anchors[:] = m.anchors.flip(0)
- m.anchor_grid[:] = m.anchor_grid.flip(0)
diff --git a/modules/facelib/detection/yolov5face/utils/datasets.py b/modules/facelib/detection/yolov5face/utils/datasets.py
deleted file mode 100644
index e672b136f..000000000
--- a/modules/facelib/detection/yolov5face/utils/datasets.py
+++ /dev/null
@@ -1,35 +0,0 @@
-import cv2
-import numpy as np
-
-
-def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scale_fill=False, scaleup=True):
- # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232
- shape = img.shape[:2] # current shape [height, width]
- if isinstance(new_shape, int):
- new_shape = (new_shape, new_shape)
-
- # Scale ratio (new / old)
- r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
- if not scaleup: # only scale down, do not scale up (for better test mAP)
- r = min(r, 1.0)
-
- # Compute padding
- ratio = r, r # width, height ratios
- new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
- dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
- if auto: # minimum rectangle
- dw, dh = np.mod(dw, 64), np.mod(dh, 64) # wh padding
- elif scale_fill: # stretch
- dw, dh = 0.0, 0.0
- new_unpad = (new_shape[1], new_shape[0])
- ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios
-
- dw /= 2 # divide padding into 2 sides
- dh /= 2
-
- if shape[::-1] != new_unpad: # resize
- img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
- top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
- left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
- img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
- return img, ratio, (dw, dh)
diff --git a/modules/facelib/detection/yolov5face/utils/extract_ckpt.py b/modules/facelib/detection/yolov5face/utils/extract_ckpt.py
deleted file mode 100644
index 413719ecc..000000000
--- a/modules/facelib/detection/yolov5face/utils/extract_ckpt.py
+++ /dev/null
@@ -1,5 +0,0 @@
-import torch
-import sys
-sys.path.insert(0,'./facelib/detection/yolov5face')
-model = torch.load('facelib/detection/yolov5face/yolov5n-face.pt', map_location='cpu')['model']
-torch.save(model.state_dict(),'weights/facelib/yolov5n-face.pth')
diff --git a/modules/facelib/detection/yolov5face/utils/general.py b/modules/facelib/detection/yolov5face/utils/general.py
deleted file mode 100644
index 1c8e14f56..000000000
--- a/modules/facelib/detection/yolov5face/utils/general.py
+++ /dev/null
@@ -1,271 +0,0 @@
-import math
-import time
-
-import numpy as np
-import torch
-import torchvision
-
-
-def check_img_size(img_size, s=32):
- # Verify img_size is a multiple of stride s
- new_size = make_divisible(img_size, int(s)) # ceil gs-multiple
- # if new_size != img_size:
- # print(f"WARNING: --img-size {img_size:g} must be multiple of max stride {s:g}, updating to {new_size:g}")
- return new_size
-
-
-def make_divisible(x, divisor):
- # Returns x evenly divisible by divisor
- return math.ceil(x / divisor) * divisor
-
-
-def xyxy2xywh(x):
- # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
- y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
- y[:, 0] = (x[:, 0] + x[:, 2]) / 2 # x center
- y[:, 1] = (x[:, 1] + x[:, 3]) / 2 # y center
- y[:, 2] = x[:, 2] - x[:, 0] # width
- y[:, 3] = x[:, 3] - x[:, 1] # height
- return y
-
-
-def xywh2xyxy(x):
- # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
- y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
- y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x
- y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y
- y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x
- y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y
- return y
-
-
-def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
- # Rescale coords (xyxy) from img1_shape to img0_shape
- if ratio_pad is None: # calculate from img0_shape
- gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
- pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
- else:
- gain = ratio_pad[0][0]
- pad = ratio_pad[1]
-
- coords[:, [0, 2]] -= pad[0] # x padding
- coords[:, [1, 3]] -= pad[1] # y padding
- coords[:, :4] /= gain
- clip_coords(coords, img0_shape)
- return coords
-
-
-def clip_coords(boxes, img_shape):
- # Clip bounding xyxy bounding boxes to image shape (height, width)
- boxes[:, 0].clamp_(0, img_shape[1]) # x1
- boxes[:, 1].clamp_(0, img_shape[0]) # y1
- boxes[:, 2].clamp_(0, img_shape[1]) # x2
- boxes[:, 3].clamp_(0, img_shape[0]) # y2
-
-
-def box_iou(box1, box2):
- # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
- """
- Return intersection-over-union (Jaccard index) of boxes.
- Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
- Arguments:
- box1 (Tensor[N, 4])
- box2 (Tensor[M, 4])
- Returns:
- iou (Tensor[N, M]): the NxM matrix containing the pairwise
- IoU values for every element in boxes1 and boxes2
- """
-
- def box_area(box):
- return (box[2] - box[0]) * (box[3] - box[1])
-
- area1 = box_area(box1.T)
- area2 = box_area(box2.T)
-
- inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
- return inter / (area1[:, None] + area2 - inter)
-
-
-def non_max_suppression_face(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, labels=()):
- """Performs Non-Maximum Suppression (NMS) on inference results
- Returns:
- detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
- """
-
- nc = prediction.shape[2] - 15 # number of classes
- xc = prediction[..., 4] > conf_thres # candidates
-
- # Settings
- # (pixels) maximum box width and height
- max_wh = 4096
- time_limit = 10.0 # seconds to quit after
- redundant = True # require redundant detections
- multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img)
- merge = False # use merge-NMS
-
- t = time.time()
- output = [torch.zeros((0, 16), device=prediction.device)] * prediction.shape[0]
- for xi, x in enumerate(prediction): # image index, image inference
- # Apply constraints
- x = x[xc[xi]] # confidence
-
- # Cat apriori labels if autolabelling
- if labels and len(labels[xi]):
- label = labels[xi]
- v = torch.zeros((len(label), nc + 15), device=x.device)
- v[:, :4] = label[:, 1:5] # box
- v[:, 4] = 1.0 # conf
- v[range(len(label)), label[:, 0].long() + 15] = 1.0 # cls
- x = torch.cat((x, v), 0)
-
- # If none remain process next image
- if not x.shape[0]:
- continue
-
- # Compute conf
- x[:, 15:] *= x[:, 4:5] # conf = obj_conf * cls_conf
-
- # Box (center x, center y, width, height) to (x1, y1, x2, y2)
- box = xywh2xyxy(x[:, :4])
-
- # Detections matrix nx6 (xyxy, conf, landmarks, cls)
- if multi_label:
- i, j = (x[:, 15:] > conf_thres).nonzero(as_tuple=False).T
- x = torch.cat((box[i], x[i, j + 15, None], x[:, 5:15], j[:, None].float()), 1)
- else: # best class only
- conf, j = x[:, 15:].max(1, keepdim=True)
- x = torch.cat((box, conf, x[:, 5:15], j.float()), 1)[conf.view(-1) > conf_thres]
-
- # Filter by class
- if classes is not None:
- x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
-
- # If none remain process next image
- n = x.shape[0] # number of boxes
- if not n:
- continue
-
- # Batched NMS
- c = x[:, 15:16] * (0 if agnostic else max_wh) # classes
- boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores
- i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS
-
- if merge and (1 < n < 3e3): # Merge NMS (boxes merged using weighted mean)
- # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
- iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix
- weights = iou * scores[None] # box weights
- x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes
- if redundant:
- i = i[iou.sum(1) > 1] # require redundancy
-
- output[xi] = x[i]
- if (time.time() - t) > time_limit:
- break # time limit exceeded
-
- return output
-
-
-def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, labels=()):
- """Performs Non-Maximum Suppression (NMS) on inference results
-
- Returns:
- detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
- """
-
- nc = prediction.shape[2] - 5 # number of classes
- xc = prediction[..., 4] > conf_thres # candidates
-
- # Settings
- # (pixels) maximum box width and height
- max_wh = 4096
- time_limit = 10.0 # seconds to quit after
- redundant = True # require redundant detections
- multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img)
- merge = False # use merge-NMS
-
- t = time.time()
- output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
- for xi, x in enumerate(prediction): # image index, image inference
- x = x[xc[xi]] # confidence
-
- # Cat apriori labels if autolabelling
- if labels and len(labels[xi]):
- label_id = labels[xi]
- v = torch.zeros((len(label_id), nc + 5), device=x.device)
- v[:, :4] = label_id[:, 1:5] # box
- v[:, 4] = 1.0 # conf
- v[range(len(label_id)), label_id[:, 0].long() + 5] = 1.0 # cls
- x = torch.cat((x, v), 0)
-
- # If none remain process next image
- if not x.shape[0]:
- continue
-
- # Compute conf
- x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf
-
- # Box (center x, center y, width, height) to (x1, y1, x2, y2)
- box = xywh2xyxy(x[:, :4])
-
- # Detections matrix nx6 (xyxy, conf, cls)
- if multi_label:
- i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
- x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
- else: # best class only
- conf, j = x[:, 5:].max(1, keepdim=True)
- x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
-
- # Filter by class
- if classes is not None:
- x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
-
- # Check shape
- n = x.shape[0] # number of boxes
- if not n: # no boxes
- continue
-
- x = x[x[:, 4].argsort(descending=True)] # sort by confidence
-
- # Batched NMS
- c = x[:, 5:6] * (0 if agnostic else max_wh) # classes
- boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores
- i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS
- if merge and (1 < n < 3e3): # Merge NMS (boxes merged using weighted mean)
- # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
- iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix
- weights = iou * scores[None] # box weights
- x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes
- if redundant:
- i = i[iou.sum(1) > 1] # require redundancy
-
- output[xi] = x[i]
- if (time.time() - t) > time_limit:
- print(f"WARNING: NMS time limit {time_limit}s exceeded")
- break # time limit exceeded
-
- return output
-
-
-def scale_coords_landmarks(img1_shape, coords, img0_shape, ratio_pad=None):
- # Rescale coords (xyxy) from img1_shape to img0_shape
- if ratio_pad is None: # calculate from img0_shape
- gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
- pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
- else:
- gain = ratio_pad[0][0]
- pad = ratio_pad[1]
-
- coords[:, [0, 2, 4, 6, 8]] -= pad[0] # x padding
- coords[:, [1, 3, 5, 7, 9]] -= pad[1] # y padding
- coords[:, :10] /= gain
- coords[:, 0].clamp_(0, img0_shape[1]) # x1
- coords[:, 1].clamp_(0, img0_shape[0]) # y1
- coords[:, 2].clamp_(0, img0_shape[1]) # x2
- coords[:, 3].clamp_(0, img0_shape[0]) # y2
- coords[:, 4].clamp_(0, img0_shape[1]) # x3
- coords[:, 5].clamp_(0, img0_shape[0]) # y3
- coords[:, 6].clamp_(0, img0_shape[1]) # x4
- coords[:, 7].clamp_(0, img0_shape[0]) # y4
- coords[:, 8].clamp_(0, img0_shape[1]) # x5
- coords[:, 9].clamp_(0, img0_shape[0]) # y5
- return coords
diff --git a/modules/facelib/detection/yolov5face/utils/torch_utils.py b/modules/facelib/detection/yolov5face/utils/torch_utils.py
deleted file mode 100644
index af2d06587..000000000
--- a/modules/facelib/detection/yolov5face/utils/torch_utils.py
+++ /dev/null
@@ -1,40 +0,0 @@
-import torch
-from torch import nn
-
-
-def fuse_conv_and_bn(conv, bn):
- # Fuse convolution and batchnorm layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/
- fusedconv = (
- nn.Conv2d(
- conv.in_channels,
- conv.out_channels,
- kernel_size=conv.kernel_size,
- stride=conv.stride,
- padding=conv.padding,
- groups=conv.groups,
- bias=True,
- )
- .requires_grad_(False)
- .to(conv.weight.device)
- )
-
- # prepare filters
- w_conv = conv.weight.clone().view(conv.out_channels, -1)
- w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
- fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size()))
-
- # prepare spatial bias
- b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias
- b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
- fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)
-
- return fusedconv
-
-
-def copy_attr(a, b, include=(), exclude=()):
- # Copy attributes from b to a, options to only include [...] and to exclude [...]
- for k, v in b.__dict__.items():
- if (include and k not in include) or k.startswith("_") or k in exclude:
- continue
-
- setattr(a, k, v)
diff --git a/modules/facelib/parsing/__init__.py b/modules/facelib/parsing/__init__.py
deleted file mode 100644
index 9017dd2b0..000000000
--- a/modules/facelib/parsing/__init__.py
+++ /dev/null
@@ -1,28 +0,0 @@
-import os
-import torch
-
-from ..utils import load_file_from_url
-from .bisenet import BiSeNet
-from .parsenet import ParseNet
-from modules import paths
-
-
-model_dir = os.path.join(paths.models_path, 'Codeformer')
-
-
-def init_parsing_model(model_name='bisenet', half=False, device='cuda'):
- if model_name == 'bisenet':
- model = BiSeNet(num_class=19)
- model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_bisenet.pth'
- elif model_name == 'parsenet':
- model = ParseNet(in_size=512, out_size=512, parsing_ch=19)
- model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth'
- else:
- raise NotImplementedError(f'{model_name} is not implemented.')
-
- model_path = load_file_from_url(url=model_url, model_dir=model_dir, progress=True, file_name=None)
- load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
- model.load_state_dict(load_net, strict=True)
- model.eval()
- model = model.to(device)
- return model
diff --git a/modules/facelib/parsing/bisenet.py b/modules/facelib/parsing/bisenet.py
deleted file mode 100644
index 3898cab76..000000000
--- a/modules/facelib/parsing/bisenet.py
+++ /dev/null
@@ -1,140 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from .resnet import ResNet18
-
-
-class ConvBNReLU(nn.Module):
-
- def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1):
- super(ConvBNReLU, self).__init__()
- self.conv = nn.Conv2d(in_chan, out_chan, kernel_size=ks, stride=stride, padding=padding, bias=False)
- self.bn = nn.BatchNorm2d(out_chan)
-
- def forward(self, x):
- x = self.conv(x)
- x = F.relu(self.bn(x))
- return x
-
-
-class BiSeNetOutput(nn.Module):
-
- def __init__(self, in_chan, mid_chan, num_class):
- super(BiSeNetOutput, self).__init__()
- self.conv = ConvBNReLU(in_chan, mid_chan, ks=3, stride=1, padding=1)
- self.conv_out = nn.Conv2d(mid_chan, num_class, kernel_size=1, bias=False)
-
- def forward(self, x):
- feat = self.conv(x)
- out = self.conv_out(feat)
- return out, feat
-
-
-class AttentionRefinementModule(nn.Module):
-
- def __init__(self, in_chan, out_chan):
- super(AttentionRefinementModule, self).__init__()
- self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1)
- self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size=1, bias=False)
- self.bn_atten = nn.BatchNorm2d(out_chan)
- self.sigmoid_atten = nn.Sigmoid()
-
- def forward(self, x):
- feat = self.conv(x)
- atten = F.avg_pool2d(feat, feat.size()[2:])
- atten = self.conv_atten(atten)
- atten = self.bn_atten(atten)
- atten = self.sigmoid_atten(atten)
- out = torch.mul(feat, atten)
- return out
-
-
-class ContextPath(nn.Module):
-
- def __init__(self):
- super(ContextPath, self).__init__()
- self.resnet = ResNet18()
- self.arm16 = AttentionRefinementModule(256, 128)
- self.arm32 = AttentionRefinementModule(512, 128)
- self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
- self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
- self.conv_avg = ConvBNReLU(512, 128, ks=1, stride=1, padding=0)
-
- def forward(self, x):
- feat8, feat16, feat32 = self.resnet(x)
- h8, w8 = feat8.size()[2:]
- h16, w16 = feat16.size()[2:]
- h32, w32 = feat32.size()[2:]
-
- avg = F.avg_pool2d(feat32, feat32.size()[2:])
- avg = self.conv_avg(avg)
- avg_up = F.interpolate(avg, (h32, w32), mode='nearest')
-
- feat32_arm = self.arm32(feat32)
- feat32_sum = feat32_arm + avg_up
- feat32_up = F.interpolate(feat32_sum, (h16, w16), mode='nearest')
- feat32_up = self.conv_head32(feat32_up)
-
- feat16_arm = self.arm16(feat16)
- feat16_sum = feat16_arm + feat32_up
- feat16_up = F.interpolate(feat16_sum, (h8, w8), mode='nearest')
- feat16_up = self.conv_head16(feat16_up)
-
- return feat8, feat16_up, feat32_up # x8, x8, x16
-
-
-class FeatureFusionModule(nn.Module):
-
- def __init__(self, in_chan, out_chan):
- super(FeatureFusionModule, self).__init__()
- self.convblk = ConvBNReLU(in_chan, out_chan, ks=1, stride=1, padding=0)
- self.conv1 = nn.Conv2d(out_chan, out_chan // 4, kernel_size=1, stride=1, padding=0, bias=False)
- self.conv2 = nn.Conv2d(out_chan // 4, out_chan, kernel_size=1, stride=1, padding=0, bias=False)
- self.relu = nn.ReLU(inplace=True)
- self.sigmoid = nn.Sigmoid()
-
- def forward(self, fsp, fcp):
- fcat = torch.cat([fsp, fcp], dim=1)
- feat = self.convblk(fcat)
- atten = F.avg_pool2d(feat, feat.size()[2:])
- atten = self.conv1(atten)
- atten = self.relu(atten)
- atten = self.conv2(atten)
- atten = self.sigmoid(atten)
- feat_atten = torch.mul(feat, atten)
- feat_out = feat_atten + feat
- return feat_out
-
-
-class BiSeNet(nn.Module):
-
- def __init__(self, num_class):
- super(BiSeNet, self).__init__()
- self.cp = ContextPath()
- self.ffm = FeatureFusionModule(256, 256)
- self.conv_out = BiSeNetOutput(256, 256, num_class)
- self.conv_out16 = BiSeNetOutput(128, 64, num_class)
- self.conv_out32 = BiSeNetOutput(128, 64, num_class)
-
- def forward(self, x, return_feat=False):
- h, w = x.size()[2:]
- feat_res8, feat_cp8, feat_cp16 = self.cp(x) # return res3b1 feature
- feat_sp = feat_res8 # replace spatial path feature with res3b1 feature
- feat_fuse = self.ffm(feat_sp, feat_cp8)
-
- out, feat = self.conv_out(feat_fuse)
- out16, feat16 = self.conv_out16(feat_cp8)
- out32, feat32 = self.conv_out32(feat_cp16)
-
- out = F.interpolate(out, (h, w), mode='bilinear', align_corners=True)
- out16 = F.interpolate(out16, (h, w), mode='bilinear', align_corners=True)
- out32 = F.interpolate(out32, (h, w), mode='bilinear', align_corners=True)
-
- if return_feat:
- feat = F.interpolate(feat, (h, w), mode='bilinear', align_corners=True)
- feat16 = F.interpolate(feat16, (h, w), mode='bilinear', align_corners=True)
- feat32 = F.interpolate(feat32, (h, w), mode='bilinear', align_corners=True)
- return out, out16, out32, feat, feat16, feat32
- else:
- return out, out16, out32
diff --git a/modules/facelib/parsing/parsenet.py b/modules/facelib/parsing/parsenet.py
deleted file mode 100644
index e178ebe43..000000000
--- a/modules/facelib/parsing/parsenet.py
+++ /dev/null
@@ -1,194 +0,0 @@
-"""Modified from https://github.com/chaofengc/PSFRGAN
-"""
-import numpy as np
-import torch.nn as nn
-from torch.nn import functional as F
-
-
-class NormLayer(nn.Module):
- """Normalization Layers.
-
- Args:
- channels: input channels, for batch norm and instance norm.
- input_size: input shape without batch size, for layer norm.
- """
-
- def __init__(self, channels, normalize_shape=None, norm_type='bn'):
- super(NormLayer, self).__init__()
- norm_type = norm_type.lower()
- self.norm_type = norm_type
- if norm_type == 'bn':
- self.norm = nn.BatchNorm2d(channels, affine=True)
- elif norm_type == 'in':
- self.norm = nn.InstanceNorm2d(channels, affine=False)
- elif norm_type == 'gn':
- self.norm = nn.GroupNorm(32, channels, affine=True)
- elif norm_type == 'pixel':
- self.norm = lambda x: F.normalize(x, p=2, dim=1)
- elif norm_type == 'layer':
- self.norm = nn.LayerNorm(normalize_shape)
- elif norm_type == 'none':
- self.norm = lambda x: x * 1.0
- else:
- assert 1 == 0, f'Norm type {norm_type} not support.'
-
- def forward(self, x, ref=None):
- if self.norm_type == 'spade':
- return self.norm(x, ref)
- else:
- return self.norm(x)
-
-
-class ReluLayer(nn.Module):
- """Relu Layer.
-
- Args:
- relu type: type of relu layer, candidates are
- - ReLU
- - LeakyReLU: default relu slope 0.2
- - PRelu
- - SELU
- - none: direct pass
- """
-
- def __init__(self, channels, relu_type='relu'):
- super(ReluLayer, self).__init__()
- relu_type = relu_type.lower()
- if relu_type == 'relu':
- self.func = nn.ReLU(True)
- elif relu_type == 'leakyrelu':
- self.func = nn.LeakyReLU(0.2, inplace=True)
- elif relu_type == 'prelu':
- self.func = nn.PReLU(channels)
- elif relu_type == 'selu':
- self.func = nn.SELU(True)
- elif relu_type == 'none':
- self.func = lambda x: x * 1.0
- else:
- assert 1 == 0, f'Relu type {relu_type} not support.'
-
- def forward(self, x):
- return self.func(x)
-
-
-class ConvLayer(nn.Module):
-
- def __init__(self,
- in_channels,
- out_channels,
- kernel_size=3,
- scale='none',
- norm_type='none',
- relu_type='none',
- use_pad=True,
- bias=True):
- super(ConvLayer, self).__init__()
- self.use_pad = use_pad
- self.norm_type = norm_type
- if norm_type in ['bn']:
- bias = False
-
- stride = 2 if scale == 'down' else 1
-
- self.scale_func = lambda x: x
- if scale == 'up':
- self.scale_func = lambda x: nn.functional.interpolate(x, scale_factor=2, mode='nearest')
-
- self.reflection_pad = nn.ReflectionPad2d(int(np.ceil((kernel_size - 1.) / 2)))
- self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, bias=bias)
-
- self.relu = ReluLayer(out_channels, relu_type)
- self.norm = NormLayer(out_channels, norm_type=norm_type)
-
- def forward(self, x):
- out = self.scale_func(x)
- if self.use_pad:
- out = self.reflection_pad(out)
- out = self.conv2d(out)
- out = self.norm(out)
- out = self.relu(out)
- return out
-
-
-class ResidualBlock(nn.Module):
- """
- Residual block recommended in: http://torch.ch/blog/2016/02/04/resnets.html
- """
-
- def __init__(self, c_in, c_out, relu_type='prelu', norm_type='bn', scale='none'):
- super(ResidualBlock, self).__init__()
-
- if scale == 'none' and c_in == c_out:
- self.shortcut_func = lambda x: x
- else:
- self.shortcut_func = ConvLayer(c_in, c_out, 3, scale)
-
- scale_config_dict = {'down': ['none', 'down'], 'up': ['up', 'none'], 'none': ['none', 'none']}
- scale_conf = scale_config_dict[scale]
-
- self.conv1 = ConvLayer(c_in, c_out, 3, scale_conf[0], norm_type=norm_type, relu_type=relu_type)
- self.conv2 = ConvLayer(c_out, c_out, 3, scale_conf[1], norm_type=norm_type, relu_type='none')
-
- def forward(self, x):
- identity = self.shortcut_func(x)
-
- res = self.conv1(x)
- res = self.conv2(res)
- return identity + res
-
-
-class ParseNet(nn.Module):
-
- def __init__(self,
- in_size=128,
- out_size=128,
- min_feat_size=32,
- base_ch=64,
- parsing_ch=19,
- res_depth=10,
- relu_type='LeakyReLU',
- norm_type='bn',
- ch_range=[32, 256]):
- super().__init__()
- self.res_depth = res_depth
- act_args = {'norm_type': norm_type, 'relu_type': relu_type}
- min_ch, max_ch = ch_range
-
- ch_clip = lambda x: max(min_ch, min(x, max_ch)) # noqa: E731
- min_feat_size = min(in_size, min_feat_size)
-
- down_steps = int(np.log2(in_size // min_feat_size))
- up_steps = int(np.log2(out_size // min_feat_size))
-
- # =============== define encoder-body-decoder ====================
- self.encoder = []
- self.encoder.append(ConvLayer(3, base_ch, 3, 1))
- head_ch = base_ch
- for i in range(down_steps):
- cin, cout = ch_clip(head_ch), ch_clip(head_ch * 2)
- self.encoder.append(ResidualBlock(cin, cout, scale='down', **act_args))
- head_ch = head_ch * 2
-
- self.body = []
- for i in range(res_depth):
- self.body.append(ResidualBlock(ch_clip(head_ch), ch_clip(head_ch), **act_args))
-
- self.decoder = []
- for i in range(up_steps):
- cin, cout = ch_clip(head_ch), ch_clip(head_ch // 2)
- self.decoder.append(ResidualBlock(cin, cout, scale='up', **act_args))
- head_ch = head_ch // 2
-
- self.encoder = nn.Sequential(*self.encoder)
- self.body = nn.Sequential(*self.body)
- self.decoder = nn.Sequential(*self.decoder)
- self.out_img_conv = ConvLayer(ch_clip(head_ch), 3)
- self.out_mask_conv = ConvLayer(ch_clip(head_ch), parsing_ch)
-
- def forward(self, x):
- feat = self.encoder(x)
- x = feat + self.body(feat)
- x = self.decoder(x)
- out_img = self.out_img_conv(x)
- out_mask = self.out_mask_conv(x)
- return out_mask, out_img
diff --git a/modules/facelib/parsing/resnet.py b/modules/facelib/parsing/resnet.py
deleted file mode 100644
index fec8e82cf..000000000
--- a/modules/facelib/parsing/resnet.py
+++ /dev/null
@@ -1,69 +0,0 @@
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-def conv3x3(in_planes, out_planes, stride=1):
- """3x3 convolution with padding"""
- return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
-
-
-class BasicBlock(nn.Module):
-
- def __init__(self, in_chan, out_chan, stride=1):
- super(BasicBlock, self).__init__()
- self.conv1 = conv3x3(in_chan, out_chan, stride)
- self.bn1 = nn.BatchNorm2d(out_chan)
- self.conv2 = conv3x3(out_chan, out_chan)
- self.bn2 = nn.BatchNorm2d(out_chan)
- self.relu = nn.ReLU(inplace=True)
- self.downsample = None
- if in_chan != out_chan or stride != 1:
- self.downsample = nn.Sequential(
- nn.Conv2d(in_chan, out_chan, kernel_size=1, stride=stride, bias=False),
- nn.BatchNorm2d(out_chan),
- )
-
- def forward(self, x):
- residual = self.conv1(x)
- residual = F.relu(self.bn1(residual))
- residual = self.conv2(residual)
- residual = self.bn2(residual)
-
- shortcut = x
- if self.downsample is not None:
- shortcut = self.downsample(x)
-
- out = shortcut + residual
- out = self.relu(out)
- return out
-
-
-def create_layer_basic(in_chan, out_chan, bnum, stride=1):
- layers = [BasicBlock(in_chan, out_chan, stride=stride)]
- for i in range(bnum - 1):
- layers.append(BasicBlock(out_chan, out_chan, stride=1))
- return nn.Sequential(*layers)
-
-
-class ResNet18(nn.Module):
-
- def __init__(self):
- super(ResNet18, self).__init__()
- self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
- self.bn1 = nn.BatchNorm2d(64)
- self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
- self.layer1 = create_layer_basic(64, 64, bnum=2, stride=1)
- self.layer2 = create_layer_basic(64, 128, bnum=2, stride=2)
- self.layer3 = create_layer_basic(128, 256, bnum=2, stride=2)
- self.layer4 = create_layer_basic(256, 512, bnum=2, stride=2)
-
- def forward(self, x):
- x = self.conv1(x)
- x = F.relu(self.bn1(x))
- x = self.maxpool(x)
-
- x = self.layer1(x)
- feat8 = self.layer2(x) # 1/8
- feat16 = self.layer3(feat8) # 1/16
- feat32 = self.layer4(feat16) # 1/32
- return feat8, feat16, feat32
diff --git a/modules/facelib/utils/__init__.py b/modules/facelib/utils/__init__.py
deleted file mode 100644
index 23ef0352c..000000000
--- a/modules/facelib/utils/__init__.py
+++ /dev/null
@@ -1,7 +0,0 @@
-from .face_utils import align_crop_face_landmarks, compute_increased_bbox, get_valid_bboxes, paste_face_back
-from .misc import img2tensor, load_file_from_url, download_pretrained_models, scandir
-
-__all__ = [
- 'align_crop_face_landmarks', 'compute_increased_bbox', 'get_valid_bboxes', 'load_file_from_url',
- 'download_pretrained_models', 'paste_face_back', 'img2tensor', 'scandir'
-]
diff --git a/modules/facelib/utils/face_restoration_helper.py b/modules/facelib/utils/face_restoration_helper.py
deleted file mode 100644
index 23132a01c..000000000
--- a/modules/facelib/utils/face_restoration_helper.py
+++ /dev/null
@@ -1,460 +0,0 @@
-import cv2
-import numpy as np
-import os
-import torch
-from torchvision.transforms.functional import normalize
-
-from ..detection import init_detection_model
-from ..parsing import init_parsing_model
-from ..utils.misc import img2tensor, imwrite, is_gray, bgr2gray
-
-
-def get_largest_face(det_faces, h, w):
-
- def get_location(val, length):
- if val < 0:
- return 0
- elif val > length:
- return length
- else:
- return val
-
- face_areas = []
- for det_face in det_faces:
- left = get_location(det_face[0], w)
- right = get_location(det_face[2], w)
- top = get_location(det_face[1], h)
- bottom = get_location(det_face[3], h)
- face_area = (right - left) * (bottom - top)
- face_areas.append(face_area)
- largest_idx = face_areas.index(max(face_areas))
- return det_faces[largest_idx], largest_idx
-
-
-def get_center_face(det_faces, h=0, w=0, center=None):
- if center is not None:
- center = np.array(center)
- else:
- center = np.array([w / 2, h / 2])
- center_dist = []
- for det_face in det_faces:
- face_center = np.array([(det_face[0] + det_face[2]) / 2, (det_face[1] + det_face[3]) / 2])
- dist = np.linalg.norm(face_center - center)
- center_dist.append(dist)
- center_idx = center_dist.index(min(center_dist))
- return det_faces[center_idx], center_idx
-
-
-class FaceRestoreHelper(object):
- """Helper for the face restoration pipeline (base class)."""
-
- def __init__(self,
- upscale_factor,
- face_size=512,
- crop_ratio=(1, 1),
- det_model='retinaface_resnet50',
- save_ext='png',
- template_3points=False,
- pad_blur=False,
- use_parse=False,
- device=None):
- self.template_3points = template_3points # improve robustness
- self.upscale_factor = int(upscale_factor)
- # the cropped face ratio based on the square face
- self.crop_ratio = crop_ratio # (h, w)
- assert (self.crop_ratio[0] >= 1 and self.crop_ratio[1] >= 1), 'crop ration only supports >=1'
- self.face_size = (int(face_size * self.crop_ratio[1]), int(face_size * self.crop_ratio[0]))
-
- if self.template_3points:
- self.face_template = np.array([[192, 240], [319, 240], [257, 371]])
- else:
- # standard 5 landmarks for FFHQ faces with 512 x 512
- # facexlib
- self.face_template = np.array([[192.98138, 239.94708], [318.90277, 240.1936], [256.63416, 314.01935],
- [201.26117, 371.41043], [313.08905, 371.15118]])
-
- # dlib: left_eye: 36:41 right_eye: 42:47 nose: 30,32,33,34 left mouth corner: 48 right mouth corner: 54
- # self.face_template = np.array([[193.65928, 242.98541], [318.32558, 243.06108], [255.67984, 328.82894],
- # [198.22603, 372.82502], [313.91018, 372.75659]])
-
-
- self.face_template = self.face_template * (face_size / 512.0)
- if self.crop_ratio[0] > 1:
- self.face_template[:, 1] += face_size * (self.crop_ratio[0] - 1) / 2
- if self.crop_ratio[1] > 1:
- self.face_template[:, 0] += face_size * (self.crop_ratio[1] - 1) / 2
- self.save_ext = save_ext
- self.pad_blur = pad_blur
- if self.pad_blur is True:
- self.template_3points = False
-
- self.all_landmarks_5 = []
- self.det_faces = []
- self.affine_matrices = []
- self.inverse_affine_matrices = []
- self.cropped_faces = []
- self.restored_faces = []
- self.pad_input_imgs = []
-
- if device is None:
- self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
- else:
- self.device = device
-
- # init face detection model
- self.face_det = init_detection_model(det_model, half=False, device=self.device)
-
- # init face parsing model
- self.use_parse = use_parse
- self.face_parse = init_parsing_model(model_name='parsenet', device=self.device)
-
- def set_upscale_factor(self, upscale_factor):
- self.upscale_factor = upscale_factor
-
- def read_image(self, img):
- """img can be image path or cv2 loaded image."""
- # self.input_img is Numpy array, (h, w, c), BGR, uint8, [0, 255]
- if isinstance(img, str):
- img = cv2.imread(img)
-
- if np.max(img) > 256: # 16-bit image
- img = img / 65535 * 255
- if len(img.shape) == 2: # gray image
- img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
- elif img.shape[2] == 4: # BGRA image with alpha channel
- img = img[:, :, 0:3]
-
- self.input_img = img
- self.is_gray = is_gray(img, threshold=5)
- if self.is_gray:
- print('Grayscale input: True')
-
- if min(self.input_img.shape[:2])<512:
- f = 512.0/min(self.input_img.shape[:2])
- self.input_img = cv2.resize(self.input_img, (0,0), fx=f, fy=f, interpolation=cv2.INTER_LINEAR)
-
- def get_face_landmarks_5(self,
- only_keep_largest=False,
- only_center_face=False,
- resize=None,
- blur_ratio=0.01,
- eye_dist_threshold=None):
- if resize is None:
- scale = 1
- input_img = self.input_img
- else:
- h, w = self.input_img.shape[0:2]
- scale = resize / min(h, w)
- scale = max(1, scale) # always scale up
- h, w = int(h * scale), int(w * scale)
- interp = cv2.INTER_AREA if scale < 1 else cv2.INTER_LINEAR
- input_img = cv2.resize(self.input_img, (w, h), interpolation=interp)
-
- with torch.no_grad():
- bboxes = self.face_det.detect_faces(input_img)
-
- if bboxes is None or bboxes.shape[0] == 0:
- return 0
- else:
- bboxes = bboxes / scale
-
- for bbox in bboxes:
- # remove faces with too small eye distance: side faces or too small faces
- eye_dist = np.linalg.norm([bbox[6] - bbox[8], bbox[7] - bbox[9]])
- if eye_dist_threshold is not None and (eye_dist < eye_dist_threshold):
- continue
-
- if self.template_3points:
- landmark = np.array([[bbox[i], bbox[i + 1]] for i in range(5, 11, 2)])
- else:
- landmark = np.array([[bbox[i], bbox[i + 1]] for i in range(5, 15, 2)])
- self.all_landmarks_5.append(landmark)
- self.det_faces.append(bbox[0:5])
-
- if len(self.det_faces) == 0:
- return 0
- if only_keep_largest:
- h, w, _ = self.input_img.shape
- self.det_faces, largest_idx = get_largest_face(self.det_faces, h, w)
- self.all_landmarks_5 = [self.all_landmarks_5[largest_idx]]
- elif only_center_face:
- h, w, _ = self.input_img.shape
- self.det_faces, center_idx = get_center_face(self.det_faces, h, w)
- self.all_landmarks_5 = [self.all_landmarks_5[center_idx]]
-
- # pad blurry images
- if self.pad_blur:
- self.pad_input_imgs = []
- for landmarks in self.all_landmarks_5:
- # get landmarks
- eye_left = landmarks[0, :]
- eye_right = landmarks[1, :]
- eye_avg = (eye_left + eye_right) * 0.5
- mouth_avg = (landmarks[3, :] + landmarks[4, :]) * 0.5
- eye_to_eye = eye_right - eye_left
- eye_to_mouth = mouth_avg - eye_avg
-
- # Get the oriented crop rectangle
- # x: half width of the oriented crop rectangle
- x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1]
- # - np.flipud(eye_to_mouth) * [-1, 1]: rotate 90 clockwise
- # norm with the hypotenuse: get the direction
- x /= np.hypot(*x) # get the hypotenuse of a right triangle
- rect_scale = 1.5
- x *= max(np.hypot(*eye_to_eye) * 2.0 * rect_scale, np.hypot(*eye_to_mouth) * 1.8 * rect_scale)
- # y: half height of the oriented crop rectangle
- y = np.flipud(x) * [-1, 1]
-
- # c: center
- c = eye_avg + eye_to_mouth * 0.1
- # quad: (left_top, left_bottom, right_bottom, right_top)
- quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])
- # qsize: side length of the square
- qsize = np.hypot(*x) * 2
- border = max(int(np.rint(qsize * 0.1)), 3)
-
- # get pad
- # pad: (width_left, height_top, width_right, height_bottom)
- pad = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),
- int(np.ceil(max(quad[:, 1]))))
- pad = [
- max(-pad[0] + border, 1),
- max(-pad[1] + border, 1),
- max(pad[2] - self.input_img.shape[0] + border, 1),
- max(pad[3] - self.input_img.shape[1] + border, 1)
- ]
-
- if max(pad) > 1:
- # pad image
- pad_img = np.pad(self.input_img, ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect')
- # modify landmark coords
- landmarks[:, 0] += pad[0]
- landmarks[:, 1] += pad[1]
- # blur pad images
- h, w, _ = pad_img.shape
- y, x, _ = np.ogrid[:h, :w, :1]
- mask = np.maximum(1.0 - np.minimum(np.float32(x) / pad[0],
- np.float32(w - 1 - x) / pad[2]),
- 1.0 - np.minimum(np.float32(y) / pad[1],
- np.float32(h - 1 - y) / pad[3]))
- blur = int(qsize * blur_ratio)
- if blur % 2 == 0:
- blur += 1
- blur_img = cv2.boxFilter(pad_img, 0, ksize=(blur, blur))
- # blur_img = cv2.GaussianBlur(pad_img, (blur, blur), 0)
-
- pad_img = pad_img.astype('float32')
- pad_img += (blur_img - pad_img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0)
- pad_img += (np.median(pad_img, axis=(0, 1)) - pad_img) * np.clip(mask, 0.0, 1.0)
- pad_img = np.clip(pad_img, 0, 255) # float32, [0, 255]
- self.pad_input_imgs.append(pad_img)
- else:
- self.pad_input_imgs.append(np.copy(self.input_img))
-
- return len(self.all_landmarks_5)
-
- def align_warp_face(self, save_cropped_path=None, border_mode='constant'):
- """Align and warp faces with face template.
- """
- if self.pad_blur:
- assert len(self.pad_input_imgs) == len(
- self.all_landmarks_5), f'Mismatched samples: {len(self.pad_input_imgs)} and {len(self.all_landmarks_5)}'
- for idx, landmark in enumerate(self.all_landmarks_5):
- # use 5 landmarks to get affine matrix
- # use cv2.LMEDS method for the equivalence to skimage transform
- # ref: https://blog.csdn.net/yichxi/article/details/115827338
- affine_matrix = cv2.estimateAffinePartial2D(landmark, self.face_template, method=cv2.LMEDS)[0]
- self.affine_matrices.append(affine_matrix)
- # warp and crop faces
- if border_mode == 'constant':
- border_mode = cv2.BORDER_CONSTANT
- elif border_mode == 'reflect101':
- border_mode = cv2.BORDER_REFLECT101
- elif border_mode == 'reflect':
- border_mode = cv2.BORDER_REFLECT
- if self.pad_blur:
- input_img = self.pad_input_imgs[idx]
- else:
- input_img = self.input_img
- cropped_face = cv2.warpAffine(
- input_img, affine_matrix, self.face_size, borderMode=border_mode, borderValue=(135, 133, 132)) # gray
- self.cropped_faces.append(cropped_face)
- # save the cropped face
- if save_cropped_path is not None:
- path = os.path.splitext(save_cropped_path)[0]
- save_path = f'{path}_{idx:02d}.{self.save_ext}'
- imwrite(cropped_face, save_path)
-
- def get_inverse_affine(self, save_inverse_affine_path=None):
- """Get inverse affine matrix."""
- for idx, affine_matrix in enumerate(self.affine_matrices):
- inverse_affine = cv2.invertAffineTransform(affine_matrix)
- inverse_affine *= self.upscale_factor
- self.inverse_affine_matrices.append(inverse_affine)
- # save inverse affine matrices
- if save_inverse_affine_path is not None:
- path, _ = os.path.splitext(save_inverse_affine_path)
- save_path = f'{path}_{idx:02d}.pth'
- torch.save(inverse_affine, save_path)
-
-
- def add_restored_face(self, face):
- if self.is_gray:
- face = bgr2gray(face) # convert img into grayscale
- self.restored_faces.append(face)
-
-
- def paste_faces_to_input_image(self, save_path=None, upsample_img=None, draw_box=False, face_upsampler=None):
- h, w, _ = self.input_img.shape
- h_up, w_up = int(h * self.upscale_factor), int(w * self.upscale_factor)
-
- if upsample_img is None:
- # simply resize the background
- # upsample_img = cv2.resize(self.input_img, (w_up, h_up), interpolation=cv2.INTER_LANCZOS4)
- upsample_img = cv2.resize(self.input_img, (w_up, h_up), interpolation=cv2.INTER_LINEAR)
- else:
- upsample_img = cv2.resize(upsample_img, (w_up, h_up), interpolation=cv2.INTER_LANCZOS4)
-
- assert len(self.restored_faces) == len(
- self.inverse_affine_matrices), ('length of restored_faces and affine_matrices are different.')
-
- inv_mask_borders = []
- for restored_face, inverse_affine in zip(self.restored_faces, self.inverse_affine_matrices):
- if face_upsampler is not None:
- restored_face = face_upsampler.enhance(restored_face, outscale=self.upscale_factor)[0]
- inverse_affine /= self.upscale_factor
- inverse_affine[:, 2] *= self.upscale_factor
- face_size = (self.face_size[0]*self.upscale_factor, self.face_size[1]*self.upscale_factor)
- else:
- # Add an offset to inverse affine matrix, for more precise back alignment
- if self.upscale_factor > 1:
- extra_offset = 0.5 * self.upscale_factor
- else:
- extra_offset = 0
- inverse_affine[:, 2] += extra_offset
- face_size = self.face_size
- inv_restored = cv2.warpAffine(restored_face, inverse_affine, (w_up, h_up))
-
- # if draw_box or not self.use_parse: # use square parse maps
- # mask = np.ones(face_size, dtype=np.float32)
- # inv_mask = cv2.warpAffine(mask, inverse_affine, (w_up, h_up))
- # # remove the black borders
- # inv_mask_erosion = cv2.erode(
- # inv_mask, np.ones((int(2 * self.upscale_factor), int(2 * self.upscale_factor)), np.uint8))
- # pasted_face = inv_mask_erosion[:, :, None] * inv_restored
- # total_face_area = np.sum(inv_mask_erosion) # // 3
- # # add border
- # if draw_box:
- # h, w = face_size
- # mask_border = np.ones((h, w, 3), dtype=np.float32)
- # border = int(1400/np.sqrt(total_face_area))
- # mask_border[border:h-border, border:w-border,:] = 0
- # inv_mask_border = cv2.warpAffine(mask_border, inverse_affine, (w_up, h_up))
- # inv_mask_borders.append(inv_mask_border)
- # if not self.use_parse:
- # # compute the fusion edge based on the area of face
- # w_edge = int(total_face_area**0.5) // 20
- # erosion_radius = w_edge * 2
- # inv_mask_center = cv2.erode(inv_mask_erosion, np.ones((erosion_radius, erosion_radius), np.uint8))
- # blur_size = w_edge * 2
- # inv_soft_mask = cv2.GaussianBlur(inv_mask_center, (blur_size + 1, blur_size + 1), 0)
- # if len(upsample_img.shape) == 2: # upsample_img is gray image
- # upsample_img = upsample_img[:, :, None]
- # inv_soft_mask = inv_soft_mask[:, :, None]
-
- # always use square mask
- mask = np.ones(face_size, dtype=np.float32)
- inv_mask = cv2.warpAffine(mask, inverse_affine, (w_up, h_up))
- # remove the black borders
- inv_mask_erosion = cv2.erode(
- inv_mask, np.ones((int(2 * self.upscale_factor), int(2 * self.upscale_factor)), np.uint8))
- pasted_face = inv_mask_erosion[:, :, None] * inv_restored
- total_face_area = np.sum(inv_mask_erosion) # // 3
- # add border
- if draw_box:
- h, w = face_size
- mask_border = np.ones((h, w, 3), dtype=np.float32)
- border = int(1400/np.sqrt(total_face_area))
- mask_border[border:h-border, border:w-border,:] = 0
- inv_mask_border = cv2.warpAffine(mask_border, inverse_affine, (w_up, h_up))
- inv_mask_borders.append(inv_mask_border)
- # compute the fusion edge based on the area of face
- w_edge = int(total_face_area**0.5) // 20
- erosion_radius = w_edge * 2
- inv_mask_center = cv2.erode(inv_mask_erosion, np.ones((erosion_radius, erosion_radius), np.uint8))
- blur_size = w_edge * 2
- inv_soft_mask = cv2.GaussianBlur(inv_mask_center, (blur_size + 1, blur_size + 1), 0)
- if len(upsample_img.shape) == 2: # upsample_img is gray image
- upsample_img = upsample_img[:, :, None]
- inv_soft_mask = inv_soft_mask[:, :, None]
-
- # parse mask
- if self.use_parse:
- # inference
- face_input = cv2.resize(restored_face, (512, 512), interpolation=cv2.INTER_LINEAR)
- face_input = img2tensor(face_input.astype('float32') / 255., bgr2rgb=True, float32=True)
- normalize(face_input, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
- face_input = torch.unsqueeze(face_input, 0).to(self.device)
- with torch.no_grad():
- out = self.face_parse(face_input)[0]
- out = out.argmax(dim=1).squeeze().cpu().numpy()
-
- parse_mask = np.zeros(out.shape)
- MASK_COLORMAP = [0, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 0, 255, 0, 0, 0]
- for idx, color in enumerate(MASK_COLORMAP):
- parse_mask[out == idx] = color
- # blur the mask
- parse_mask = cv2.GaussianBlur(parse_mask, (101, 101), 11)
- parse_mask = cv2.GaussianBlur(parse_mask, (101, 101), 11)
- # remove the black borders
- thres = 10
- parse_mask[:thres, :] = 0
- parse_mask[-thres:, :] = 0
- parse_mask[:, :thres] = 0
- parse_mask[:, -thres:] = 0
- parse_mask = parse_mask / 255.
-
- parse_mask = cv2.resize(parse_mask, face_size)
- parse_mask = cv2.warpAffine(parse_mask, inverse_affine, (w_up, h_up), flags=3)
- inv_soft_parse_mask = parse_mask[:, :, None]
- # pasted_face = inv_restored
- fuse_mask = (inv_soft_parse_mask 256: # 16-bit image
- upsample_img = upsample_img.astype(np.uint16)
- else:
- upsample_img = upsample_img.astype(np.uint8)
-
- # draw bounding box
- if draw_box:
- # upsample_input_img = cv2.resize(input_img, (w_up, h_up))
- img_color = np.ones([*upsample_img.shape], dtype=np.float32)
- img_color[:,:,0] = 0
- img_color[:,:,1] = 255
- img_color[:,:,2] = 0
- for inv_mask_border in inv_mask_borders:
- upsample_img = inv_mask_border * img_color + (1 - inv_mask_border) * upsample_img
- # upsample_input_img = inv_mask_border * img_color + (1 - inv_mask_border) * upsample_input_img
-
- if save_path is not None:
- path = os.path.splitext(save_path)[0]
- save_path = f'{path}.{self.save_ext}'
- imwrite(upsample_img, save_path)
- return upsample_img
-
- def clean_all(self):
- self.all_landmarks_5 = []
- self.restored_faces = []
- self.affine_matrices = []
- self.cropped_faces = []
- self.inverse_affine_matrices = []
- self.det_faces = []
- self.pad_input_imgs = []
diff --git a/modules/facelib/utils/face_utils.py b/modules/facelib/utils/face_utils.py
deleted file mode 100644
index 3470a4c7e..000000000
--- a/modules/facelib/utils/face_utils.py
+++ /dev/null
@@ -1,208 +0,0 @@
-import os
-import cv2
-import numpy as np
-import torch
-
-
-def compute_increased_bbox(bbox, increase_area, preserve_aspect=True):
- left, top, right, bot = bbox
- width = right - left
- height = bot - top
-
- if preserve_aspect:
- width_increase = max(increase_area, ((1 + 2 * increase_area) * height - width) / (2 * width))
- height_increase = max(increase_area, ((1 + 2 * increase_area) * width - height) / (2 * height))
- else:
- width_increase = height_increase = increase_area
- left = int(left - width_increase * width)
- top = int(top - height_increase * height)
- right = int(right + width_increase * width)
- bot = int(bot + height_increase * height)
- return (left, top, right, bot)
-
-
-def get_valid_bboxes(bboxes, h, w):
- left = max(bboxes[0], 0)
- top = max(bboxes[1], 0)
- right = min(bboxes[2], w)
- bottom = min(bboxes[3], h)
- return (left, top, right, bottom)
-
-
-def align_crop_face_landmarks(img,
- landmarks,
- output_size,
- transform_size=None,
- enable_padding=True,
- return_inverse_affine=False,
- shrink_ratio=(1, 1)):
- """Align and crop face with landmarks.
-
- The output_size and transform_size are based on width. The height is
- adjusted based on shrink_ratio_h/shring_ration_w.
-
- Modified from:
- https://github.com/NVlabs/ffhq-dataset/blob/master/download_ffhq.py
-
- Args:
- img (Numpy array): Input image.
- landmarks (Numpy array): 5 or 68 or 98 landmarks.
- output_size (int): Output face size.
- transform_size (ing): Transform size. Usually the four time of
- output_size.
- enable_padding (float): Default: True.
- shrink_ratio (float | tuple[float] | list[float]): Shring the whole
- face for height and width (crop larger area). Default: (1, 1).
-
- Returns:
- (Numpy array): Cropped face.
- """
- lm_type = 'retinaface_5' # Options: dlib_5, retinaface_5
-
- if isinstance(shrink_ratio, (float, int)):
- shrink_ratio = (shrink_ratio, shrink_ratio)
- if transform_size is None:
- transform_size = output_size * 4
-
- # Parse landmarks
- lm = np.array(landmarks)
- if lm.shape[0] == 5 and lm_type == 'retinaface_5':
- eye_left = lm[0]
- eye_right = lm[1]
- mouth_avg = (lm[3] + lm[4]) * 0.5
- elif lm.shape[0] == 5 and lm_type == 'dlib_5':
- lm_eye_left = lm[2:4]
- lm_eye_right = lm[0:2]
- eye_left = np.mean(lm_eye_left, axis=0)
- eye_right = np.mean(lm_eye_right, axis=0)
- mouth_avg = lm[4]
- elif lm.shape[0] == 68:
- lm_eye_left = lm[36:42]
- lm_eye_right = lm[42:48]
- eye_left = np.mean(lm_eye_left, axis=0)
- eye_right = np.mean(lm_eye_right, axis=0)
- mouth_avg = (lm[48] + lm[54]) * 0.5
- elif lm.shape[0] == 98:
- lm_eye_left = lm[60:68]
- lm_eye_right = lm[68:76]
- eye_left = np.mean(lm_eye_left, axis=0)
- eye_right = np.mean(lm_eye_right, axis=0)
- mouth_avg = (lm[76] + lm[82]) * 0.5
-
- eye_avg = (eye_left + eye_right) * 0.5
- eye_to_eye = eye_right - eye_left
- eye_to_mouth = mouth_avg - eye_avg
-
- # Get the oriented crop rectangle
- # x: half width of the oriented crop rectangle
- x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1]
- # - np.flipud(eye_to_mouth) * [-1, 1]: rotate 90 clockwise
- # norm with the hypotenuse: get the direction
- x /= np.hypot(*x) # get the hypotenuse of a right triangle
- rect_scale = 1
- x *= max(np.hypot(*eye_to_eye) * 2.0 * rect_scale, np.hypot(*eye_to_mouth) * 1.8 * rect_scale)
- # y: half height of the oriented crop rectangle
- y = np.flipud(x) * [-1, 1]
-
- x *= shrink_ratio[1] # width
- y *= shrink_ratio[0] # height
-
- # c: center
- c = eye_avg + eye_to_mouth * 0.1
- # quad: (left_top, left_bottom, right_bottom, right_top)
- quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])
- # qsize: side length of the square
- qsize = np.hypot(*x) * 2
-
- quad_ori = np.copy(quad)
- # Shrink, for large face
- shrink = int(np.floor(qsize / output_size * 0.5))
- if shrink > 1:
- h, w = img.shape[0:2]
- rsize = (int(np.rint(float(w) / shrink)), int(np.rint(float(h) / shrink)))
- img = cv2.resize(img, rsize, interpolation=cv2.INTER_AREA)
- quad /= shrink
- qsize /= shrink
-
- # Crop
- h, w = img.shape[0:2]
- border = max(int(np.rint(qsize * 0.1)), 3)
- crop = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),
- int(np.ceil(max(quad[:, 1]))))
- crop = (max(crop[0] - border, 0), max(crop[1] - border, 0), min(crop[2] + border, w), min(crop[3] + border, h))
- if crop[2] - crop[0] < w or crop[3] - crop[1] < h:
- img = img[crop[1]:crop[3], crop[0]:crop[2], :]
- quad -= crop[0:2]
-
- # Pad
- # pad: (width_left, height_top, width_right, height_bottom)
- h, w = img.shape[0:2]
- pad = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),
- int(np.ceil(max(quad[:, 1]))))
- pad = (max(-pad[0] + border, 0), max(-pad[1] + border, 0), max(pad[2] - w + border, 0), max(pad[3] - h + border, 0))
- if enable_padding and max(pad) > border - 4:
- pad = np.maximum(pad, int(np.rint(qsize * 0.3)))
- img = np.pad(img, ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect')
- h, w = img.shape[0:2]
- y, x, _ = np.ogrid[:h, :w, :1]
- mask = np.maximum(1.0 - np.minimum(np.float32(x) / pad[0],
- np.float32(w - 1 - x) / pad[2]),
- 1.0 - np.minimum(np.float32(y) / pad[1],
- np.float32(h - 1 - y) / pad[3]))
- blur = int(qsize * 0.02)
- if blur % 2 == 0:
- blur += 1
- blur_img = cv2.boxFilter(img, 0, ksize=(blur, blur))
-
- img = img.astype('float32')
- img += (blur_img - img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0)
- img += (np.median(img, axis=(0, 1)) - img) * np.clip(mask, 0.0, 1.0)
- img = np.clip(img, 0, 255) # float32, [0, 255]
- quad += pad[:2]
-
- # Transform use cv2
- h_ratio = shrink_ratio[0] / shrink_ratio[1]
- dst_h, dst_w = int(transform_size * h_ratio), transform_size
- template = np.array([[0, 0], [0, dst_h], [dst_w, dst_h], [dst_w, 0]])
- # use cv2.LMEDS method for the equivalence to skimage transform
- # ref: https://blog.csdn.net/yichxi/article/details/115827338
- affine_matrix = cv2.estimateAffinePartial2D(quad, template, method=cv2.LMEDS)[0]
- cropped_face = cv2.warpAffine(
- img, affine_matrix, (dst_w, dst_h), borderMode=cv2.BORDER_CONSTANT, borderValue=(135, 133, 132)) # gray
-
- if output_size < transform_size:
- cropped_face = cv2.resize(
- cropped_face, (output_size, int(output_size * h_ratio)), interpolation=cv2.INTER_LINEAR)
-
- if return_inverse_affine:
- dst_h, dst_w = int(output_size * h_ratio), output_size
- template = np.array([[0, 0], [0, dst_h], [dst_w, dst_h], [dst_w, 0]])
- # use cv2.LMEDS method for the equivalence to skimage transform
- # ref: https://blog.csdn.net/yichxi/article/details/115827338
- affine_matrix = cv2.estimateAffinePartial2D(
- quad_ori, np.array([[0, 0], [0, output_size], [dst_w, dst_h], [dst_w, 0]]), method=cv2.LMEDS)[0]
- inverse_affine = cv2.invertAffineTransform(affine_matrix)
- else:
- inverse_affine = None
- return cropped_face, inverse_affine
-
-
-def paste_face_back(img, face, inverse_affine):
- h, w = img.shape[0:2]
- face_h, face_w = face.shape[0:2]
- inv_restored = cv2.warpAffine(face, inverse_affine, (w, h))
- mask = np.ones((face_h, face_w, 3), dtype=np.float32)
- inv_mask = cv2.warpAffine(mask, inverse_affine, (w, h))
- # remove the black borders
- inv_mask_erosion = cv2.erode(inv_mask, np.ones((2, 2), np.uint8))
- inv_restored_remove_border = inv_mask_erosion * inv_restored
- total_face_area = np.sum(inv_mask_erosion) // 3
- # compute the fusion edge based on the area of face
- w_edge = int(total_face_area**0.5) // 20
- erosion_radius = w_edge * 2
- inv_mask_center = cv2.erode(inv_mask_erosion, np.ones((erosion_radius, erosion_radius), np.uint8))
- blur_size = w_edge * 2
- inv_soft_mask = cv2.GaussianBlur(inv_mask_center, (blur_size + 1, blur_size + 1), 0)
- img = inv_soft_mask * inv_restored_remove_border + (1 - inv_soft_mask) * img
- # float32, [0, 255]
- return img
diff --git a/modules/facelib/utils/misc.py b/modules/facelib/utils/misc.py
deleted file mode 100644
index 1f14e6c37..000000000
--- a/modules/facelib/utils/misc.py
+++ /dev/null
@@ -1,174 +0,0 @@
-import cv2
-import os
-import os.path as osp
-import numpy as np
-from PIL import Image
-import torch
-from torch.hub import download_url_to_file, get_dir
-from urllib.parse import urlparse
-# from basicsr.utils.download_util import download_file_from_google_drive
-
-ROOT_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-
-def download_pretrained_models(file_ids, save_path_root):
- import gdown
-
- os.makedirs(save_path_root, exist_ok=True)
-
- for file_name, file_id in file_ids.items():
- file_url = 'https://drive.google.com/uc?id='+file_id
- save_path = osp.abspath(osp.join(save_path_root, file_name))
- if osp.exists(save_path):
- user_response = input(f'{file_name} already exist. Do you want to cover it? Y/N\n')
- if user_response.lower() == 'y':
- print(f'Covering {file_name} to {save_path}')
- gdown.download(file_url, save_path, quiet=False)
- # download_file_from_google_drive(file_id, save_path)
- elif user_response.lower() == 'n':
- print(f'Skipping {file_name}')
- else:
- raise ValueError('Wrong input. Only accepts Y/N.')
- else:
- print(f'Downloading {file_name} to {save_path}')
- gdown.download(file_url, save_path, quiet=False)
- # download_file_from_google_drive(file_id, save_path)
-
-
-def imwrite(img, file_path, params=None, auto_mkdir=True):
- """Write image to file.
-
- Args:
- img (ndarray): Image array to be written.
- file_path (str): Image file path.
- params (None or list): Same as opencv's :func:`imwrite` interface.
- auto_mkdir (bool): If the parent folder of `file_path` does not exist,
- whether to create it automatically.
-
- Returns:
- bool: Successful or not.
- """
- if auto_mkdir:
- dir_name = os.path.abspath(os.path.dirname(file_path))
- os.makedirs(dir_name, exist_ok=True)
- return cv2.imwrite(file_path, img, params)
-
-
-def img2tensor(imgs, bgr2rgb=True, float32=True):
- """Numpy array to tensor.
-
- Args:
- imgs (list[ndarray] | ndarray): Input images.
- bgr2rgb (bool): Whether to change bgr to rgb.
- float32 (bool): Whether to change to float32.
-
- Returns:
- list[tensor] | tensor: Tensor images. If returned results only have
- one element, just return tensor.
- """
-
- def _totensor(img, bgr2rgb, float32):
- if img.shape[2] == 3 and bgr2rgb:
- if img.dtype == 'float64':
- img = img.astype('float32')
- img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
- img = torch.from_numpy(img.transpose(2, 0, 1))
- if float32:
- img = img.float()
- return img
-
- if isinstance(imgs, list):
- return [_totensor(img, bgr2rgb, float32) for img in imgs]
- else:
- return _totensor(imgs, bgr2rgb, float32)
-
-
-def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
- """Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
- """
- if model_dir is None:
- hub_dir = get_dir()
- model_dir = os.path.join(hub_dir, 'checkpoints')
-
- os.makedirs(os.path.join(ROOT_DIR, model_dir), exist_ok=True)
-
- parts = urlparse(url)
- filename = os.path.basename(parts.path)
- if file_name is not None:
- filename = file_name
- cached_file = os.path.abspath(os.path.join(ROOT_DIR, model_dir, filename))
- if not os.path.exists(cached_file):
- print(f'Downloading: "{url}" to {cached_file}\n')
- download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
- return cached_file
-
-
-def scandir(dir_path, suffix=None, recursive=False, full_path=False):
- """Scan a directory to find the interested files.
- Args:
- dir_path (str): Path of the directory.
- suffix (str | tuple(str), optional): File suffix that we are
- interested in. Default: None.
- recursive (bool, optional): If set to True, recursively scan the
- directory. Default: False.
- full_path (bool, optional): If set to True, include the dir_path.
- Default: False.
- Returns:
- A generator for all the interested files with relative paths.
- """
-
- if (suffix is not None) and not isinstance(suffix, (str, tuple)):
- raise TypeError('"suffix" must be a string or tuple of strings')
-
- root = dir_path
-
- def _scandir(dir_path, suffix, recursive):
- for entry in os.scandir(dir_path):
- if not entry.name.startswith('.') and entry.is_file():
- if full_path:
- return_path = entry.path
- else:
- return_path = osp.relpath(entry.path, root)
-
- if suffix is None:
- yield return_path
- elif return_path.endswith(suffix):
- yield return_path
- else:
- if recursive:
- yield from _scandir(entry.path, suffix=suffix, recursive=recursive)
- else:
- continue
-
- return _scandir(dir_path, suffix=suffix, recursive=recursive)
-
-
-def is_gray(img, threshold=10):
- img = Image.fromarray(img)
- if len(img.getbands()) == 1:
- return True
- img1 = np.asarray(img.getchannel(channel=0), dtype=np.int16)
- img2 = np.asarray(img.getchannel(channel=1), dtype=np.int16)
- img3 = np.asarray(img.getchannel(channel=2), dtype=np.int16)
- diff1 = (img1 - img2).var()
- diff2 = (img2 - img3).var()
- diff3 = (img3 - img1).var()
- diff_sum = (diff1 + diff2 + diff3) / 3.0
- if diff_sum <= threshold:
- return True
- else:
- return False
-
-def rgb2gray(img, out_channel=3):
- r, g, b = img[:,:,0], img[:,:,1], img[:,:,2]
- gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
- if out_channel == 3:
- gray = gray[:,:,np.newaxis].repeat(3, axis=2)
- return gray
-
-def bgr2gray(img, out_channel=3):
- b, g, r = img[:,:,0], img[:,:,1], img[:,:,2]
- gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
- if out_channel == 3:
- gray = gray[:,:,np.newaxis].repeat(3, axis=2)
- return gray
diff --git a/modules/modelloader.py b/modules/modelloader.py
index 26cb228c7..ea45bea82 100644
--- a/modules/modelloader.py
+++ b/modules/modelloader.py
@@ -393,9 +393,6 @@ def cleanup_models():
src_path = os.path.join(models_path, "BSRGAN")
dest_path = os.path.join(models_path, "ESRGAN")
move_files(src_path, dest_path, ".pth")
- src_path = os.path.join(root_path, "gfpgan")
- dest_path = os.path.join(models_path, "GFPGAN")
- move_files(src_path, dest_path)
src_path = os.path.join(root_path, "SwinIR")
dest_path = os.path.join(models_path, "SwinIR")
move_files(src_path, dest_path)
diff --git a/modules/postprocess/codeformer_arch.py b/modules/postprocess/codeformer_arch.py
deleted file mode 100644
index f1a6bf601..000000000
--- a/modules/postprocess/codeformer_arch.py
+++ /dev/null
@@ -1,272 +0,0 @@
-# this file is copied from CodeFormer repository. Please see comment in modules/codeformer_model.py
-
-import math
-from typing import Optional
-import torch
-from torch import nn, Tensor
-import torch.nn.functional as F
-from modules.postprocess.vqgan_arch import VQAutoEncoder, ResBlock
-
-def calc_mean_std(feat, eps=1e-5):
- """Calculate mean and std for adaptive_instance_normalization.
-
- Args:
- feat (Tensor): 4D tensor.
- eps (float): A small value added to the variance to avoid
- divide-by-zero. Default: 1e-5.
- """
- size = feat.size()
- assert len(size) == 4, 'The input feature should be 4D tensor.'
- b, c = size[:2]
- feat_var = feat.view(b, c, -1).var(dim=2) + eps
- feat_std = feat_var.sqrt().view(b, c, 1, 1)
- feat_mean = feat.view(b, c, -1).mean(dim=2).view(b, c, 1, 1)
- return feat_mean, feat_std
-
-
-def adaptive_instance_normalization(content_feat, style_feat):
- """Adaptive instance normalization.
-
- Adjust the reference features to have the similar color and illuminations
- as those in the degradate features.
-
- Args:
- content_feat (Tensor): The reference feature.
- style_feat (Tensor): The degradate features.
- """
- size = content_feat.size()
- style_mean, style_std = calc_mean_std(style_feat)
- content_mean, content_std = calc_mean_std(content_feat)
- normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)
- return normalized_feat * style_std.expand(size) + style_mean.expand(size)
-
-
-class PositionEmbeddingSine(nn.Module):
- """
- This is a more standard version of the position embedding, very similar to the one
- used by the Attention is all you need paper, generalized to work on images.
- """
-
- def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
- super().__init__()
- self.num_pos_feats = num_pos_feats
- self.temperature = temperature
- self.normalize = normalize
- if scale is not None and normalize is False:
- raise ValueError("normalize should be True if scale is passed")
- if scale is None:
- scale = 2 * math.pi
- self.scale = scale
-
- def forward(self, x, mask=None):
- if mask is None:
- mask = torch.zeros((x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool)
- not_mask = ~mask
- y_embed = not_mask.cumsum(1, dtype=torch.float32)
- x_embed = not_mask.cumsum(2, dtype=torch.float32)
- if self.normalize:
- eps = 1e-6
- y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
- x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
-
- dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
- dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
-
- pos_x = x_embed[:, :, :, None] / dim_t
- pos_y = y_embed[:, :, :, None] / dim_t
- pos_x = torch.stack(
- (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4
- ).flatten(3)
- pos_y = torch.stack(
- (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4
- ).flatten(3)
- pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
- return pos
-
-def _get_activation_fn(activation):
- """Return an activation function given a string"""
- if activation == "relu":
- return F.relu
- if activation == "gelu":
- return F.gelu
- if activation == "glu":
- return F.glu
- raise RuntimeError(F"activation should be relu/gelu, not {activation}.")
-
-
-class TransformerSALayer(nn.Module):
- def __init__(self, embed_dim, nhead=8, dim_mlp=2048, dropout=0.0, activation="gelu"):
- super().__init__()
- self.self_attn = nn.MultiheadAttention(embed_dim, nhead, dropout=dropout)
- # Implementation of Feedforward model - MLP
- self.linear1 = nn.Linear(embed_dim, dim_mlp)
- self.dropout = nn.Dropout(dropout)
- self.linear2 = nn.Linear(dim_mlp, embed_dim)
-
- self.norm1 = nn.LayerNorm(embed_dim)
- self.norm2 = nn.LayerNorm(embed_dim)
- self.dropout1 = nn.Dropout(dropout)
- self.dropout2 = nn.Dropout(dropout)
-
- self.activation = _get_activation_fn(activation)
-
- def with_pos_embed(self, tensor, pos: Optional[Tensor]):
- return tensor if pos is None else tensor + pos
-
- def forward(self, tgt,
- tgt_mask: Optional[Tensor] = None,
- tgt_key_padding_mask: Optional[Tensor] = None,
- query_pos: Optional[Tensor] = None):
- # self attention
- tgt2 = self.norm1(tgt)
- q = k = self.with_pos_embed(tgt2, query_pos)
- tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
- key_padding_mask=tgt_key_padding_mask)[0]
- tgt = tgt + self.dropout1(tgt2)
-
- # ffn
- tgt2 = self.norm2(tgt)
- tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
- tgt = tgt + self.dropout2(tgt2)
- return tgt
-
-class Fuse_sft_block(nn.Module):
- def __init__(self, in_ch, out_ch):
- super().__init__()
- self.encode_enc = ResBlock(2*in_ch, out_ch)
-
- self.scale = nn.Sequential(
- nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
- nn.LeakyReLU(0.2, True),
- nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1))
-
- self.shift = nn.Sequential(
- nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
- nn.LeakyReLU(0.2, True),
- nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1))
-
- def forward(self, enc_feat, dec_feat, w=1):
- enc_feat = self.encode_enc(torch.cat([enc_feat, dec_feat], dim=1))
- scale = self.scale(enc_feat)
- shift = self.shift(enc_feat)
- residual = w * (dec_feat * scale + shift)
- out = dec_feat + residual
- return out
-
-
-class CodeFormer(VQAutoEncoder):
- def __init__(self, dim_embd=512, n_head=8, n_layers=9,
- codebook_size=1024, latent_size=256,
- connect_list=('32', '64', '128', '256'),
- fix_modules=('quantize', 'generator')):
- super(CodeFormer, self).__init__(512, 64, [1, 2, 2, 4, 4, 8], 'nearest',2, [16], codebook_size)
-
- if fix_modules is not None:
- for module in fix_modules:
- for param in getattr(self, module).parameters():
- param.requires_grad = False
-
- self.connect_list = connect_list
- self.n_layers = n_layers
- self.dim_embd = dim_embd
- self.dim_mlp = dim_embd*2
-
- self.position_emb = nn.Parameter(torch.zeros(latent_size, self.dim_embd))
- self.feat_emb = nn.Linear(256, self.dim_embd)
-
- # transformer
- self.ft_layers = nn.Sequential(*[TransformerSALayer(embed_dim=dim_embd, nhead=n_head, dim_mlp=self.dim_mlp, dropout=0.0)
- for _ in range(self.n_layers)])
-
- # logits_predict head
- self.idx_pred_layer = nn.Sequential(
- nn.LayerNorm(dim_embd),
- nn.Linear(dim_embd, codebook_size, bias=False))
-
- self.channels = {
- '16': 512,
- '32': 256,
- '64': 256,
- '128': 128,
- '256': 128,
- '512': 64,
- }
-
- # after second residual block for > 16, before attn layer for ==16
- self.fuse_encoder_block = {'512':2, '256':5, '128':8, '64':11, '32':14, '16':18}
- # after first residual block for > 16, before attn layer for ==16
- self.fuse_generator_block = {'16':6, '32': 9, '64':12, '128':15, '256':18, '512':21}
-
- # fuse_convs_dict
- self.fuse_convs_dict = nn.ModuleDict()
- for f_size in self.connect_list:
- in_ch = self.channels[f_size]
- self.fuse_convs_dict[f_size] = Fuse_sft_block(in_ch, in_ch)
-
- def _init_weights(self, module):
- if isinstance(module, (nn.Linear, nn.Embedding)):
- module.weight.data.normal_(mean=0.0, std=0.02)
- if isinstance(module, nn.Linear) and module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
- def forward(self, x, w=0, detach_16=True, code_only=False, adain=False):
- # ################### Encoder #####################
- enc_feat_dict = {}
- out_list = [self.fuse_encoder_block[f_size] for f_size in self.connect_list]
- for i, block in enumerate(self.encoder.blocks):
- x = block(x)
- if i in out_list:
- enc_feat_dict[str(x.shape[-1])] = x.clone()
-
- lq_feat = x
- # ################# Transformer ###################
- # quant_feat, codebook_loss, quant_stats = self.quantize(lq_feat)
- pos_emb = self.position_emb.unsqueeze(1).repeat(1,x.shape[0],1)
- # BCHW -> BC(HW) -> (HW)BC
- feat_emb = self.feat_emb(lq_feat.flatten(2).permute(2,0,1))
- query_emb = feat_emb
- # Transformer encoder
- for layer in self.ft_layers:
- query_emb = layer(query_emb, query_pos=pos_emb)
-
- # output logits
- logits = self.idx_pred_layer(query_emb) # (hw)bn
- logits = logits.permute(1,0,2) # (hw)bn -> b(hw)n
-
- if code_only: # for training stage II
- # logits doesn't need softmax before cross_entropy loss
- return logits, lq_feat
-
- # ################# Quantization ###################
- # if self.training:
- # quant_feat = torch.einsum('btn,nc->btc', [soft_one_hot, self.quantize.embedding.weight])
- # # b(hw)c -> bc(hw) -> bchw
- # quant_feat = quant_feat.permute(0,2,1).view(lq_feat.shape)
- # ------------
- soft_one_hot = F.softmax(logits, dim=2)
- _, top_idx = torch.topk(soft_one_hot, 1, dim=2)
- quant_feat = self.quantize.get_codebook_feat(top_idx, shape=[x.shape[0],16,16,256])
- # preserve gradients
- # quant_feat = lq_feat + (quant_feat - lq_feat).detach()
-
- if detach_16:
- quant_feat = quant_feat.detach() # for training stage III
- if adain:
- quant_feat = adaptive_instance_normalization(quant_feat, lq_feat)
-
- # ################## Generator ####################
- x = quant_feat
- fuse_list = [self.fuse_generator_block[f_size] for f_size in self.connect_list]
-
- for i, block in enumerate(self.generator.blocks):
- x = block(x)
- if i in fuse_list: # fuse after i-th block
- f_size = str(x.shape[-1])
- if w>0:
- x = self.fuse_convs_dict[f_size](enc_feat_dict[f_size].detach(), x, w)
- out = x
- # logits doesn't need softmax before cross_entropy loss
- return out, logits, lq_feat
diff --git a/modules/postprocess/gfpgan_model.py b/modules/postprocess/gfpgan_model.py
deleted file mode 100644
index 8763f411e..000000000
--- a/modules/postprocess/gfpgan_model.py
+++ /dev/null
@@ -1,115 +0,0 @@
-import os
-
-from installer import install
-from modules import paths, shared, devices, modelloader, errors
-
-model_dir = "GFPGAN"
-user_path = None
-model_path = os.path.join(paths.models_path, model_dir)
-model_url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
-have_gfpgan = False
-loaded_gfpgan_model = None
-
-
-def gfpgann():
- import facexlib
- import gfpgan # pylint: disable=unused-import
- global loaded_gfpgan_model # pylint: disable=global-statement
- if loaded_gfpgan_model is not None:
- loaded_gfpgan_model.gfpgan.to(devices.device)
- return loaded_gfpgan_model
- if gfpgan_constructor is None:
- return None
- models = modelloader.load_models(model_path, model_url, user_path, ext_filter="GFPGAN")
- if len(models) == 1 and "http" in models[0]:
- model_file = models[0]
- elif len(models) != 0:
- latest_file = max(models, key=os.path.getctime)
- model_file = latest_file
- else:
- shared.log.error(f"Model failed loading: type=GFPGAN model={model_file}")
- return None
- if hasattr(facexlib.detection.retinaface, 'device'):
- facexlib.detection.retinaface.device = devices.device
- model = gfpgan_constructor(model_path=model_file, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None, device=devices.device)
- loaded_gfpgan_model = model
- shared.log.info(f"Model loaded: type=GFPGAN model={model_file}")
- return model
-
-
-def send_model_to(model, device):
- model.gfpgan.to(device)
- model.face_helper.face_det.to(device)
- model.face_helper.face_parse.to(device)
-
-
-def gfpgan_fix_faces(np_image):
- model = gfpgann()
- if model is None:
- return np_image
-
- send_model_to(model, devices.device)
-
- np_image_bgr = np_image[:, :, ::-1]
- _cropped_faces, _restored_faces, gfpgan_output_bgr = model.enhance(np_image_bgr, has_aligned=False, only_center_face=False, paste_back=True)
- np_image = gfpgan_output_bgr[:, :, ::-1]
-
- model.face_helper.clean_all()
-
- if shared.opts.detailer_unload:
- send_model_to(model, devices.cpu)
-
- return np_image
-
-
-gfpgan_constructor = None
-
-
-def setup_model(dirname):
- try:
- if not os.path.exists(model_path):
- os.makedirs(model_path)
- except Exception:
- pass
- try:
- install('--no-build-isolation git+https://github.com/Disty0/BasicSR@23c1fb6f5c559ef5ce7ad657f2fa56e41b121754', 'basicsr')
- install('--no-build-isolation git+https://github.com/Disty0/GFPGAN@ae0f7e44fafe0ef4716f3c10067f8f379b74c21c', 'gfpgan')
- import gfpgan
- import facexlib
- import modules.detailer
-
- global user_path # pylint: disable=global-statement
- global have_gfpgan # pylint: disable=global-statement
- global gfpgan_constructor # pylint: disable=global-statement
- load_file_from_url_orig = gfpgan.utils.load_file_from_url
- facex_load_file_from_url_orig = facexlib.detection.load_file_from_url
- facex_load_file_from_url_orig2 = facexlib.parsing.load_file_from_url
-
- def my_load_file_from_url(**kwargs):
- return load_file_from_url_orig(**dict(kwargs, model_dir=model_path))
-
- def facex_load_file_from_url(**kwargs):
- return facex_load_file_from_url_orig(**dict(kwargs, save_dir=model_path, model_dir=None))
-
- def facex_load_file_from_url2(**kwargs):
- return facex_load_file_from_url_orig2(**dict(kwargs, save_dir=model_path, model_dir=None))
-
- gfpgan.utils.load_file_from_url = my_load_file_from_url
- facexlib.detection.load_file_from_url = facex_load_file_from_url
- facexlib.parsing.load_file_from_url = facex_load_file_from_url2
- user_path = dirname
- have_gfpgan = True
- gfpgan_constructor = gfpgan.GFPGANer
-
- class FaceRestorerGFPGAN(modules.detailer.Detailer):
- cmd_dir = model_path
-
- def name(self):
- return "GFPGAN"
-
- def restore(self, np_image, p=None): # pylint: disable=unused-argument
- return gfpgan_fix_faces(np_image)
-
- shared.face_restorers.append(FaceRestorerGFPGAN())
- except Exception as e:
- errors.log.error(f'GFPGan failed to initialize: {e}')
diff --git a/modules/postprocess/restorer.py b/modules/postprocess/restorer.py
deleted file mode 100644
index ca424f306..000000000
--- a/modules/postprocess/restorer.py
+++ /dev/null
@@ -1,61 +0,0 @@
-import time
-import cv2
-import numpy as np
-from modules import shared, devices
-
-
-face_helper = None
-
-
-def restore(np_image, name, session, strength): # pylint: disable=unused-argument
- t0 = time.time()
- global face_helper # pylint: disable=global-statement
- try:
- from modules.facelib.utils.face_restoration_helper import FaceRestoreHelper
- from modules.facelib.detection.retinaface import retinaface
- except Exception as e:
- shared.log.error(f"FaceRestorer error: {e}")
- return np_image
- if hasattr(retinaface, 'device'):
- retinaface.device = devices.device
- if face_helper is None:
- face_helper = FaceRestoreHelper(1, face_size=512, crop_ratio=(1, 1), det_model='retinaface_resnet50', save_ext='png', use_parse=True, device=devices.device)
-
- np_image = np_image[:, :, ::-1]
- original_resolution = np_image.shape[0:2]
- resolution = session.get_inputs()[0].shape[-2:]
-
- if face_helper is None or session is None:
- return np_image
- face_helper.clean_all()
- face_helper.read_image(np_image)
- face_helper.get_face_landmarks_5(only_center_face=False, eye_dist_threshold=5)
- face_helper.align_warp_face()
-
- detected_faces = len(face_helper.cropped_faces)
- for cropped_face in face_helper.cropped_faces:
- cropped_face = cv2.resize(cropped_face, resolution, interpolation=cv2.INTER_LANCZOS4)
- cropped_face = cropped_face.astype(np.float16)[:,:,::-1] / 255.0
- cropped_face = cropped_face.transpose((2, 0, 1))
- cropped_face = (cropped_face - 0.5) / 0.5
- cropped_face = np.expand_dims(cropped_face, axis=0).astype(np.float16)
- w = np.array([strength], dtype=np.double)
- if 'codeformer' in name:
- restored_face = session.run(None, {'x':cropped_face, 'w':w})[0][0]
- else:
- restored_face = session.run(None, {'input':cropped_face})[0][0]
- restored_face = (restored_face.transpose(1,2,0).clip(-1,1) + 1) * 0.5
- restored_face = (restored_face * 255)[:,:,::-1]
- restored_face = restored_face.clip(0, 255).astype('uint8')
- face_helper.add_restored_face(restored_face)
- face_helper.get_inverse_affine(None)
- restored_img = face_helper.paste_faces_to_input_image()
- restored_img = restored_img[:, :, ::-1]
- if original_resolution != restored_img.shape[0:2]:
- restored_img = cv2.resize(restored_img, (0, 0), fx=original_resolution[1]/restored_img.shape[1], fy=original_resolution[0]/restored_img.shape[0], interpolation=cv2.INTER_LANCZOS4)
-
- face_helper.clean_all()
- t1 = time.time()
- shared.log.info(f'Detailer: model="{name}" faces={detected_faces} strength={strength} time={t1-t0:.3f}')
-
- return restored_img
diff --git a/modules/postprocess/vqgan_arch.py b/modules/postprocess/vqgan_arch.py
deleted file mode 100644
index 87911f8cf..000000000
--- a/modules/postprocess/vqgan_arch.py
+++ /dev/null
@@ -1,435 +0,0 @@
-# this file is copied from CodeFormer repository. Please see comment in modules/codeformer_model.py
-
-'''
-VQGAN code, adapted from the original created by the Unleashing Transformers authors:
-https://github.com/samb-t/unleashing-transformers/blob/master/models/vqgan.py
-
-'''
-import torch
-from torch import nn
-import torch.nn.functional as F
-from basicsr.utils import get_root_logger
-from basicsr.utils.registry import ARCH_REGISTRY
-
-def normalize(in_channels):
- return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
-
-
-@torch.jit.script
-def swish(x):
- return x*torch.sigmoid(x)
-
-
-# Define VQVAE classes
-class VectorQuantizer(nn.Module):
- def __init__(self, codebook_size, emb_dim, beta):
- super(VectorQuantizer, self).__init__()
- self.codebook_size = codebook_size # number of embeddings
- self.emb_dim = emb_dim # dimension of embedding
- self.beta = beta # commitment cost used in loss term, beta * ||z_e(x)-sg[e]||^2
- self.embedding = nn.Embedding(self.codebook_size, self.emb_dim)
- self.embedding.weight.data.uniform_(-1.0 / self.codebook_size, 1.0 / self.codebook_size)
-
- def forward(self, z):
- # reshape z -> (batch, height, width, channel) and flatten
- z = z.permute(0, 2, 3, 1).contiguous()
- z_flattened = z.view(-1, self.emb_dim)
-
- # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
- d = (z_flattened ** 2).sum(dim=1, keepdim=True) + (self.embedding.weight**2).sum(1) - \
- 2 * torch.matmul(z_flattened, self.embedding.weight.t())
-
- mean_distance = torch.mean(d)
- # find closest encodings
- # min_encoding_indices = torch.argmin(d, dim=1).unsqueeze(1)
- min_encoding_scores, min_encoding_indices = torch.topk(d, 1, dim=1, largest=False)
- # [0-1], higher score, higher confidence
- min_encoding_scores = torch.exp(-min_encoding_scores/10)
-
- min_encodings = torch.zeros(min_encoding_indices.shape[0], self.codebook_size).to(z)
- min_encodings.scatter_(1, min_encoding_indices, 1)
-
- # get quantized latent vectors
- z_q = torch.matmul(min_encodings, self.embedding.weight).view(z.shape)
- # compute loss for embedding
- loss = torch.mean((z_q.detach()-z)**2) + self.beta * torch.mean((z_q - z.detach()) ** 2)
- # preserve gradients
- z_q = z + (z_q - z).detach()
-
- # perplexity
- e_mean = torch.mean(min_encodings, dim=0)
- perplexity = torch.exp(-torch.sum(e_mean * torch.log(e_mean + 1e-10)))
- # reshape back to match original input shape
- z_q = z_q.permute(0, 3, 1, 2).contiguous()
-
- return z_q, loss, {
- "perplexity": perplexity,
- "min_encodings": min_encodings,
- "min_encoding_indices": min_encoding_indices,
- "min_encoding_scores": min_encoding_scores,
- "mean_distance": mean_distance
- }
-
- def get_codebook_feat(self, indices, shape):
- # input indices: batch*token_num -> (batch*token_num)*1
- # shape: batch, height, width, channel
- indices = indices.view(-1,1)
- min_encodings = torch.zeros(indices.shape[0], self.codebook_size).to(indices)
- min_encodings.scatter_(1, indices, 1)
- # get quantized latent vectors
- z_q = torch.matmul(min_encodings.float(), self.embedding.weight)
-
- if shape is not None: # reshape back to match original input shape
- z_q = z_q.view(shape).permute(0, 3, 1, 2).contiguous()
-
- return z_q
-
-
-class GumbelQuantizer(nn.Module):
- def __init__(self, codebook_size, emb_dim, num_hiddens, straight_through=False, kl_weight=5e-4, temp_init=1.0):
- super().__init__()
- self.codebook_size = codebook_size # number of embeddings
- self.emb_dim = emb_dim # dimension of embedding
- self.straight_through = straight_through
- self.temperature = temp_init
- self.kl_weight = kl_weight
- self.proj = nn.Conv2d(num_hiddens, codebook_size, 1) # projects last encoder layer to quantized logits
- self.embed = nn.Embedding(codebook_size, emb_dim)
-
- def forward(self, z):
- hard = self.straight_through if self.training else True
-
- logits = self.proj(z)
-
- soft_one_hot = F.gumbel_softmax(logits, tau=self.temperature, dim=1, hard=hard)
-
- z_q = torch.einsum("b n h w, n d -> b d h w", soft_one_hot, self.embed.weight)
-
- # + kl divergence to the prior loss
- qy = F.softmax(logits, dim=1)
- diff = self.kl_weight * torch.sum(qy * torch.log(qy * self.codebook_size + 1e-10), dim=1).mean()
- min_encoding_indices = soft_one_hot.argmax(dim=1)
-
- return z_q, diff, {
- "min_encoding_indices": min_encoding_indices
- }
-
-
-class Downsample(nn.Module):
- def __init__(self, in_channels):
- super().__init__()
- self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
-
- def forward(self, x):
- pad = (0, 1, 0, 1)
- x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
- x = self.conv(x)
- return x
-
-
-class Upsample(nn.Module):
- def __init__(self, in_channels):
- super().__init__()
- self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
-
- def forward(self, x):
- x = F.interpolate(x, scale_factor=2.0, mode="nearest")
- x = self.conv(x)
-
- return x
-
-
-class ResBlock(nn.Module):
- def __init__(self, in_channels, out_channels=None):
- super(ResBlock, self).__init__()
- self.in_channels = in_channels
- self.out_channels = in_channels if out_channels is None else out_channels
- self.norm1 = normalize(in_channels)
- self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
- self.norm2 = normalize(out_channels)
- self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
- if self.in_channels != self.out_channels:
- self.conv_out = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
-
- def forward(self, x_in):
- x = x_in
- x = self.norm1(x)
- x = swish(x)
- x = self.conv1(x)
- x = self.norm2(x)
- x = swish(x)
- x = self.conv2(x)
- if self.in_channels != self.out_channels:
- x_in = self.conv_out(x_in)
-
- return x + x_in
-
-
-class AttnBlock(nn.Module):
- def __init__(self, in_channels):
- super().__init__()
- self.in_channels = in_channels
-
- self.norm = normalize(in_channels)
- self.q = torch.nn.Conv2d(
- in_channels,
- in_channels,
- kernel_size=1,
- stride=1,
- padding=0
- )
- self.k = torch.nn.Conv2d(
- in_channels,
- in_channels,
- kernel_size=1,
- stride=1,
- padding=0
- )
- self.v = torch.nn.Conv2d(
- in_channels,
- in_channels,
- kernel_size=1,
- stride=1,
- padding=0
- )
- self.proj_out = torch.nn.Conv2d(
- in_channels,
- in_channels,
- kernel_size=1,
- stride=1,
- padding=0
- )
-
- def forward(self, x):
- h_ = x
- h_ = self.norm(h_)
- q = self.q(h_)
- k = self.k(h_)
- v = self.v(h_)
-
- # compute attention
- b, c, h, w = q.shape
- q = q.reshape(b, c, h*w)
- q = q.permute(0, 2, 1)
- k = k.reshape(b, c, h*w)
- w_ = torch.bmm(q, k)
- w_ = w_ * (int(c)**(-0.5))
- w_ = F.softmax(w_, dim=2)
-
- # attend to values
- v = v.reshape(b, c, h*w)
- w_ = w_.permute(0, 2, 1)
- h_ = torch.bmm(v, w_)
- h_ = h_.reshape(b, c, h, w)
-
- h_ = self.proj_out(h_)
-
- return x+h_
-
-
-class Encoder(nn.Module):
- def __init__(self, in_channels, nf, emb_dim, ch_mult, num_res_blocks, resolution, attn_resolutions):
- super().__init__()
- self.nf = nf
- self.num_resolutions = len(ch_mult)
- self.num_res_blocks = num_res_blocks
- self.resolution = resolution
- self.attn_resolutions = attn_resolutions
-
- curr_res = self.resolution
- in_ch_mult = (1,)+tuple(ch_mult)
-
- blocks = []
- # initial convultion
- blocks.append(nn.Conv2d(in_channels, nf, kernel_size=3, stride=1, padding=1))
-
- # residual and downsampling blocks, with attention on smaller res (16x16)
- for i in range(self.num_resolutions):
- block_in_ch = nf * in_ch_mult[i]
- block_out_ch = nf * ch_mult[i]
- for _ in range(self.num_res_blocks):
- blocks.append(ResBlock(block_in_ch, block_out_ch))
- block_in_ch = block_out_ch
- if curr_res in attn_resolutions:
- blocks.append(AttnBlock(block_in_ch))
-
- if i != self.num_resolutions - 1:
- blocks.append(Downsample(block_in_ch))
- curr_res = curr_res // 2
-
- # non-local attention block
- blocks.append(ResBlock(block_in_ch, block_in_ch))
- blocks.append(AttnBlock(block_in_ch))
- blocks.append(ResBlock(block_in_ch, block_in_ch))
-
- # normalise and convert to latent size
- blocks.append(normalize(block_in_ch))
- blocks.append(nn.Conv2d(block_in_ch, emb_dim, kernel_size=3, stride=1, padding=1))
- self.blocks = nn.ModuleList(blocks)
-
- def forward(self, x):
- for block in self.blocks:
- x = block(x)
-
- return x
-
-
-class Generator(nn.Module):
- def __init__(self, nf, emb_dim, ch_mult, res_blocks, img_size, attn_resolutions):
- super().__init__()
- self.nf = nf
- self.ch_mult = ch_mult
- self.num_resolutions = len(self.ch_mult)
- self.num_res_blocks = res_blocks
- self.resolution = img_size
- self.attn_resolutions = attn_resolutions
- self.in_channels = emb_dim
- self.out_channels = 3
- block_in_ch = self.nf * self.ch_mult[-1]
- curr_res = self.resolution // 2 ** (self.num_resolutions-1)
-
- blocks = []
- # initial conv
- blocks.append(nn.Conv2d(self.in_channels, block_in_ch, kernel_size=3, stride=1, padding=1))
-
- # non-local attention block
- blocks.append(ResBlock(block_in_ch, block_in_ch))
- blocks.append(AttnBlock(block_in_ch))
- blocks.append(ResBlock(block_in_ch, block_in_ch))
-
- for i in reversed(range(self.num_resolutions)):
- block_out_ch = self.nf * self.ch_mult[i]
-
- for _ in range(self.num_res_blocks):
- blocks.append(ResBlock(block_in_ch, block_out_ch))
- block_in_ch = block_out_ch
-
- if curr_res in self.attn_resolutions:
- blocks.append(AttnBlock(block_in_ch))
-
- if i != 0:
- blocks.append(Upsample(block_in_ch))
- curr_res = curr_res * 2
-
- blocks.append(normalize(block_in_ch))
- blocks.append(nn.Conv2d(block_in_ch, self.out_channels, kernel_size=3, stride=1, padding=1))
-
- self.blocks = nn.ModuleList(blocks)
-
-
- def forward(self, x):
- for block in self.blocks:
- x = block(x)
-
- return x
-
-
-@ARCH_REGISTRY.register()
-class VQAutoEncoder(nn.Module):
- def __init__(self, img_size, nf, ch_mult, quantizer="nearest", res_blocks=2, attn_resolutions=None, codebook_size=1024, emb_dim=256,
- beta=0.25, gumbel_straight_through=False, gumbel_kl_weight=1e-8, model_path=None):
- super().__init__()
- logger = get_root_logger()
- self.in_channels = 3
- self.nf = nf
- self.n_blocks = res_blocks
- self.codebook_size = codebook_size
- self.embed_dim = emb_dim
- self.ch_mult = ch_mult
- self.resolution = img_size
- self.attn_resolutions = attn_resolutions or [16]
- self.quantizer_type = quantizer
- self.encoder = Encoder(
- self.in_channels,
- self.nf,
- self.embed_dim,
- self.ch_mult,
- self.n_blocks,
- self.resolution,
- self.attn_resolutions
- )
- if self.quantizer_type == "nearest":
- self.beta = beta #0.25
- self.quantize = VectorQuantizer(self.codebook_size, self.embed_dim, self.beta)
- elif self.quantizer_type == "gumbel":
- self.gumbel_num_hiddens = emb_dim
- self.straight_through = gumbel_straight_through
- self.kl_weight = gumbel_kl_weight
- self.quantize = GumbelQuantizer(
- self.codebook_size,
- self.embed_dim,
- self.gumbel_num_hiddens,
- self.straight_through,
- self.kl_weight
- )
- self.generator = Generator(
- self.nf,
- self.embed_dim,
- self.ch_mult,
- self.n_blocks,
- self.resolution,
- self.attn_resolutions
- )
-
- if model_path is not None:
- chkpt = torch.load(model_path, map_location='cpu')
- if 'params_ema' in chkpt:
- self.load_state_dict(torch.load(model_path, map_location='cpu')['params_ema'])
- logger.info(f'vqgan is loaded from: {model_path} [params_ema]')
- elif 'params' in chkpt:
- self.load_state_dict(torch.load(model_path, map_location='cpu')['params'])
- logger.info(f'vqgan is loaded from: {model_path} [params]')
- else:
- raise ValueError('Wrong params!')
-
-
- def forward(self, x):
- x = self.encoder(x)
- quant, codebook_loss, quant_stats = self.quantize(x)
- x = self.generator(quant)
- return x, codebook_loss, quant_stats
-
-
-
-# patch based discriminator
-@ARCH_REGISTRY.register()
-class VQGANDiscriminator(nn.Module):
- def __init__(self, nc=3, ndf=64, n_layers=4, model_path=None):
- super().__init__()
-
- layers = [nn.Conv2d(nc, ndf, kernel_size=4, stride=2, padding=1), nn.LeakyReLU(0.2, True)]
- ndf_mult = 1
- ndf_mult_prev = 1
- for n in range(1, n_layers): # gradually increase the number of filters
- ndf_mult_prev = ndf_mult
- ndf_mult = min(2 ** n, 8)
- layers += [
- nn.Conv2d(ndf * ndf_mult_prev, ndf * ndf_mult, kernel_size=4, stride=2, padding=1, bias=False),
- nn.BatchNorm2d(ndf * ndf_mult),
- nn.LeakyReLU(0.2, True)
- ]
-
- ndf_mult_prev = ndf_mult
- ndf_mult = min(2 ** n_layers, 8)
-
- layers += [
- nn.Conv2d(ndf * ndf_mult_prev, ndf * ndf_mult, kernel_size=4, stride=1, padding=1, bias=False),
- nn.BatchNorm2d(ndf * ndf_mult),
- nn.LeakyReLU(0.2, True)
- ]
-
- layers += [
- nn.Conv2d(ndf * ndf_mult, 1, kernel_size=4, stride=1, padding=1)] # output 1 channel prediction map
- self.main = nn.Sequential(*layers)
-
- if model_path is not None:
- chkpt = torch.load(model_path, map_location='cpu')
- if 'params_d' in chkpt:
- self.load_state_dict(torch.load(model_path, map_location='cpu')['params_d'])
- elif 'params' in chkpt:
- self.load_state_dict(torch.load(model_path, map_location='cpu')['params'])
- else:
- raise ValueError('Wrong params!')
-
- def forward(self, x):
- return self.main(x)
diff --git a/modules/postprocess/yolo.py b/modules/postprocess/yolo.py
index bb0fd4caf..61f5f5661 100644
--- a/modules/postprocess/yolo.py
+++ b/modules/postprocess/yolo.py
@@ -21,10 +21,6 @@ predefined = [ #
'https://huggingface.co/vladmandic/yolo-detailers/resolve/main/anzhc-eyes-seg.pt',
'https://huggingface.co/vladmandic/yolo-detailers/resolve/main/anzhc-face-1024-seg-8n.pt',
'https://huggingface.co/vladmandic/yolo-detailers/resolve/main/anzhc-head-seg-8n.pt',
- 'https://huggingface.co/netrunner-exe/Face-Upscalers-onnx/resolve/main/codeformer.fp16.onnx',
- 'https://huggingface.co/netrunner-exe/Face-Upscalers-onnx/resolve/main/restoreformer.fp16.onnx',
- 'https://huggingface.co/netrunner-exe/Face-Upscalers-onnx/resolve/main/GFPGANv1.4.fp16.onnx',
- 'https://huggingface.co/netrunner-exe/Face-Upscalers-onnx/resolve/main/GPEN-BFR-512.fp16.onnx',
]
load_lock = threading.Lock()
@@ -314,12 +310,6 @@ class YoloRestorer(Detailer):
shared.log.warning(f'Detailer: model="{name}" not loaded')
continue
- if name.endswith('.fp16'): # run gfpgan or codeformer directly and skip detailer processing
- from modules.postprocess import restorer
- np_image = restorer.restore(np_image, name, model, p.detailer_strength)
- image = Image.fromarray(np_image)
- continue
-
if image is None:
image = Image.fromarray(np_image)
items = self.predict(model, image)
diff --git a/modules/postprocessing.py b/modules/postprocessing.py
index 1f04905c5..b624e3a7e 100644
--- a/modules/postprocessing.py
+++ b/modules/postprocessing.py
@@ -100,7 +100,7 @@ def run_postprocessing(extras_mode, image, image_folder: List[tempfile.NamedTemp
return outputs, info, params
-def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_dir, show_extras_results, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, upscaling_resize_w, upscaling_resize_h, upscaling_crop, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility, save_output: bool = True):
+def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_dir, show_extras_results, upscaling_resize, upscaling_resize_w, upscaling_resize_h, upscaling_crop, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility, save_output: bool = True):
"""old handler for API"""
args = scripts_manager.scripts_postproc.create_args_for_run({
@@ -114,13 +114,6 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_
"upscaler_2_name": extras_upscaler_2,
"upscaler_2_visibility": extras_upscaler_2_visibility,
},
- "GFPGAN": {
- "gfpgan_visibility": gfpgan_visibility,
- },
- "CodeFormer": {
- "codeformer_visibility": codeformer_visibility,
- "codeformer_weight": codeformer_weight,
- },
})
return run_postprocessing(extras_mode, image, image_folder, input_dir, output_dir, show_extras_results, *args, save_output=save_output)
diff --git a/modules/processing.py b/modules/processing.py
index 32bbf3545..13f2c275c 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -3,7 +3,7 @@ import json
import time
import numpy as np
from PIL import Image, ImageOps
-from modules import shared, devices, errors, images, scripts_manager, memstats, script_callbacks, extra_networks, detailer, sd_models, sd_checkpoint, sd_vae, processing_helpers, timer, face_restoration
+from modules import shared, devices, errors, images, scripts_manager, memstats, script_callbacks, extra_networks, detailer, sd_models, sd_checkpoint, sd_vae, processing_helpers, timer
from modules.sd_hijack_hypertile import context_hypertile_vae, context_hypertile_unet
from modules.processing_class import StableDiffusionProcessing, StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, StableDiffusionProcessingControl, StableDiffusionProcessingVideo # pylint: disable=unused-import
from modules.processing_info import create_infotext
@@ -55,8 +55,6 @@ class Processed:
self.audio = audio
- self.restore_faces = p.restore_faces or False
- self.face_restoration_model = shared.opts.face_restoration_model if p.restore_faces else None
self.detailer = p.detailer_enabled or False
self.detailer_model = shared.opts.detailer_model if p.detailer_enabled else None
self.seed_resize_from_w = p.seed_resize_from_w
@@ -296,15 +294,6 @@ def process_samples(p: StableDiffusionProcessing, samples):
if not shared.state.interrupted and not shared.state.skipped:
- if p.restore_faces:
- p.ops.append('restore')
- if not p.do_not_save_samples and shared.opts.save_images_before_detailer:
- info = create_infotext(p, p.prompts, p.seeds, p.subseeds, index=i)
- images.save_image(Image.fromarray(sample), path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-restore")
- sample = face_restoration.restore_faces(sample, p)
- if sample is not None:
- image = Image.fromarray(sample)
-
if p.detailer_enabled:
p.ops.append('detailer')
if not p.do_not_save_samples and shared.opts.save_images_before_detailer:
diff --git a/modules/processing_class.py b/modules/processing_class.py
index 09c0bf5c1..a6f9fa5dd 100644
--- a/modules/processing_class.py
+++ b/modules/processing_class.py
@@ -57,7 +57,6 @@ class StableDiffusionProcessing:
# other
hidiffusion: bool = False,
do_not_reload_embeddings: bool = False,
- restore_faces: bool = False,
# detailer
detailer_enabled: bool = False,
detailer_prompt: str = '',
@@ -214,7 +213,6 @@ class StableDiffusionProcessing:
self.detailer_steps = detailer_steps
self.detailer_strength = detailer_strength
self.detailer_resolution = detailer_resolution
- self.restore_faces = restore_faces
self.init_images = init_images
self.init_control = init_control
self.resize_mode = resize_mode
diff --git a/modules/shared.py b/modules/shared.py
index 5893ca980..2237b3fd9 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -50,7 +50,6 @@ xformers_available = False
compiled_model_state = None
sd_upscalers = []
detailers = []
-face_restorers = []
yolo = None
tab_names = []
extra_networks: list[ExtraNetworksPage] = []
@@ -493,8 +492,6 @@ options_templates.update(options_section(('system-paths', "System Paths"), {
"embeddings_dir": OptionInfo(os.path.join(paths.models_path, 'embeddings'), "Folder with textual inversion embeddings", folder=True),
"control_dir": OptionInfo(os.path.join(paths.models_path, 'control'), "Folder with Control models", folder=True),
"yolo_dir": OptionInfo(os.path.join(paths.models_path, 'yolo'), "Folder with Yolo models", folder=True),
- "codeformer_models_path": OptionInfo(os.path.join(paths.models_path, 'Codeformer'), "Folder with codeformer models", folder=True),
- "gfpgan_models_path": OptionInfo(os.path.join(paths.models_path, 'GFPGAN'), "Folder with GFPGAN models", folder=True),
"esrgan_models_path": OptionInfo(os.path.join(paths.models_path, 'ESRGAN'), "Folder with ESRGAN models", folder=True),
"bsrgan_models_path": OptionInfo(os.path.join(paths.models_path, 'BSRGAN'), "Folder with BSRGAN models", folder=True),
"realesrgan_models_path": OptionInfo(os.path.join(paths.models_path, 'RealESRGAN'), "Folder with RealESRGAN models", folder=True),
@@ -662,9 +659,6 @@ options_templates.update(options_section(('postprocessing', "Postprocessing"), {
"postprocessing_sep_seedvt": OptionInfo("SeedVT
", "", gr.HTML),
"seedvt_cfg_scale": OptionInfo(3.5, "SeedVR CFG Scale", gr.Slider, {"minimum": 1, "maximum": 15, "step": 1}),
- "postprocessing_sep_face_restore": OptionInfo("Face Restore
", "", gr.HTML),
- "face_restoration_model": OptionInfo("None", "Face restoration", gr.Radio, lambda: {"choices": ['None'] + [x.name() for x in face_restorers]}),
- "code_former_weight": OptionInfo(0.2, "CodeFormer weight parameter", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}),
"postprocessing_sep_upscalers": OptionInfo("Upscaling
", "", gr.HTML),
"upscaler_unload": OptionInfo(False, "Unload upscaler after processing"),
diff --git a/pyproject.toml b/pyproject.toml
index 2ec740e67..6ff739fc8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,7 +16,6 @@ exclude = [
".vscode",
"modules/cfgzero",
- "modules/facelib",
"modules/flash_attn_triton_amd",
"modules/hidiffusion",
"modules/intel/ipex",
@@ -137,7 +136,6 @@ main.ignore-paths=[
"modules/control/proc",
"modules/control/units",
"modules/dml",
- "modules/facelib",
"modules/flash_attn_triton_amd",
"modules/ggml",
"modules/hidiffusion",
diff --git a/scripts/postprocessing_codeformer.py b/scripts/postprocessing_codeformer.py
deleted file mode 100644
index 062eb5a64..000000000
--- a/scripts/postprocessing_codeformer.py
+++ /dev/null
@@ -1,28 +0,0 @@
-from PIL import Image
-import numpy as np
-import gradio as gr
-from modules import scripts_postprocessing
-from modules.postprocess import codeformer_model
-
-
-class ScriptPostprocessingCodeFormer(scripts_postprocessing.ScriptPostprocessing):
- name = "CodeFormer"
- order = 3000
-
- def ui(self):
- with gr.Accordion('Restore faces: CodeFormer', open = False, elem_id="postprocess_codeformer_accordion"):
- with gr.Row():
- codeformer_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label="Strength", value=0.0, elem_id="extras_codeformer_visibility")
- codeformer_weight = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label="Weight", value=0.2, elem_id="extras_codeformer_weight")
- return { "codeformer_visibility": codeformer_visibility, "codeformer_weight": codeformer_weight }
-
- def process(self, pp: scripts_postprocessing.PostprocessedImage, codeformer_visibility, codeformer_weight): # pylint: disable=arguments-differ
- if codeformer_visibility == 0:
- return
- restored_img = codeformer_model.codeformer.restore(np.array(pp.image, dtype=np.uint8), w=codeformer_weight)
- res = Image.fromarray(restored_img)
- if codeformer_visibility < 1.0:
- res = Image.blend(pp.image, res, codeformer_visibility)
- pp.image = res
- pp.info["CodeFormer visibility"] = round(codeformer_visibility, 3)
- pp.info["CodeFormer weight"] = round(codeformer_weight, 3)
diff --git a/scripts/postprocessing_gfpgan.py b/scripts/postprocessing_gfpgan.py
deleted file mode 100644
index 9e42416b1..000000000
--- a/scripts/postprocessing_gfpgan.py
+++ /dev/null
@@ -1,29 +0,0 @@
-from PIL import Image
-import numpy as np
-import gradio as gr
-from modules import scripts_postprocessing
-
-
-class ScriptPostprocessingGfpGan(scripts_postprocessing.ScriptPostprocessing):
- name = "GFPGAN"
- order = 2000
-
- def ui(self):
- with gr.Accordion('Restore faces: GFPGan', open = False, elem_id="postprocess_gfpgan_accordion"):
- with gr.Row():
- gfpgan_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Strength", value=0, elem_id="extras_gfpgan_visibility")
- return { "gfpgan_visibility": gfpgan_visibility }
-
- def process(self, pp: scripts_postprocessing.PostprocessedImage, gfpgan_visibility): # pylint: disable=arguments-differ
- from installer import install
- install("facexlib")
- install("gfpgan")
- if gfpgan_visibility == 0:
- return
- from modules.postprocess import gfpgan_model
- restored_img = gfpgan_model.gfpgan_fix_faces(np.array(pp.image, dtype=np.uint8))
- res = Image.fromarray(restored_img)
- if gfpgan_visibility < 1.0:
- res = Image.blend(pp.image, res, gfpgan_visibility)
- pp.image = res
- pp.info["GFPGAN visibility"] = round(gfpgan_visibility, 3)
diff --git a/scripts/xyz/xyz_grid_classes.py b/scripts/xyz/xyz_grid_classes.py
index 44faeb503..a6c6678b7 100644
--- a/scripts/xyz/xyz_grid_classes.py
+++ b/scripts/xyz/xyz_grid_classes.py
@@ -246,7 +246,7 @@ axis_options = [
AxisOption("[Refine] Refiner steps", float, apply_field("refiner_steps")),
AxisOption("[Postprocess] Upscaler", str, apply_upscaler, cost=0.4, choices=lambda: [x.name for x in shared.sd_upscalers]),
AxisOption("[Postprocess] Context", str, apply_context, choices=lambda: ["Add with forward", "Remove with forward", "Add with backward", "Remove with backward"]),
- AxisOption("[Postprocess] Detailer", str, apply_detailer, fmt=format_value_add_label),
+ AxisOption("[Postprocess] Detailer", bool, apply_detailer, fmt=format_bool, choices=lambda: [False, True]),
AxisOption("[Postprocess] Detailer strength", str, apply_field("detailer_strength")),
AxisOption("[Quant] SDNQ quant mode", str, apply_sdnq_quant, cost=0.9, fmt=format_value_add_label, choices=lambda: ['none'] + sorted(shared.sdnq_quant_modes)),
AxisOption("[Quant] SDNQ quant mode TE", str, apply_sdnq_quant_te, cost=0.9, fmt=format_value_add_label, choices=lambda: ['none'] + sorted(shared.sdnq_quant_modes)),
diff --git a/scripts/xyz/xyz_grid_shared.py b/scripts/xyz/xyz_grid_shared.py
index b8e856014..d11c8bee4 100644
--- a/scripts/xyz/xyz_grid_shared.py
+++ b/scripts/xyz/xyz_grid_shared.py
@@ -284,17 +284,8 @@ def apply_context(p: processing.StableDiffusionProcessingTxt2Img, opt, x):
def apply_detailer(p, opt, x):
- opt = opt.lower()
- if opt == 'codeformer':
- is_active = True
- p.detailer_model = 'CodeFormer'
- elif opt == 'gfpgan':
- is_active = True
- p.detailer_model = 'GFPGAN'
- else:
- is_active = opt in ('true', 'yes', 'y', '1')
- p.detailer_enabled = is_active
- shared.log.debug(f'XYZ grid apply face-restore: "{x}"')
+ p.detailer_enabled = bool(opt)
+ shared.log.debug(f'XYZ grid apply detailer: "{x}"')
def apply_control(field):
diff --git a/scripts/xyz_grid_on.py b/scripts/xyz_grid_on.py
index 6e4c245f2..dfbd963b8 100644
--- a/scripts/xyz_grid_on.py
+++ b/scripts/xyz_grid_on.py
@@ -452,7 +452,6 @@ class Script(scripts_manager.Script):
def process_images(self, p, *args): # pylint: disable=W0221, W0613
if xyz_results_cache is not None and len(xyz_results_cache.images) > 0:
- p.restore_faces = False
p.detailer_enabled = False
p.color_corrections = None
# p.scripts = None
diff --git a/webui.py b/webui.py
index de6ecf832..1347b3387 100644
--- a/webui.py
+++ b/webui.py
@@ -100,11 +100,6 @@ def initialize():
shared.prompt_styles.reload()
timer.startup.record("styles")
- import modules.postprocess.codeformer_model as codeformer
- codeformer.setup_model(shared.opts.codeformer_models_path)
- sys.modules["modules.codeformer_model"] = codeformer
- import modules.postprocess.gfpgan_model as gfpgan
- gfpgan.setup_model(shared.opts.gfpgan_models_path)
import modules.postprocess.yolo as yolo
yolo.initialize()
timer.startup.record("detailer")