update face detector.

Author: sczhou

.gitignore

@@ -8,6 +8,7 @@ version.py
 # *.png
 # *.jpeg
 # *.jpg
+*.pt
 *.gif
 *.pth
 *.dat
@@ -121,4 +122,5 @@ venv.bak/
 .mypy_cache/
 
 # project
-results/
+results/
+*_old*

README.md

@@ -16,9 +16,10 @@ S-Lab, Nanyang Technological University
 
 ### Updates
 
+- **2022.07.29**: The face detector is upgraded with the family of `['YOLOv5', 'RetinaFace']`. :hugs: 
 - **2022.07.17**: The Colab demo of CodeFormer is available now. <a href="https://colab.research.google.com/drive/1m52PNveE4PBhYrecj34cnpEeiHcC5LTb?usp=sharing"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="google colab logo"></a>
-- **2022.07.16**:  Test code for face restoration is released. :blush:
-- **2022.06.21**:  This repo is created.
+- **2022.07.16**: Test code for face restoration is released. :blush:
+- **2022.06.21**: This repo is created.
 
 
 
@@ -54,16 +55,16 @@ source activate codeformer
 # install python dependencies
 pip3 install -r requirements.txt
 python basicsr/setup.py develop
-conda install -c conda-forge dlib
 ```
+<!-- conda install -c conda-forge dlib -->
 
 ### Quick Inference
 
 ##### Download Pre-trained Models:
-Download the dlib pretrained models from [[Google Drive](https://drive.google.com/drive/folders/1YCqeuNDGCsJBAm90eGh7M_WWKTt19yIY?usp=sharing) | [OneDrive](https://entuedu-my.sharepoint.com/:f:/g/personal/s200094_e_ntu_edu_sg/Em2BaKU2OjhDolr11ngbrUgBu8q6SPn8E0jW-AC7nJF0Ig?e=HkjYrF)] to the `weights/dlib` folder. 
+Download the facelib pretrained models from [[Google Drive](https://drive.google.com/drive/folders/1b_3qwrzY_kTQh0-SnBoGBgOrJ_PLZSKm?usp=sharing) | [OneDrive](https://entuedu-my.sharepoint.com/:f:/g/personal/s200094_e_ntu_edu_sg/EvDxR7FcAbZMp_MA9ouq7aQB8XTppMb3-T0uGZ_2anI2mg?e=DXsJFo)] to the `weights/facelib` folder. 
 You can download by run the following command OR manually download the pretrained models.
 ```
-python scripts/download_pretrained_models.py dlib
+python scripts/download_pretrained_models.py facelib
 ```
 
 Download the CodeFormer pretrained models from [[Google Drive](https://drive.google.com/drive/folders/1CNNByjHDFt0b95q54yMVp6Ifo5iuU6QS?usp=sharing) | [OneDrive](https://entuedu-my.sharepoint.com/:f:/g/personal/s200094_e_ntu_edu_sg/EoKFj4wo8cdIn2-TY2IV6CYBhZ0pIG4kUOeHdPR_A5nlbg?e=AO8UN9)] to the `weights/CodeFormer` folder. 
@@ -82,8 +83,7 @@ You can put the testing images in the `inputs/TestWhole` folder. If you would li
 python inference_codeformer.py --w 0.5 --has_aligned --test_path [input folder]
 
 # For the whole images
-# Please set `--upsample_num_times 2` when faces are small and failed detected
-python inference_codeformer.py --w 0.7 --upsample_num_times 1 --test_path [input folder]
+python inference_codeformer.py --w 0.7 --test_path [input folder]
 ```
 
 NOTE that *w* is in [0, 1]. Generally, smaller *w* tends to produce a higher-quality result, while larger *w* yields a higher-fidelity result. 

basicsr/archs/vqgan_arch.py

@@ -228,7 +228,7 @@ def forward(self, x):
 
 
 class Encoder(nn.Module):
-    def __init__(self, in_channels, nf, out_channels, ch_mult, num_res_blocks, resolution, attn_resolutions):
+    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)
@@ -264,7 +264,7 @@ def __init__(self, in_channels, nf, out_channels, ch_mult, num_res_blocks, resol
 
         # normalise and convert to latent size
         blocks.append(normalize(block_in_ch))
-        blocks.append(nn.Conv2d(block_in_ch, out_channels, kernel_size=3, stride=1, padding=1))
+        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):
@@ -275,7 +275,7 @@ def forward(self, x):
 
 
 class Generator(nn.Module):
-    def __init__(self, nf, ch_mult, res_blocks, img_size, attn_resolutions, emb_dim):
+    def __init__(self, nf, emb_dim, ch_mult, res_blocks, img_size, attn_resolutions):
         super().__init__()
         self.nf = nf 
         self.ch_mult = ch_mult 
@@ -362,7 +362,14 @@ def __init__(self, img_size, nf, ch_mult, quantizer="nearest", res_blocks=2, att
                 self.straight_through,
                 self.kl_weight
             )
-        self.generator = Generator(nf, ch_mult, res_blocks, img_size, attn_resolutions, emb_dim)
+        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')

basicsr/utils/face_util.py

@@ -0,0 +1,74 @@
+import os
+import torch
+from torch import nn
+from copy import deepcopy
+
+from facelib.utils import load_file_from_url
+from facelib.utils import download_pretrained_models
+from facelib.detection.yolov5face.models.common import Conv
+
+from .retinaface.retinaface import RetinaFace
+from .yolov5face.face_detector import YoloDetector
+
+
+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/xinntao/facexlib/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/xinntao/facexlib/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='weights/facelib', 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)
+        f_id = {'yolov5l-face.pth': '131578zMA6B2x8VQHyHfa6GEPtulMCNzV'}
+    elif model_name == 'YOLOv5n':
+        model = YoloDetector(config_name='facelib/detection/yolov5face/models/yolov5n.yaml', device=device)
+        f_id = {'yolov5n-face.pth': '1fhcpFvWZqghpGXjYPIne2sw1Fy4yhw6o'}
+    else:
+        raise NotImplementedError(f'{model_name} is not implemented.')
+
+    model_path = os.path.join('weights/facelib', list(f_id.keys())[0])
+    if not os.path.exists(model_path):
+        download_pretrained_models(file_ids=f_id, save_path_root='weights/facelib')
+
+    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