-
-
-
[](https://huggingface.co/spaces/sczhou/CodeFormer) [](https://replicate.com/sczhou/codeformer) [](https://openxlab.org.cn/apps/detail/ShangchenZhou/CodeFormer) 
-
-
-[Shangchen Zhou](https://shangchenzhou.com/), [Kelvin C.K. Chan](https://ckkelvinchan.github.io/), [Chongyi Li](https://li-chongyi.github.io/), [Chen Change Loy](https://www.mmlab-ntu.com/person/ccloy/)
-
-S-Lab, Nanyang Technological University
-
-
-
-
-:star: If CodeFormer is helpful to your images or projects, please help star this repo. Thanks! :hugs:
-
-
-### Update
-- **2023.07.20**: Integrated to :panda_face: [OpenXLab](https://openxlab.org.cn/apps). Try out online demo! [](https://openxlab.org.cn/apps/detail/ShangchenZhou/CodeFormer)
-- **2023.04.19**: :whale: Training codes and config files are public available now.
-- **2023.04.09**: Add features of inpainting and colorization for cropped and aligned face images.
-- **2023.02.10**: Include `dlib` as a new face detector option, it produces more accurate face identity.
-- **2022.10.05**: Support video input `--input_path [YOUR_VIDEO.mp4]`. Try it to enhance your videos! :clapper:
-- **2022.09.14**: Integrated to :hugs: [Hugging Face](https://huggingface.co/spaces). Try out online demo! [](https://huggingface.co/spaces/sczhou/CodeFormer)
-- **2022.09.09**: Integrated to :rocket: [Replicate](https://replicate.com/explore). Try out online demo! [](https://replicate.com/sczhou/codeformer)
-- [**More**](docs/history_changelog.md)
-
-### TODO
-- [x] Add training code and config files
-- [x] Add checkpoint and script for face inpainting
-- [x] Add checkpoint and script for face colorization
-- [x] ~~Add background image enhancement~~
-
-#### :panda_face: Try Enhancing Old Photos / Fixing AI-arts
-[
](https://imgsli.com/MTI3NTE2) [
](https://imgsli.com/MTI3NTE1) [
](https://imgsli.com/MTI3NTIw)
-
-#### Face Restoration
-
-
-
-
-#### Face Color Enhancement and Restoration
-
-
-
-#### Face Inpainting
-
-
-
-
-
-### Dependencies and Installation
-
-- Pytorch >= 1.7.1
-- CUDA >= 10.1
-- Other required packages in `requirements.txt`
+```bash
+ conda create -n hallo python=3.10
+ conda activate hallo
```
-# git clone this repository
-git clone https://github.com/sczhou/CodeFormer
-cd CodeFormer
-
-# create new anaconda env
-conda create -n codeformer python=3.8 -y
-conda activate codeformer
-
-# install python dependencies
-pip3 install -r requirements.txt
-python basicsr/setup.py develop
-conda install -c conda-forge dlib (only for face detection or cropping with dlib)
-```
-
-### Quick Inference
+Install packages with `pip`
-#### Download Pre-trained Models:
-Download the facelib and dlib pretrained models from [[Releases](https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0) | [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 manually download the pretrained models OR download by running the following command:
-```
-python scripts/download_pretrained_models.py facelib
-python scripts/download_pretrained_models.py dlib (only for dlib face detector)
-```
-
-Download the CodeFormer pretrained models from [[Releases](https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0) | [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. You can manually download the pretrained models OR download by running the following command:
-```
-python scripts/download_pretrained_models.py CodeFormer
+```bash
+ pip install torch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 --index-url https://download.pytorch.org/whl/cu118
+ pip install -r requirements.txt
```
-#### Prepare Testing Data:
-You can put the testing images in the `inputs/TestWhole` folder. If you would like to test on cropped and aligned faces, you can put them in the `inputs/cropped_faces` folder. You can get the cropped and aligned faces by running the following command:
+Besides, ffmpeg is also needed:
+```bash
+ apt-get install ffmpeg
```
-# you may need to install dlib via: conda install -c conda-forge dlib
-python scripts/crop_align_face.py -i [input folder] -o [output folder]
-```
-
-#### Testing:
-[Note] If you want to compare CodeFormer in your paper, please run the following command indicating `--has_aligned` (for cropped and aligned face), as the command for the whole image will involve a process of face-background fusion that may damage hair texture on the boundary, which leads to unfair comparison.
+### đĨ Download Pretrained Models
-Fidelity weight *w* lays in [0, 1]. Generally, smaller *w* tends to produce a higher-quality result, while larger *w* yields a higher-fidelity result. The results will be saved in the `results` folder.
+You can easily get all pretrained models required by inference from our [HuggingFace repo](https://huggingface.co/fudan-generative-ai/hallo2).
+Clone the pretrained models into `${PROJECT_ROOT}/pretrained_models` directory by cmd below:
-đ§đģ Face Restoration (cropped and aligned face)
-```
-# For cropped and aligned faces (512x512)
-python inference_codeformer.py -w 0.5 --has_aligned --input_path [image folder]|[image path]
+```shell
+git lfs install
+git clone https://huggingface.co/fudan-generative-ai/hallo2 pretrained_models
```
-:framed_picture: Whole Image Enhancement
-```
-# For whole image
-# Add '--bg_upsampler realesrgan' to enhance the background regions with Real-ESRGAN
-# Add '--face_upsample' to further upsample restorated face with Real-ESRGAN
-python inference_codeformer.py -w 0.7 --input_path [image folder]|[image path]
-```
+Or you can download them separately from their source repo:
+- [hallo2](https://huggingface.co/fudan-generative-ai/hallo2/blob/main/hallo2/net_g.pth): Our checkpoint of video super-resolution.
+- [facelib](https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0): pretrained face parse models
+- [realesrgan](https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/RealESRGAN_x2plus.pth): background upsample model
+- [CodeFormer](https://github.com/sczhou/CodeFormer/releases/download/v0.1.0): pretrained [Codeformer](https://github.com/sczhou/CodeFormer) model, it's optional to download it, only if you want to train our video super-resolution model from scratch
-:clapper: Video Enhancement
-```
-# For Windows/Mac users, please install ffmpeg first
-conda install -c conda-forge ffmpeg
-```
-```
-# For video clips
-# Video path should end with '.mp4'|'.mov'|'.avi'
-python inference_codeformer.py --bg_upsampler realesrgan --face_upsample -w 1.0 --input_path [video path]
-```
+Finally, these pretrained models should be organized as follows:
-đ Face Colorization (cropped and aligned face)
-```
-# For cropped and aligned faces (512x512)
-# Colorize black and white or faded photo
-python inference_colorization.py --input_path [image folder]|[image path]
+```text
+./pretrained_models/
+|-- CodeFormer/
+| |-- codeformer.pth
+| `-- vqgan_code1024.pth
+|-- facelib
+| |-- detection_mobilenet0.25_Final.pth
+| |-- detection_Resnet50_Final.pth
+| |-- parsing_parsenet.pth
+| |-- yolov5l-face.pth
+| `-- yolov5n-face.pth
+|-- hallo2
+| `-- net_g.pth
+`-- realesrgan
+ `-- RealESRGAN_x2plus.pth
```
-đ¨ Face Inpainting (cropped and aligned face)
-```
-# For cropped and aligned faces (512x512)
-# Inputs could be masked by white brush using an image editing app (e.g., Photoshop)
-# (check out the examples in inputs/masked_faces)
-python inference_inpainting.py --input_path [image folder]|[image path]
-```
-### Training:
-The training commands can be found in the documents: [English](docs/train.md) **|** [įŽäŊä¸æ](docs/train_CN.md).
+### đŽ Run Inference
+#### High-Resolution animation
+Simply to run the `scripts/video_sr.py` and pass `input_path` and `output_path`:
-### Citation
-If our work is useful for your research, please consider citing:
+```bash
+python scripts/video_sr.py --input_path [input_video] --output_path [output_dir] --bg_upsampler realesrgan --face_upsample -w 1 -s 4
+```
- @inproceedings{zhou2022codeformer,
- author = {Zhou, Shangchen and Chan, Kelvin C.K. and Li, Chongyi and Loy, Chen Change},
- title = {Towards Robust Blind Face Restoration with Codebook Lookup TransFormer},
- booktitle = {NeurIPS},
- year = {2022}
- }
+Animation results will be saved at `output_dir`.
-### License
+For more options:
-This project is licensed under NTU S-Lab License 1.0. Redistribution and use should follow this license.
+```shell
+usage: video_sr.py [-h] [-i INPUT_PATH] [-o OUTPUT_PATH] [-w FIDELITY_WEIGHT] [-s UPSCALE] [--has_aligned] [--only_center_face] [--draw_box]
+ [--detection_model DETECTION_MODEL] [--bg_upsampler BG_UPSAMPLER] [--face_upsample] [--bg_tile BG_TILE] [--suffix SUFFIX]
-### Acknowledgement
+options:
+ -h, --help show this help message and exit
+ -i INPUT_PATH, --input_path INPUT_PATH
+ Input video
+ -o OUTPUT_PATH, --output_path OUTPUT_PATH
+ Output folder.
+ -w FIDELITY_WEIGHT, --fidelity_weight FIDELITY_WEIGHT
+ Balance the quality and fidelity. Default: 0.5
+ -s UPSCALE, --upscale UPSCALE
+ The final upsampling scale of the image. Default: 2
+ --has_aligned Input are cropped and aligned faces. Default: False
+ --only_center_face Only restore the center face. Default: False
+ --draw_box Draw the bounding box for the detected faces. Default: False
+ --detection_model DETECTION_MODEL
+ Face detector. Optional: retinaface_resnet50, retinaface_mobile0.25, YOLOv5l, YOLOv5n. Default: retinaface_resnet50
+ --bg_upsampler BG_UPSAMPLER
+ Background upsampler. Optional: realesrgan
+ --face_upsample Face upsampler after enhancement. Default: False
+ --bg_tile BG_TILE Tile size for background sampler. Default: 400
+ --suffix SUFFIX Suffix of the restored faces. Default: None
+```
-This project is based on [BasicSR](https://github.com/XPixelGroup/BasicSR). Some codes are brought from [Unleashing Transformers](https://github.com/samb-t/unleashing-transformers), [YOLOv5-face](https://github.com/deepcam-cn/yolov5-face), and [FaceXLib](https://github.com/xinntao/facexlib). We also adopt [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN) to support background image enhancement. Thanks for their awesome works.
+## Training
+##### prepare data for training
+We use the VFHQ dataset for training, you can download from its [homepage](https://liangbinxie.github.io/projects/vfhq/). Then updata `dataroot_gt` in `./configs/train/video_sr.yaml`.
-### Contact
-If you have any questions, please feel free to reach me out at `shangchenzhou@gmail.com`.
+#### training
+Start training with the following command:
+```shell
+python -m torch.distributed.launch --nproc_per_node=8 --master_port=4652 \
+basicsr/train.py -opt ./configs/train/video_sr.yaml \
+--launcher pytorch
+```
diff --git a/assets/CodeFormer_logo.png b/assets/CodeFormer_logo.png
deleted file mode 100644
index 024cb724..00000000
Binary files a/assets/CodeFormer_logo.png and /dev/null differ
diff --git a/assets/color_enhancement_result1.png b/assets/color_enhancement_result1.png
deleted file mode 100644
index 34433db6..00000000
Binary files a/assets/color_enhancement_result1.png and /dev/null differ
diff --git a/assets/color_enhancement_result2.png b/assets/color_enhancement_result2.png
deleted file mode 100644
index 228690ac..00000000
Binary files a/assets/color_enhancement_result2.png and /dev/null differ
diff --git a/assets/framework.png b/assets/framework.png
new file mode 100755
index 00000000..2ed98dbb
Binary files /dev/null and b/assets/framework.png differ
diff --git a/assets/framework_1.jpg b/assets/framework_1.jpg
new file mode 100755
index 00000000..72de2f86
Binary files /dev/null and b/assets/framework_1.jpg differ
diff --git a/assets/framework_2.jpg b/assets/framework_2.jpg
new file mode 100755
index 00000000..bf9344c2
Binary files /dev/null and b/assets/framework_2.jpg differ
diff --git a/assets/imgsli_1.jpg b/assets/imgsli_1.jpg
deleted file mode 100644
index 313438a6..00000000
Binary files a/assets/imgsli_1.jpg and /dev/null differ
diff --git a/assets/imgsli_2.jpg b/assets/imgsli_2.jpg
deleted file mode 100644
index 42dd7f43..00000000
Binary files a/assets/imgsli_2.jpg and /dev/null differ
diff --git a/assets/imgsli_3.jpg b/assets/imgsli_3.jpg
deleted file mode 100644
index c3f67d9d..00000000
Binary files a/assets/imgsli_3.jpg and /dev/null differ
diff --git a/assets/inpainting_result1.png b/assets/inpainting_result1.png
deleted file mode 100644
index 2c6fa68a..00000000
Binary files a/assets/inpainting_result1.png and /dev/null differ
diff --git a/assets/inpainting_result2.png b/assets/inpainting_result2.png
deleted file mode 100644
index 2945f9f9..00000000
Binary files a/assets/inpainting_result2.png and /dev/null differ
diff --git a/assets/network.jpg b/assets/network.jpg
deleted file mode 100644
index 5aaa6bd1..00000000
Binary files a/assets/network.jpg and /dev/null differ
diff --git a/assets/restoration_result1.png b/assets/restoration_result1.png
deleted file mode 100644
index 8fd3b67e..00000000
Binary files a/assets/restoration_result1.png and /dev/null differ
diff --git a/assets/restoration_result2.png b/assets/restoration_result2.png
deleted file mode 100644
index a2ff2827..00000000
Binary files a/assets/restoration_result2.png and /dev/null differ
diff --git a/assets/restoration_result3.png b/assets/restoration_result3.png
deleted file mode 100644
index 022d7642..00000000
Binary files a/assets/restoration_result3.png and /dev/null differ
diff --git a/assets/restoration_result4.png b/assets/restoration_result4.png
deleted file mode 100644
index 5e965076..00000000
Binary files a/assets/restoration_result4.png and /dev/null differ
diff --git a/assets/wechat.jpeg b/assets/wechat.jpeg
new file mode 100755
index 00000000..f641fd9c
Binary files /dev/null and b/assets/wechat.jpeg differ
diff --git a/basicsr/VERSION b/basicsr/VERSION
old mode 100644
new mode 100755
diff --git a/basicsr/__init__.py b/basicsr/__init__.py
old mode 100644
new mode 100755
diff --git a/basicsr/archs/__init__.py b/basicsr/archs/__init__.py
old mode 100644
new mode 100755
diff --git a/basicsr/archs/arcface_arch.py b/basicsr/archs/arcface_arch.py
old mode 100644
new mode 100755
diff --git a/basicsr/archs/arch_util.py b/basicsr/archs/arch_util.py
old mode 100644
new mode 100755
diff --git a/basicsr/archs/codeformer_arch.py b/basicsr/archs/codeformer_arch.py
old mode 100644
new mode 100755
index de66937d..e74dea82
--- a/basicsr/archs/codeformer_arch.py
+++ b/basicsr/archs/codeformer_arch.py
@@ -9,6 +9,10 @@
from basicsr.utils import get_root_logger
from basicsr.utils.registry import ARCH_REGISTRY
+from einops import rearrange
+
+from tqdm import tqdm
+
def calc_mean_std(feat, eps=1e-5):
"""Calculate mean and std for adaptive_instance_normalization.
@@ -110,15 +114,25 @@ def __init__(self, embed_dim, nhead=8, dim_mlp=2048, dropout=0.0, activation="ge
self.dropout1 = nn.Dropout(dropout)
self.dropout2 = nn.Dropout(dropout)
+ self.temp_norm = nn.LayerNorm(embed_dim)
+ self.temp_dropout = nn.Dropout(dropout)
+ self.temp_attn = nn.MultiheadAttention(embed_dim, nhead, dropout=dropout)
+
+ self.temp_ffn_norm = nn.LayerNorm(embed_dim)
+ self.temp_linear1 = nn.Linear(embed_dim, dim_mlp)
+ self.temp_ffn_dropout = nn.Dropout(dropout)
+ self.temp_linear2 = nn.Linear(dim_mlp, embed_dim)
+
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,
+ def forward(self, tgt, video_length, batch_size,
tgt_mask: Optional[Tensor] = None,
tgt_key_padding_mask: Optional[Tensor] = None,
- query_pos: Optional[Tensor] = None):
+ query_pos: Optional[Tensor] = None
+ ):
# self attention
tgt2 = self.norm1(tgt)
@@ -131,6 +145,22 @@ def forward(self, tgt,
tgt2 = self.norm2(tgt)
tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
tgt = tgt + self.dropout2(tgt2)
+
+ # tmp attn
+ tgt = rearrange(tgt, "d (b f) c -> f (b d) c", f=video_length)
+ tgt2 = self.temp_norm(tgt)
+ query_pos = rearrange(query_pos, "d (b f) c -> f (b d) c", f=video_length)
+ q = k = self.with_pos_embed(tgt2, query_pos)
+ tgt2 = self.temp_attn(q, k, value=tgt2, attn_mask=tgt_mask,
+ key_padding_mask=tgt_key_padding_mask)[0]
+ tgt = tgt + self.temp_dropout(tgt2)
+ tgt = rearrange(tgt, "f (b d) c -> d (b f) c", b=batch_size)
+
+ # ffn
+ tgt2 = self.temp_ffn_norm(tgt)
+ tgt2 = self.temp_linear2(self.temp_ffn_dropout(self.activation(self.temp_linear1(tgt2))))
+ tgt = tgt + self.temp_ffn_dropout(tgt2)
+
return tgt
class Fuse_sft_block(nn.Module):
@@ -151,9 +181,9 @@ def __init__(self, in_ch, out_ch):
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
+ shift = self.shift(enc_feat)
+ out = w * (dec_feat * scale + shift) + dec_feat
+
return out
@@ -166,13 +196,13 @@ def __init__(self, dim_embd=512, n_head=8, n_layers=9,
super(CodeFormer, self).__init__(512, 64, [1, 2, 2, 4, 4, 8], 'nearest',2, [16], codebook_size)
if vqgan_path is not None:
- self.load_state_dict(
+ m, n = self.load_state_dict(
torch.load(vqgan_path, map_location='cpu')['params_ema'])
- if fix_modules is not None:
- for module in fix_modules:
- for param in getattr(self, module).parameters():
- param.requires_grad = False
+ # 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
@@ -221,6 +251,8 @@ def _init_weights(self, module):
module.weight.data.fill_(1.0)
def forward(self, x, w=0, detach_16=True, code_only=False, adain=False):
+ b, f, _, _, _ = x.shape
+ x = rearrange(x, "b f c h w -> (b f) c h w")
# ################### Encoder #####################
enc_feat_dict = {}
out_list = [self.fuse_encoder_block[f_size] for f_size in self.connect_list]
@@ -238,7 +270,7 @@ def forward(self, x, w=0, detach_16=True, code_only=False, adain=False):
query_emb = feat_emb
# Transformer encoder
for layer in self.ft_layers:
- query_emb = layer(query_emb, query_pos=pos_emb)
+ query_emb = layer(query_emb, query_pos=pos_emb, video_length=f, batch_size=b)
# output logits
logits = self.idx_pred_layer(query_emb) # (hw)bn
@@ -248,17 +280,10 @@ def forward(self, x, w=0, detach_16=True, code_only=False, adain=False):
# 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
@@ -269,12 +294,69 @@ def forward(self, x, w=0, detach_16=True, code_only=False, adain=False):
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)
+
+ x = rearrange(x, "(b f) c h w -> b f c h w", f=f)
out = x
# logits doesn't need softmax before cross_entropy loss
- return out, logits, lq_feat
\ No newline at end of file
+ return out, logits, lq_feat
+
+
+ def inference(self, x, w=0, detach_16=True, adain=False):
+ with torch.no_grad():
+ b, f, _, _, _ = x.shape
+ x = rearrange(x, "b f c h w -> (b f) c h w")
+ # ################### 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.detach().cpu().clone()
+
+ lq_feat = x.detach()
+ 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, video_length=f, batch_size=b)
+
+ # output logits
+ logits = self.idx_pred_layer(query_emb) # (hw)bn
+ logits = logits.permute(1,0,2) # (hw)bn -> b(hw)n
+
+
+ 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])
+
+
+ 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].to(x.device), x, w)
+
+ x = rearrange(x, "(b f) c h w -> b f c h w", f=f)
+ # logits doesn't need softmax before cross_entropy loss
+ return x, top_idx
+
\ No newline at end of file
diff --git a/basicsr/archs/rrdbnet_arch.py b/basicsr/archs/rrdbnet_arch.py
old mode 100644
new mode 100755
diff --git a/basicsr/archs/vgg_arch.py b/basicsr/archs/vgg_arch.py
old mode 100644
new mode 100755
diff --git a/basicsr/archs/vqgan_arch.py b/basicsr/archs/vqgan_arch.py
old mode 100644
new mode 100755
index 5ac69263..3a65de10
--- a/basicsr/archs/vqgan_arch.py
+++ b/basicsr/archs/vqgan_arch.py
@@ -11,6 +11,7 @@
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)
diff --git a/basicsr/data/__init__.py b/basicsr/data/__init__.py
old mode 100644
new mode 100755
diff --git a/basicsr/data/data_sampler.py b/basicsr/data/data_sampler.py
old mode 100644
new mode 100755
diff --git a/basicsr/data/data_util.py b/basicsr/data/data_util.py
old mode 100644
new mode 100755
diff --git a/basicsr/data/ffhq_blind_joint_dataset.py b/basicsr/data/ffhq_blind_joint_dataset.py
deleted file mode 100755
index 0dc845f7..00000000
--- a/basicsr/data/ffhq_blind_joint_dataset.py
+++ /dev/null
@@ -1,324 +0,0 @@
-import cv2
-import math
-import random
-import numpy as np
-import os.path as osp
-from scipy.io import loadmat
-import torch
-import torch.utils.data as data
-from torchvision.transforms.functional import (adjust_brightness, adjust_contrast,
- adjust_hue, adjust_saturation, normalize)
-from basicsr.data import gaussian_kernels as gaussian_kernels
-from basicsr.data.transforms import augment
-from basicsr.data.data_util import paths_from_folder
-from basicsr.utils import FileClient, get_root_logger, imfrombytes, img2tensor
-from basicsr.utils.registry import DATASET_REGISTRY
-
-@DATASET_REGISTRY.register()
-class FFHQBlindJointDataset(data.Dataset):
-
- def __init__(self, opt):
- super(FFHQBlindJointDataset, self).__init__()
- logger = get_root_logger()
- self.opt = opt
- # file client (io backend)
- self.file_client = None
- self.io_backend_opt = opt['io_backend']
-
- self.gt_folder = opt['dataroot_gt']
- self.gt_size = opt.get('gt_size', 512)
- self.in_size = opt.get('in_size', 512)
- assert self.gt_size >= self.in_size, 'Wrong setting.'
-
- self.mean = opt.get('mean', [0.5, 0.5, 0.5])
- self.std = opt.get('std', [0.5, 0.5, 0.5])
-
- self.component_path = opt.get('component_path', None)
- self.latent_gt_path = opt.get('latent_gt_path', None)
-
- if self.component_path is not None:
- self.crop_components = True
- self.components_dict = torch.load(self.component_path)
- self.eye_enlarge_ratio = opt.get('eye_enlarge_ratio', 1.4)
- self.nose_enlarge_ratio = opt.get('nose_enlarge_ratio', 1.1)
- self.mouth_enlarge_ratio = opt.get('mouth_enlarge_ratio', 1.3)
- else:
- self.crop_components = False
-
- if self.latent_gt_path is not None:
- self.load_latent_gt = True
- self.latent_gt_dict = torch.load(self.latent_gt_path)
- else:
- self.load_latent_gt = False
-
- if self.io_backend_opt['type'] == 'lmdb':
- self.io_backend_opt['db_paths'] = self.gt_folder
- if not self.gt_folder.endswith('.lmdb'):
- raise ValueError("'dataroot_gt' should end with '.lmdb', "f'but received {self.gt_folder}')
- with open(osp.join(self.gt_folder, 'meta_info.txt')) as fin:
- self.paths = [line.split('.')[0] for line in fin]
- else:
- self.paths = paths_from_folder(self.gt_folder)
-
- # perform corrupt
- self.use_corrupt = opt.get('use_corrupt', True)
- self.use_motion_kernel = False
- # self.use_motion_kernel = opt.get('use_motion_kernel', True)
-
- if self.use_motion_kernel:
- self.motion_kernel_prob = opt.get('motion_kernel_prob', 0.001)
- motion_kernel_path = opt.get('motion_kernel_path', 'basicsr/data/motion-blur-kernels-32.pth')
- self.motion_kernels = torch.load(motion_kernel_path)
-
- if self.use_corrupt:
- # degradation configurations
- self.blur_kernel_size = self.opt['blur_kernel_size']
- self.kernel_list = self.opt['kernel_list']
- self.kernel_prob = self.opt['kernel_prob']
- # Small degradation
- self.blur_sigma = self.opt['blur_sigma']
- self.downsample_range = self.opt['downsample_range']
- self.noise_range = self.opt['noise_range']
- self.jpeg_range = self.opt['jpeg_range']
- # Large degradation
- self.blur_sigma_large = self.opt['blur_sigma_large']
- self.downsample_range_large = self.opt['downsample_range_large']
- self.noise_range_large = self.opt['noise_range_large']
- self.jpeg_range_large = self.opt['jpeg_range_large']
-
- # print
- logger.info(f'Blur: blur_kernel_size {self.blur_kernel_size}, sigma: [{", ".join(map(str, self.blur_sigma))}]')
- logger.info(f'Downsample: downsample_range [{", ".join(map(str, self.downsample_range))}]')
- logger.info(f'Noise: [{", ".join(map(str, self.noise_range))}]')
- logger.info(f'JPEG compression: [{", ".join(map(str, self.jpeg_range))}]')
-
- # color jitter
- self.color_jitter_prob = opt.get('color_jitter_prob', None)
- self.color_jitter_pt_prob = opt.get('color_jitter_pt_prob', None)
- self.color_jitter_shift = opt.get('color_jitter_shift', 20)
- if self.color_jitter_prob is not None:
- logger.info(f'Use random color jitter. Prob: {self.color_jitter_prob}, shift: {self.color_jitter_shift}')
-
- # to gray
- self.gray_prob = opt.get('gray_prob', 0.0)
- if self.gray_prob is not None:
- logger.info(f'Use random gray. Prob: {self.gray_prob}')
- self.color_jitter_shift /= 255.
-
- @staticmethod
- def color_jitter(img, shift):
- """jitter color: randomly jitter the RGB values, in numpy formats"""
- jitter_val = np.random.uniform(-shift, shift, 3).astype(np.float32)
- img = img + jitter_val
- img = np.clip(img, 0, 1)
- return img
-
- @staticmethod
- def color_jitter_pt(img, brightness, contrast, saturation, hue):
- """jitter color: randomly jitter the brightness, contrast, saturation, and hue, in torch Tensor formats"""
- fn_idx = torch.randperm(4)
- for fn_id in fn_idx:
- if fn_id == 0 and brightness is not None:
- brightness_factor = torch.tensor(1.0).uniform_(brightness[0], brightness[1]).item()
- img = adjust_brightness(img, brightness_factor)
-
- if fn_id == 1 and contrast is not None:
- contrast_factor = torch.tensor(1.0).uniform_(contrast[0], contrast[1]).item()
- img = adjust_contrast(img, contrast_factor)
-
- if fn_id == 2 and saturation is not None:
- saturation_factor = torch.tensor(1.0).uniform_(saturation[0], saturation[1]).item()
- img = adjust_saturation(img, saturation_factor)
-
- if fn_id == 3 and hue is not None:
- hue_factor = torch.tensor(1.0).uniform_(hue[0], hue[1]).item()
- img = adjust_hue(img, hue_factor)
- return img
-
-
- def get_component_locations(self, name, status):
- components_bbox = self.components_dict[name]
- if status[0]: # hflip
- # exchange right and left eye
- tmp = components_bbox['left_eye']
- components_bbox['left_eye'] = components_bbox['right_eye']
- components_bbox['right_eye'] = tmp
- # modify the width coordinate
- components_bbox['left_eye'][0] = self.gt_size - components_bbox['left_eye'][0]
- components_bbox['right_eye'][0] = self.gt_size - components_bbox['right_eye'][0]
- components_bbox['nose'][0] = self.gt_size - components_bbox['nose'][0]
- components_bbox['mouth'][0] = self.gt_size - components_bbox['mouth'][0]
-
- locations_gt = {}
- locations_in = {}
- for part in ['left_eye', 'right_eye', 'nose', 'mouth']:
- mean = components_bbox[part][0:2]
- half_len = components_bbox[part][2]
- if 'eye' in part:
- half_len *= self.eye_enlarge_ratio
- elif part == 'nose':
- half_len *= self.nose_enlarge_ratio
- elif part == 'mouth':
- half_len *= self.mouth_enlarge_ratio
- loc = np.hstack((mean - half_len + 1, mean + half_len))
- loc = torch.from_numpy(loc).float()
- locations_gt[part] = loc
- loc_in = loc/(self.gt_size//self.in_size)
- locations_in[part] = loc_in
- return locations_gt, locations_in
-
-
- def __getitem__(self, index):
- if self.file_client is None:
- self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
-
- # load gt image
- gt_path = self.paths[index]
- name = osp.basename(gt_path)[:-4]
- img_bytes = self.file_client.get(gt_path)
- img_gt = imfrombytes(img_bytes, float32=True)
-
- # random horizontal flip
- img_gt, status = augment(img_gt, hflip=self.opt['use_hflip'], rotation=False, return_status=True)
-
- if self.load_latent_gt:
- if status[0]:
- latent_gt = self.latent_gt_dict['hflip'][name]
- else:
- latent_gt = self.latent_gt_dict['orig'][name]
-
- if self.crop_components:
- locations_gt, locations_in = self.get_component_locations(name, status)
-
- # generate in image
- img_in = img_gt
- if self.use_corrupt:
- # motion blur
- if self.use_motion_kernel and random.random() < self.motion_kernel_prob:
- m_i = random.randint(0,31)
- k = self.motion_kernels[f'{m_i:02d}']
- img_in = cv2.filter2D(img_in,-1,k)
-
- # gaussian blur
- kernel = gaussian_kernels.random_mixed_kernels(
- self.kernel_list,
- self.kernel_prob,
- self.blur_kernel_size,
- self.blur_sigma,
- self.blur_sigma,
- [-math.pi, math.pi],
- noise_range=None)
- img_in = cv2.filter2D(img_in, -1, kernel)
-
- # downsample
- scale = np.random.uniform(self.downsample_range[0], self.downsample_range[1])
- img_in = cv2.resize(img_in, (int(self.gt_size // scale), int(self.gt_size // scale)), interpolation=cv2.INTER_LINEAR)
-
- # noise
- if self.noise_range is not None:
- noise_sigma = np.random.uniform(self.noise_range[0] / 255., self.noise_range[1] / 255.)
- noise = np.float32(np.random.randn(*(img_in.shape))) * noise_sigma
- img_in = img_in + noise
- img_in = np.clip(img_in, 0, 1)
-
- # jpeg
- if self.jpeg_range is not None:
- jpeg_p = np.random.uniform(self.jpeg_range[0], self.jpeg_range[1])
- encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), int(jpeg_p)]
- _, encimg = cv2.imencode('.jpg', img_in * 255., encode_param)
- img_in = np.float32(cv2.imdecode(encimg, 1)) / 255.
-
- # resize to in_size
- img_in = cv2.resize(img_in, (self.in_size, self.in_size), interpolation=cv2.INTER_LINEAR)
-
-
- # generate in_large with large degradation
- img_in_large = img_gt
-
- if self.use_corrupt:
- # motion blur
- if self.use_motion_kernel and random.random() < self.motion_kernel_prob:
- m_i = random.randint(0,31)
- k = self.motion_kernels[f'{m_i:02d}']
- img_in_large = cv2.filter2D(img_in_large,-1,k)
-
- # gaussian blur
- kernel = gaussian_kernels.random_mixed_kernels(
- self.kernel_list,
- self.kernel_prob,
- self.blur_kernel_size,
- self.blur_sigma_large,
- self.blur_sigma_large,
- [-math.pi, math.pi],
- noise_range=None)
- img_in_large = cv2.filter2D(img_in_large, -1, kernel)
-
- # downsample
- scale = np.random.uniform(self.downsample_range_large[0], self.downsample_range_large[1])
- img_in_large = cv2.resize(img_in_large, (int(self.gt_size // scale), int(self.gt_size // scale)), interpolation=cv2.INTER_LINEAR)
-
- # noise
- if self.noise_range_large is not None:
- noise_sigma = np.random.uniform(self.noise_range_large[0] / 255., self.noise_range_large[1] / 255.)
- noise = np.float32(np.random.randn(*(img_in_large.shape))) * noise_sigma
- img_in_large = img_in_large + noise
- img_in_large = np.clip(img_in_large, 0, 1)
-
- # jpeg
- if self.jpeg_range_large is not None:
- jpeg_p = np.random.uniform(self.jpeg_range_large[0], self.jpeg_range_large[1])
- encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), int(jpeg_p)]
- _, encimg = cv2.imencode('.jpg', img_in_large * 255., encode_param)
- img_in_large = np.float32(cv2.imdecode(encimg, 1)) / 255.
-
- # resize to in_size
- img_in_large = cv2.resize(img_in_large, (self.in_size, self.in_size), interpolation=cv2.INTER_LINEAR)
-
-
- # random color jitter (only for lq)
- if self.color_jitter_prob is not None and (np.random.uniform() < self.color_jitter_prob):
- img_in = self.color_jitter(img_in, self.color_jitter_shift)
- img_in_large = self.color_jitter(img_in_large, self.color_jitter_shift)
- # random to gray (only for lq)
- if self.gray_prob and np.random.uniform() < self.gray_prob:
- img_in = cv2.cvtColor(img_in, cv2.COLOR_BGR2GRAY)
- img_in = np.tile(img_in[:, :, None], [1, 1, 3])
- img_in_large = cv2.cvtColor(img_in_large, cv2.COLOR_BGR2GRAY)
- img_in_large = np.tile(img_in_large[:, :, None], [1, 1, 3])
-
- # BGR to RGB, HWC to CHW, numpy to tensor
- img_in, img_in_large, img_gt = img2tensor([img_in, img_in_large, img_gt], bgr2rgb=True, float32=True)
-
- # random color jitter (pytorch version) (only for lq)
- if self.color_jitter_pt_prob is not None and (np.random.uniform() < self.color_jitter_pt_prob):
- brightness = self.opt.get('brightness', (0.5, 1.5))
- contrast = self.opt.get('contrast', (0.5, 1.5))
- saturation = self.opt.get('saturation', (0, 1.5))
- hue = self.opt.get('hue', (-0.1, 0.1))
- img_in = self.color_jitter_pt(img_in, brightness, contrast, saturation, hue)
- img_in_large = self.color_jitter_pt(img_in_large, brightness, contrast, saturation, hue)
-
- # round and clip
- img_in = np.clip((img_in * 255.0).round(), 0, 255) / 255.
- img_in_large = np.clip((img_in_large * 255.0).round(), 0, 255) / 255.
-
- # Set vgg range_norm=True if use the normalization here
- # normalize
- normalize(img_in, self.mean, self.std, inplace=True)
- normalize(img_in_large, self.mean, self.std, inplace=True)
- normalize(img_gt, self.mean, self.std, inplace=True)
-
- return_dict = {'in': img_in, 'in_large_de': img_in_large, 'gt': img_gt, 'gt_path': gt_path}
-
- if self.crop_components:
- return_dict['locations_in'] = locations_in
- return_dict['locations_gt'] = locations_gt
-
- if self.load_latent_gt:
- return_dict['latent_gt'] = latent_gt
-
- return return_dict
-
-
- def __len__(self):
- return len(self.paths)
diff --git a/basicsr/data/gaussian_kernels.py b/basicsr/data/gaussian_kernels.py
index 0ce57f0a..a7c05a33 100755
--- a/basicsr/data/gaussian_kernels.py
+++ b/basicsr/data/gaussian_kernels.py
@@ -632,16 +632,16 @@ def show_one_kernel():
plt.show()
-def show_plateau_kernel():
- import matplotlib.pyplot as plt
- kernel_size = 21
+# def show_plateau_kernel():
+# import matplotlib.pyplot as plt
+# kernel_size = 21
- kernel = plateau_type1(kernel_size, 2, 4, -math.pi / 8, 2, grid=None)
- kernel_norm = bivariate_isotropic_Gaussian(kernel_size, 5)
- kernel_gau = bivariate_generalized_Gaussian(
- kernel_size, 2, 4, -math.pi / 8, 2, grid=None)
- delta_h, delta_w = mass_center_shift(kernel_size, kernel)
- print(delta_h, delta_w)
+# kernel = plateau_type1(kernel_size, 2, 4, -math.pi / 8, 2, grid=None)
+# kernel_norm = bivariate_isotropic_Gaussian(kernel_size, 5)
+# kernel_gau = bivariate_generalized_Gaussian(
+# kernel_size, 2, 4, -math.pi / 8, 2, grid=None)
+# delta_h, delta_w = mass_center_shift(kernel_size, kernel)
+# print(delta_h, delta_w)
# kernel_slice = kernel[10, :]
# kernel_gau_slice = kernel_gau[10, :]
@@ -662,29 +662,29 @@ def show_plateau_kernel():
# ax.plot(t, y2)
# plt.show()
- fig, axs = plt.subplots(nrows=2, ncols=2)
- # axs.set_axis_off()
- ax = axs[0][0]
- im = ax.matshow(kernel, cmap='jet', origin='upper')
- fig.colorbar(im, ax=ax)
-
- # image
- ax = axs[0][1]
- kernel_vis = kernel - np.min(kernel)
- kernel_vis = kernel_vis / np.max(kernel_vis) * 255.
- ax.imshow(kernel_vis, interpolation='nearest')
-
- _, xx, yy = mesh_grid(kernel_size)
- # contour
- ax = axs[1][0]
- CS = ax.contour(xx, yy, kernel, origin='upper')
- ax.clabel(CS, inline=1, fontsize=3)
+ # fig, axs = plt.subplots(nrows=2, ncols=2)
+ # # axs.set_axis_off()
+ # ax = axs[0][0]
+ # im = ax.matshow(kernel, cmap='jet', origin='upper')
+ # fig.colorbar(im, ax=ax)
+
+ # # image
+ # ax = axs[0][1]
+ # kernel_vis = kernel - np.min(kernel)
+ # kernel_vis = kernel_vis / np.max(kernel_vis) * 255.
+ # ax.imshow(kernel_vis, interpolation='nearest')
+
+ # _, xx, yy = mesh_grid(kernel_size)
+ # # contour
+ # ax = axs[1][0]
+ # CS = ax.contour(xx, yy, kernel, origin='upper')
+ # ax.clabel(CS, inline=1, fontsize=3)
+
+ # # contourf
+ # ax = axs[1][1]
+ # kernel = kernel / np.max(kernel)
+ # p = ax.contourf(
+ # xx, yy, kernel, origin='upper', levels=np.linspace(-0.05, 1.05, 10))
+ # fig.colorbar(p)
- # contourf
- ax = axs[1][1]
- kernel = kernel / np.max(kernel)
- p = ax.contourf(
- xx, yy, kernel, origin='upper', levels=np.linspace(-0.05, 1.05, 10))
- fig.colorbar(p)
-
- plt.show()
+ # plt.show()
diff --git a/basicsr/data/paired_image_dataset.py b/basicsr/data/paired_image_dataset.py
deleted file mode 100644
index c6a6c07b..00000000
--- a/basicsr/data/paired_image_dataset.py
+++ /dev/null
@@ -1,101 +0,0 @@
-from torch.utils import data as data
-from torchvision.transforms.functional import normalize
-
-from basicsr.data.data_util import paired_paths_from_folder, paired_paths_from_lmdb, paired_paths_from_meta_info_file
-from basicsr.data.transforms import augment, paired_random_crop
-from basicsr.utils import FileClient, imfrombytes, img2tensor
-from basicsr.utils.registry import DATASET_REGISTRY
-
-
-@DATASET_REGISTRY.register()
-class PairedImageDataset(data.Dataset):
- """Paired image dataset for image restoration.
-
- Read LQ (Low Quality, e.g. LR (Low Resolution), blurry, noisy, etc) and
- GT image pairs.
-
- There are three modes:
- 1. 'lmdb': Use lmdb files.
- If opt['io_backend'] == lmdb.
- 2. 'meta_info_file': Use meta information file to generate paths.
- If opt['io_backend'] != lmdb and opt['meta_info_file'] is not None.
- 3. 'folder': Scan folders to generate paths.
- The rest.
-
- Args:
- opt (dict): Config for train datasets. It contains the following keys:
- dataroot_gt (str): Data root path for gt.
- dataroot_lq (str): Data root path for lq.
- meta_info_file (str): Path for meta information file.
- io_backend (dict): IO backend type and other kwarg.
- filename_tmpl (str): Template for each filename. Note that the
- template excludes the file extension. Default: '{}'.
- gt_size (int): Cropped patched size for gt patches.
- use_flip (bool): Use horizontal flips.
- use_rot (bool): Use rotation (use vertical flip and transposing h
- and w for implementation).
-
- scale (bool): Scale, which will be added automatically.
- phase (str): 'train' or 'val'.
- """
-
- def __init__(self, opt):
- super(PairedImageDataset, self).__init__()
- self.opt = opt
- # file client (io backend)
- self.file_client = None
- self.io_backend_opt = opt['io_backend']
- self.mean = opt['mean'] if 'mean' in opt else None
- self.std = opt['std'] if 'std' in opt else None
-
- self.gt_folder, self.lq_folder = opt['dataroot_gt'], opt['dataroot_lq']
- if 'filename_tmpl' in opt:
- self.filename_tmpl = opt['filename_tmpl']
- else:
- self.filename_tmpl = '{}'
-
- if self.io_backend_opt['type'] == 'lmdb':
- self.io_backend_opt['db_paths'] = [self.lq_folder, self.gt_folder]
- self.io_backend_opt['client_keys'] = ['lq', 'gt']
- self.paths = paired_paths_from_lmdb([self.lq_folder, self.gt_folder], ['lq', 'gt'])
- elif 'meta_info_file' in self.opt and self.opt['meta_info_file'] is not None:
- self.paths = paired_paths_from_meta_info_file([self.lq_folder, self.gt_folder], ['lq', 'gt'],
- self.opt['meta_info_file'], self.filename_tmpl)
- else:
- self.paths = paired_paths_from_folder([self.lq_folder, self.gt_folder], ['lq', 'gt'], self.filename_tmpl)
-
- def __getitem__(self, index):
- if self.file_client is None:
- self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
-
- scale = self.opt['scale']
-
- # Load gt and lq images. Dimension order: HWC; channel order: BGR;
- # image range: [0, 1], float32.
- gt_path = self.paths[index]['gt_path']
- img_bytes = self.file_client.get(gt_path, 'gt')
- img_gt = imfrombytes(img_bytes, float32=True)
- lq_path = self.paths[index]['lq_path']
- img_bytes = self.file_client.get(lq_path, 'lq')
- img_lq = imfrombytes(img_bytes, float32=True)
-
- # augmentation for training
- if self.opt['phase'] == 'train':
- gt_size = self.opt['gt_size']
- # random crop
- img_gt, img_lq = paired_random_crop(img_gt, img_lq, gt_size, scale, gt_path)
- # flip, rotation
- img_gt, img_lq = augment([img_gt, img_lq], self.opt['use_flip'], self.opt['use_rot'])
-
- # TODO: color space transform
- # BGR to RGB, HWC to CHW, numpy to tensor
- img_gt, img_lq = img2tensor([img_gt, img_lq], bgr2rgb=True, float32=True)
- # normalize
- if self.mean is not None or self.std is not None:
- normalize(img_lq, self.mean, self.std, inplace=True)
- normalize(img_gt, self.mean, self.std, inplace=True)
-
- return {'lq': img_lq, 'gt': img_gt, 'lq_path': lq_path, 'gt_path': gt_path}
-
- def __len__(self):
- return len(self.paths)
diff --git a/basicsr/data/prefetch_dataloader.py b/basicsr/data/prefetch_dataloader.py
old mode 100644
new mode 100755
diff --git a/basicsr/data/transforms.py b/basicsr/data/transforms.py
old mode 100644
new mode 100755
diff --git a/basicsr/data/ffhq_blind_dataset.py b/basicsr/data/vfhq_dataset.py
similarity index 57%
rename from basicsr/data/ffhq_blind_dataset.py
rename to basicsr/data/vfhq_dataset.py
index 9f900606..eadbdeb5 100755
--- a/basicsr/data/ffhq_blind_dataset.py
+++ b/basicsr/data/vfhq_dataset.py
@@ -15,16 +15,20 @@
from basicsr.utils import FileClient, get_root_logger, imfrombytes, img2tensor
from basicsr.utils.registry import DATASET_REGISTRY
+from pathlib import Path
+import torchvision.transforms as transforms
+
@DATASET_REGISTRY.register()
-class FFHQBlindDataset(data.Dataset):
+class VFHQBlindDataset(data.Dataset):
def __init__(self, opt):
- super(FFHQBlindDataset, self).__init__()
+ super(VFHQBlindDataset, self).__init__()
logger = get_root_logger()
self.opt = opt
# file client (io backend)
self.file_client = None
self.io_backend_opt = opt['io_backend']
+ self.video_length = opt['video_length']
self.gt_folder = opt['dataroot_gt']
self.gt_size = opt.get('gt_size', 512)
@@ -59,26 +63,25 @@ def __init__(self, opt):
with open(osp.join(self.gt_folder, 'meta_info.txt')) as fin:
self.paths = [line.split('.')[0] for line in fin]
else:
- self.paths = paths_from_folder(self.gt_folder)
+ gt_folder = Path(self.gt_folder)
+ sub_dir = gt_folder.iterdir()
+ self.paths = []
+ for p in sub_dir:
+ if p.is_dir():
+ l = list(p.glob('*.png'))
+ if len(l) > self.video_length:
+ self.paths.append(str(p))
+
# inpainting mask
self.gen_inpaint_mask = opt.get('gen_inpaint_mask', False)
if self.gen_inpaint_mask:
logger.info(f'generate mask ...')
- # self.mask_max_angle = opt.get('mask_max_angle', 10)
- # self.mask_max_len = opt.get('mask_max_len', 150)
- # self.mask_max_width = opt.get('mask_max_width', 50)
- # self.mask_draw_times = opt.get('mask_draw_times', 10)
- # # print
- # logger.info(f'mask_max_angle: {self.mask_max_angle}')
- # logger.info(f'mask_max_len: {self.mask_max_len}')
- # logger.info(f'mask_max_width: {self.mask_max_width}')
- # logger.info(f'mask_draw_times: {self.mask_draw_times}')
+
# perform corrupt
self.use_corrupt = opt.get('use_corrupt', True)
self.use_motion_kernel = False
- # self.use_motion_kernel = opt.get('use_motion_kernel', True)
if self.use_motion_kernel:
self.motion_kernel_prob = opt.get('motion_kernel_prob', 0.001)
@@ -113,6 +116,7 @@ def __init__(self, opt):
logger.info(f'Use random gray. Prob: {self.gray_prob}')
self.color_jitter_shift /= 255.
+
@staticmethod
def color_jitter(img, shift):
"""jitter color: randomly jitter the RGB values, in numpy formats"""
@@ -182,118 +186,113 @@ def __getitem__(self, index):
# load gt image
gt_path = self.paths[index]
- name = osp.basename(gt_path)[:-4]
- img_bytes = self.file_client.get(gt_path)
- img_gt = imfrombytes(img_bytes, float32=True)
-
- # random horizontal flip
- img_gt, status = augment(img_gt, hflip=self.opt['use_hflip'], rotation=False, return_status=True)
-
- if self.load_latent_gt:
- if status[0]:
- latent_gt = self.latent_gt_dict['hflip'][name]
- else:
- latent_gt = self.latent_gt_dict['orig'][name]
-
- if self.crop_components:
- locations_gt, locations_in = self.get_component_locations(name, status)
-
- # generate in image
- img_in = img_gt
- if self.use_corrupt and not self.gen_inpaint_mask:
- # motion blur
- if self.use_motion_kernel and random.random() < self.motion_kernel_prob:
- m_i = random.randint(0,31)
- k = self.motion_kernels[f'{m_i:02d}']
- img_in = cv2.filter2D(img_in,-1,k)
-
- # gaussian blur
- kernel = gaussian_kernels.random_mixed_kernels(
- self.kernel_list,
- self.kernel_prob,
- self.blur_kernel_size,
- self.blur_sigma,
- self.blur_sigma,
- [-math.pi, math.pi],
- noise_range=None)
- img_in = cv2.filter2D(img_in, -1, kernel)
-
- # downsample
- scale = np.random.uniform(self.downsample_range[0], self.downsample_range[1])
- img_in = cv2.resize(img_in, (int(self.gt_size // scale), int(self.gt_size // scale)), interpolation=cv2.INTER_LINEAR)
-
- # noise
- if self.noise_range is not None:
- noise_sigma = np.random.uniform(self.noise_range[0] / 255., self.noise_range[1] / 255.)
- noise = np.float32(np.random.randn(*(img_in.shape))) * noise_sigma
- img_in = img_in + noise
- img_in = np.clip(img_in, 0, 1)
-
- # jpeg
- if self.jpeg_range is not None:
- jpeg_p = np.random.uniform(self.jpeg_range[0], self.jpeg_range[1])
- encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), int(jpeg_p)]
- _, encimg = cv2.imencode('.jpg', img_in * 255., encode_param)
- img_in = np.float32(cv2.imdecode(encimg, 1)) / 255.
-
- # resize to in_size
- img_in = cv2.resize(img_in, (self.in_size, self.in_size), interpolation=cv2.INTER_LINEAR)
-
- # if self.gen_inpaint_mask:
- # inpaint_mask = random_ff_mask(shape=(self.gt_size,self.gt_size),
- # max_angle = self.mask_max_angle, max_len = self.mask_max_len,
- # max_width = self.mask_max_width, times = self.mask_draw_times)
- # img_in = img_in * (1 - inpaint_mask.reshape(self.gt_size,self.gt_size,1)) + \
- # 1.0 * inpaint_mask.reshape(self.gt_size,self.gt_size,1)
-
- # inpaint_mask = torch.from_numpy(inpaint_mask).view(1,self.gt_size,self.gt_size)
-
- if self.gen_inpaint_mask:
- img_in = (img_in*255).astype('uint8')
- img_in = brush_stroke_mask(Image.fromarray(img_in))
- img_in = np.array(img_in) / 255.
-
- # random color jitter (only for lq)
- if self.color_jitter_prob is not None and (np.random.uniform() < self.color_jitter_prob):
- img_in = self.color_jitter(img_in, self.color_jitter_shift)
- # random to gray (only for lq)
- if self.gray_prob and np.random.uniform() < self.gray_prob:
- img_in = cv2.cvtColor(img_in, cv2.COLOR_BGR2GRAY)
- img_in = np.tile(img_in[:, :, None], [1, 1, 3])
-
- # BGR to RGB, HWC to CHW, numpy to tensor
- img_in, img_gt = img2tensor([img_in, img_gt], bgr2rgb=True, float32=True)
- # random color jitter (pytorch version) (only for lq)
- if self.color_jitter_pt_prob is not None and (np.random.uniform() < self.color_jitter_pt_prob):
- brightness = self.opt.get('brightness', (0.5, 1.5))
- contrast = self.opt.get('contrast', (0.5, 1.5))
- saturation = self.opt.get('saturation', (0, 1.5))
- hue = self.opt.get('hue', (-0.1, 0.1))
- img_in = self.color_jitter_pt(img_in, brightness, contrast, saturation, hue)
+ image_list = list(Path(gt_path).glob('*.png'))
+ lenght = len(image_list)
- # round and clip
- img_in = np.clip((img_in * 255.0).round(), 0, 255) / 255.
+ start_idx = random.randint(0, lenght-self.video_length-1)
+ in_list = []
+ gt_list = []
- # Set vgg range_norm=True if use the normalization here
- # normalize
- normalize(img_in, self.mean, self.std, inplace=True)
- normalize(img_gt, self.mean, self.std, inplace=True)
+ for i in range(start_idx, start_idx+self.video_length):
+ gt_path_idx = image_list[i]
- return_dict = {'in': img_in, 'gt': img_gt, 'gt_path': gt_path}
-
- if self.crop_components:
- return_dict['locations_in'] = locations_in
- return_dict['locations_gt'] = locations_gt
-
- if self.load_latent_gt:
- return_dict['latent_gt'] = latent_gt
-
- # if self.gen_inpaint_mask:
- # return_dict['inpaint_mask'] = inpaint_mask
+ img_bytes = self.file_client.get(gt_path_idx)
+ img_gt = imfrombytes(img_bytes, float32=True)
+
+ # random horizontal flip
+ img_gt, status = augment(img_gt, hflip=self.opt['use_hflip'], rotation=False, return_status=True)
+
+
+ # generate in image
+ img_in = img_gt
+ if self.use_corrupt and not self.gen_inpaint_mask:
+ # motion blur
+ if self.use_motion_kernel and random.random() < self.motion_kernel_prob:
+ m_i = random.randint(0,31)
+ k = self.motion_kernels[f'{m_i:02d}']
+ img_in = cv2.filter2D(img_in,-1,k)
+
+ # gaussian blur
+ kernel = gaussian_kernels.random_mixed_kernels(
+ self.kernel_list,
+ self.kernel_prob,
+ self.blur_kernel_size,
+ self.blur_sigma,
+ self.blur_sigma,
+ [-math.pi, math.pi],
+ noise_range=None)
+ img_in = cv2.filter2D(img_in, -1, kernel)
+
+ # downsample
+ scale = np.random.uniform(self.downsample_range[0], self.downsample_range[1])
+ img_in = cv2.resize(img_in, (int(self.gt_size // scale), int(self.gt_size // scale)), interpolation=cv2.INTER_LINEAR)
+
+ # noise
+ if self.noise_range is not None:
+ noise_sigma = np.random.uniform(self.noise_range[0] / 255., self.noise_range[1] / 255.)
+ noise = np.float32(np.random.randn(*(img_in.shape))) * noise_sigma
+ img_in = img_in + noise
+ img_in = np.clip(img_in, 0, 1)
+
+ # jpeg
+ if self.jpeg_range is not None:
+ jpeg_p = np.random.uniform(self.jpeg_range[0], self.jpeg_range[1])
+ encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), int(jpeg_p)]
+ _, encimg = cv2.imencode('.jpg', img_in * 255., encode_param)
+ img_in = np.float32(cv2.imdecode(encimg, 1)) / 255.
+
+ # resize to in_size
+ img_in = cv2.resize(img_in, (self.in_size, self.in_size), interpolation=cv2.INTER_LINEAR)
+
+
+ if self.gen_inpaint_mask:
+ img_in = (img_in*255).astype('uint8')
+ img_in = brush_stroke_mask(Image.fromarray(img_in))
+ img_in = np.array(img_in) / 255.
+
+ # random color jitter (only for lq)
+ if self.color_jitter_prob is not None and (np.random.uniform() < self.color_jitter_prob):
+ img_in = self.color_jitter(img_in, self.color_jitter_shift)
+ # random to gray (only for lq)
+ if self.gray_prob and np.random.uniform() < self.gray_prob:
+ img_in = cv2.cvtColor(img_in, cv2.COLOR_BGR2GRAY)
+ img_in = np.tile(img_in[:, :, None], [1, 1, 3])
+
+ # BGR to RGB, HWC to CHW, numpy to tensor
+ img_in, img_gt = img2tensor([img_in, img_gt], bgr2rgb=True, float32=True)
+
+ # random color jitter (pytorch version) (only for lq)
+ if self.color_jitter_pt_prob is not None and (np.random.uniform() < self.color_jitter_pt_prob):
+ brightness = self.opt.get('brightness', (0.5, 1.5))
+ contrast = self.opt.get('contrast', (0.5, 1.5))
+ saturation = self.opt.get('saturation', (0, 1.5))
+ hue = self.opt.get('hue', (-0.1, 0.1))
+ img_in = self.color_jitter_pt(img_in, brightness, contrast, saturation, hue)
+
+ # round and clip
+ img_in = np.clip((img_in * 255.0).round(), 0, 255) / 255.
+
+ # Set vgg range_norm=True if use the normalization here
+ # normalize
+ normalize(img_in, self.mean, self.std, inplace=True)
+ normalize(img_gt, self.mean, self.std, inplace=True)
+
+ img_in = img_in.unsqueeze(0)
+ img_gt = img_gt.unsqueeze(0)
+
+ in_list.append(img_in)
+ gt_list.append(img_gt)
+
+ in_video = torch.cat(in_list, dim=0)
+ gt_video = torch.cat(gt_list, dim=0)
+
+ return_dict = {'in': in_video, 'gt': gt_video}
return return_dict
def __len__(self):
- return len(self.paths)
\ No newline at end of file
+ return len(self.paths)
+
+
\ No newline at end of file
diff --git a/basicsr/losses/__init__.py b/basicsr/losses/__init__.py
old mode 100644
new mode 100755
diff --git a/basicsr/losses/loss_util.py b/basicsr/losses/loss_util.py
old mode 100644
new mode 100755
diff --git a/basicsr/losses/losses.py b/basicsr/losses/losses.py
old mode 100644
new mode 100755
diff --git a/basicsr/metrics/__init__.py b/basicsr/metrics/__init__.py
old mode 100644
new mode 100755
diff --git a/basicsr/metrics/metric_util.py b/basicsr/metrics/metric_util.py
old mode 100644
new mode 100755
diff --git a/basicsr/metrics/psnr_ssim.py b/basicsr/metrics/psnr_ssim.py
old mode 100644
new mode 100755
diff --git a/basicsr/models/__init__.py b/basicsr/models/__init__.py
old mode 100644
new mode 100755
diff --git a/basicsr/models/base_model.py b/basicsr/models/base_model.py
old mode 100644
new mode 100755
diff --git a/basicsr/models/codeformer_joint_model.py b/basicsr/models/codeformer_joint_model.py
deleted file mode 100644
index 9e4b9d12..00000000
--- a/basicsr/models/codeformer_joint_model.py
+++ /dev/null
@@ -1,350 +0,0 @@
-import torch
-from collections import OrderedDict
-from os import path as osp
-from tqdm import tqdm
-
-
-from basicsr.archs import build_network
-from basicsr.losses import build_loss
-from basicsr.metrics import calculate_metric
-from basicsr.utils import get_root_logger, imwrite, tensor2img
-from basicsr.utils.registry import MODEL_REGISTRY
-import torch.nn.functional as F
-from .sr_model import SRModel
-
-
-@MODEL_REGISTRY.register()
-class CodeFormerJointModel(SRModel):
- def feed_data(self, data):
- self.gt = data['gt'].to(self.device)
- self.input = data['in'].to(self.device)
- self.input_large_de = data['in_large_de'].to(self.device)
- self.b = self.gt.shape[0]
-
- if 'latent_gt' in data:
- self.idx_gt = data['latent_gt'].to(self.device)
- self.idx_gt = self.idx_gt.view(self.b, -1)
- else:
- self.idx_gt = None
-
- def init_training_settings(self):
- logger = get_root_logger()
- train_opt = self.opt['train']
-
- self.ema_decay = train_opt.get('ema_decay', 0)
- if self.ema_decay > 0:
- logger.info(f'Use Exponential Moving Average with decay: {self.ema_decay}')
- # define network net_g with Exponential Moving Average (EMA)
- # net_g_ema is used only for testing on one GPU and saving
- # There is no need to wrap with DistributedDataParallel
- self.net_g_ema = build_network(self.opt['network_g']).to(self.device)
- # load pretrained model
- load_path = self.opt['path'].get('pretrain_network_g', None)
- if load_path is not None:
- self.load_network(self.net_g_ema, load_path, self.opt['path'].get('strict_load_g', True), 'params_ema')
- else:
- self.model_ema(0) # copy net_g weight
- self.net_g_ema.eval()
-
- if self.opt['datasets']['train'].get('latent_gt_path', None) is not None:
- self.generate_idx_gt = False
- elif self.opt.get('network_vqgan', None) is not None:
- self.hq_vqgan_fix = build_network(self.opt['network_vqgan']).to(self.device)
- self.hq_vqgan_fix.eval()
- self.generate_idx_gt = True
- for param in self.hq_vqgan_fix.parameters():
- param.requires_grad = False
- else:
- raise NotImplementedError(f'Shoule have network_vqgan config or pre-calculated latent code.')
-
- logger.info(f'Need to generate latent GT code: {self.generate_idx_gt}')
-
- self.hq_feat_loss = train_opt.get('use_hq_feat_loss', True)
- self.feat_loss_weight = train_opt.get('feat_loss_weight', 1.0)
- self.cross_entropy_loss = train_opt.get('cross_entropy_loss', True)
- self.entropy_loss_weight = train_opt.get('entropy_loss_weight', 0.5)
- self.scale_adaptive_gan_weight = train_opt.get('scale_adaptive_gan_weight', 0.8)
-
- # define network net_d
- self.net_d = build_network(self.opt['network_d'])
- self.net_d = self.model_to_device(self.net_d)
- self.print_network(self.net_d)
-
- # load pretrained models
- load_path = self.opt['path'].get('pretrain_network_d', None)
- if load_path is not None:
- self.load_network(self.net_d, load_path, self.opt['path'].get('strict_load_d', True))
-
- self.net_g.train()
- self.net_d.train()
-
- # define losses
- if train_opt.get('pixel_opt'):
- self.cri_pix = build_loss(train_opt['pixel_opt']).to(self.device)
- else:
- self.cri_pix = None
-
- if train_opt.get('perceptual_opt'):
- self.cri_perceptual = build_loss(train_opt['perceptual_opt']).to(self.device)
- else:
- self.cri_perceptual = None
-
- if train_opt.get('gan_opt'):
- self.cri_gan = build_loss(train_opt['gan_opt']).to(self.device)
-
-
- self.fix_generator = train_opt.get('fix_generator', True)
- logger.info(f'fix_generator: {self.fix_generator}')
-
- self.net_g_start_iter = train_opt.get('net_g_start_iter', 0)
- self.net_d_iters = train_opt.get('net_d_iters', 1)
- self.net_d_start_iter = train_opt.get('net_d_start_iter', 0)
-
- # set up optimizers and schedulers
- self.setup_optimizers()
- self.setup_schedulers()
-
- def calculate_adaptive_weight(self, recon_loss, g_loss, last_layer, disc_weight_max):
- recon_grads = torch.autograd.grad(recon_loss, last_layer, retain_graph=True)[0]
- g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
-
- d_weight = torch.norm(recon_grads) / (torch.norm(g_grads) + 1e-4)
- d_weight = torch.clamp(d_weight, 0.0, disc_weight_max).detach()
- return d_weight
-
- def setup_optimizers(self):
- train_opt = self.opt['train']
- # optimizer g
- optim_params_g = []
- for k, v in self.net_g.named_parameters():
- if v.requires_grad:
- optim_params_g.append(v)
- else:
- logger = get_root_logger()
- logger.warning(f'Params {k} will not be optimized.')
- optim_type = train_opt['optim_g'].pop('type')
- self.optimizer_g = self.get_optimizer(optim_type, optim_params_g, **train_opt['optim_g'])
- self.optimizers.append(self.optimizer_g)
- # optimizer d
- optim_type = train_opt['optim_d'].pop('type')
- self.optimizer_d = self.get_optimizer(optim_type, self.net_d.parameters(), **train_opt['optim_d'])
- self.optimizers.append(self.optimizer_d)
-
- def gray_resize_for_identity(self, out, size=128):
- out_gray = (0.2989 * out[:, 0, :, :] + 0.5870 * out[:, 1, :, :] + 0.1140 * out[:, 2, :, :])
- out_gray = out_gray.unsqueeze(1)
- out_gray = F.interpolate(out_gray, (size, size), mode='bilinear', align_corners=False)
- return out_gray
-
- def optimize_parameters(self, current_iter):
- logger = get_root_logger()
- # optimize net_g
- for p in self.net_d.parameters():
- p.requires_grad = False
-
- self.optimizer_g.zero_grad()
-
- if self.generate_idx_gt:
- x = self.hq_vqgan_fix.encoder(self.gt)
- output, _, quant_stats = self.hq_vqgan_fix.quantize(x)
- min_encoding_indices = quant_stats['min_encoding_indices']
- self.idx_gt = min_encoding_indices.view(self.b, -1)
-
- if current_iter <= 40000: # small degradation
- small_per_n = 1
- w = 1
- elif current_iter <= 80000: # small degradation
- small_per_n = 1
- w = 1.3
- elif current_iter <= 120000: # large degradation
- small_per_n = 120000
- w = 0
- else: # mixed degradation
- small_per_n = 15
- w = 1.3
-
- if current_iter % small_per_n == 0:
- self.output, logits, lq_feat = self.net_g(self.input, w=w, detach_16=True)
- large_de = False
- else:
- logits, lq_feat = self.net_g(self.input_large_de, code_only=True)
- large_de = True
-
- if self.hq_feat_loss:
- # quant_feats
- quant_feat_gt = self.net_g.module.quantize.get_codebook_feat(self.idx_gt, shape=[self.b,16,16,256])
-
- l_g_total = 0
- loss_dict = OrderedDict()
- if current_iter % self.net_d_iters == 0 and current_iter > self.net_g_start_iter:
- # hq_feat_loss
- if not 'transformer' in self.opt['network_g']['fix_modules']:
- if self.hq_feat_loss: # codebook loss
- l_feat_encoder = torch.mean((quant_feat_gt.detach()-lq_feat)**2) * self.feat_loss_weight
- l_g_total += l_feat_encoder
- loss_dict['l_feat_encoder'] = l_feat_encoder
-
- # cross_entropy_loss
- if self.cross_entropy_loss:
- # b(hw)n -> bn(hw)
- cross_entropy_loss = F.cross_entropy(logits.permute(0, 2, 1), self.idx_gt) * self.entropy_loss_weight
- l_g_total += cross_entropy_loss
- loss_dict['cross_entropy_loss'] = cross_entropy_loss
-
- # pixel loss
- if not large_de: # when large degradation don't need image-level loss
- if self.cri_pix:
- l_g_pix = self.cri_pix(self.output, self.gt)
- l_g_total += l_g_pix
- loss_dict['l_g_pix'] = l_g_pix
-
- # perceptual loss
- if self.cri_perceptual:
- l_g_percep = self.cri_perceptual(self.output, self.gt)
- l_g_total += l_g_percep
- loss_dict['l_g_percep'] = l_g_percep
-
- # gan loss
- if current_iter > self.net_d_start_iter:
- fake_g_pred = self.net_d(self.output)
- l_g_gan = self.cri_gan(fake_g_pred, True, is_disc=False)
- recon_loss = l_g_pix + l_g_percep
- if not self.fix_generator:
- last_layer = self.net_g.module.generator.blocks[-1].weight
- d_weight = self.calculate_adaptive_weight(recon_loss, l_g_gan, last_layer, disc_weight_max=1.0)
- else:
- largest_fuse_size = self.opt['network_g']['connect_list'][-1]
- last_layer = self.net_g.module.fuse_convs_dict[largest_fuse_size].shift[-1].weight
- d_weight = self.calculate_adaptive_weight(recon_loss, l_g_gan, last_layer, disc_weight_max=1.0)
-
- d_weight *= self.scale_adaptive_gan_weight # 0.8
- loss_dict['d_weight'] = d_weight
- l_g_total += d_weight * l_g_gan
- loss_dict['l_g_gan'] = d_weight * l_g_gan
-
- l_g_total.backward()
- self.optimizer_g.step()
-
- if self.ema_decay > 0:
- self.model_ema(decay=self.ema_decay)
-
- # optimize net_d
- if not large_de:
- if current_iter > self.net_d_start_iter:
- for p in self.net_d.parameters():
- p.requires_grad = True
-
- self.optimizer_d.zero_grad()
- # real
- real_d_pred = self.net_d(self.gt)
- l_d_real = self.cri_gan(real_d_pred, True, is_disc=True)
- loss_dict['l_d_real'] = l_d_real
- loss_dict['out_d_real'] = torch.mean(real_d_pred.detach())
- l_d_real.backward()
- # fake
- fake_d_pred = self.net_d(self.output.detach())
- l_d_fake = self.cri_gan(fake_d_pred, False, is_disc=True)
- loss_dict['l_d_fake'] = l_d_fake
- loss_dict['out_d_fake'] = torch.mean(fake_d_pred.detach())
- l_d_fake.backward()
-
- self.optimizer_d.step()
-
- self.log_dict = self.reduce_loss_dict(loss_dict)
-
-
- def test(self):
- with torch.no_grad():
- if hasattr(self, 'net_g_ema'):
- self.net_g_ema.eval()
- self.output, _, _ = self.net_g_ema(self.input, w=1)
- else:
- logger = get_root_logger()
- logger.warning('Do not have self.net_g_ema, use self.net_g.')
- self.net_g.eval()
- self.output, _, _ = self.net_g(self.input, w=1)
- self.net_g.train()
-
-
- def dist_validation(self, dataloader, current_iter, tb_logger, save_img):
- if self.opt['rank'] == 0:
- self.nondist_validation(dataloader, current_iter, tb_logger, save_img)
-
-
- def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
- dataset_name = dataloader.dataset.opt['name']
- with_metrics = self.opt['val'].get('metrics') is not None
- if with_metrics:
- self.metric_results = {metric: 0 for metric in self.opt['val']['metrics'].keys()}
- pbar = tqdm(total=len(dataloader), unit='image')
-
- for idx, val_data in enumerate(dataloader):
- img_name = osp.splitext(osp.basename(val_data['lq_path'][0]))[0]
- self.feed_data(val_data)
- self.test()
-
- visuals = self.get_current_visuals()
- sr_img = tensor2img([visuals['result']])
- if 'gt' in visuals:
- gt_img = tensor2img([visuals['gt']])
- del self.gt
-
- # tentative for out of GPU memory
- del self.lq
- del self.output
- torch.cuda.empty_cache()
-
- if save_img:
- if self.opt['is_train']:
- save_img_path = osp.join(self.opt['path']['visualization'], img_name,
- f'{img_name}_{current_iter}.png')
- else:
- if self.opt['val']['suffix']:
- save_img_path = osp.join(self.opt['path']['visualization'], dataset_name,
- f'{img_name}_{self.opt["val"]["suffix"]}.png')
- else:
- save_img_path = osp.join(self.opt['path']['visualization'], dataset_name,
- f'{img_name}_{self.opt["name"]}.png')
- imwrite(sr_img, save_img_path)
-
- if with_metrics:
- # calculate metrics
- for name, opt_ in self.opt['val']['metrics'].items():
- metric_data = dict(img1=sr_img, img2=gt_img)
- self.metric_results[name] += calculate_metric(metric_data, opt_)
- pbar.update(1)
- pbar.set_description(f'Test {img_name}')
- pbar.close()
-
- if with_metrics:
- for metric in self.metric_results.keys():
- self.metric_results[metric] /= (idx + 1)
-
- self._log_validation_metric_values(current_iter, dataset_name, tb_logger)
-
-
- def _log_validation_metric_values(self, current_iter, dataset_name, tb_logger):
- log_str = f'Validation {dataset_name}\n'
- for metric, value in self.metric_results.items():
- log_str += f'\t # {metric}: {value:.4f}\n'
- logger = get_root_logger()
- logger.info(log_str)
- if tb_logger:
- for metric, value in self.metric_results.items():
- tb_logger.add_scalar(f'metrics/{metric}', value, current_iter)
-
-
- def get_current_visuals(self):
- out_dict = OrderedDict()
- out_dict['gt'] = self.gt.detach().cpu()
- out_dict['result'] = self.output.detach().cpu()
- return out_dict
-
-
- def save(self, epoch, current_iter):
- if self.ema_decay > 0:
- self.save_network([self.net_g, self.net_g_ema], 'net_g', current_iter, param_key=['params', 'params_ema'])
- else:
- self.save_network(self.net_g, 'net_g', current_iter)
- self.save_network(self.net_d, 'net_d', current_iter)
- self.save_training_state(epoch, current_iter)
diff --git a/basicsr/models/codeformer_model.py b/basicsr/models/codeformer_model.py
deleted file mode 100644
index 61829f12..00000000
--- a/basicsr/models/codeformer_model.py
+++ /dev/null
@@ -1,332 +0,0 @@
-import torch
-from collections import OrderedDict
-from os import path as osp
-from tqdm import tqdm
-
-from basicsr.archs import build_network
-from basicsr.losses import build_loss
-from basicsr.metrics import calculate_metric
-from basicsr.utils import get_root_logger, imwrite, tensor2img
-from basicsr.utils.registry import MODEL_REGISTRY
-import torch.nn.functional as F
-from .sr_model import SRModel
-
-
-@MODEL_REGISTRY.register()
-class CodeFormerModel(SRModel):
- def feed_data(self, data):
- self.gt = data['gt'].to(self.device)
- self.input = data['in'].to(self.device)
- self.b = self.gt.shape[0]
-
- if 'latent_gt' in data:
- self.idx_gt = data['latent_gt'].to(self.device)
- self.idx_gt = self.idx_gt.view(self.b, -1)
- else:
- self.idx_gt = None
-
- def init_training_settings(self):
- logger = get_root_logger()
- train_opt = self.opt['train']
-
- self.ema_decay = train_opt.get('ema_decay', 0)
- if self.ema_decay > 0:
- logger.info(f'Use Exponential Moving Average with decay: {self.ema_decay}')
- # define network net_g with Exponential Moving Average (EMA)
- # net_g_ema is used only for testing on one GPU and saving
- # There is no need to wrap with DistributedDataParallel
- self.net_g_ema = build_network(self.opt['network_g']).to(self.device)
- # load pretrained model
- load_path = self.opt['path'].get('pretrain_network_g', None)
- if load_path is not None:
- self.load_network(self.net_g_ema, load_path, self.opt['path'].get('strict_load_g', True), 'params_ema')
- else:
- self.model_ema(0) # copy net_g weight
- self.net_g_ema.eval()
-
- if self.opt.get('network_vqgan', None) is not None and self.opt['datasets'].get('latent_gt_path') is None:
- self.hq_vqgan_fix = build_network(self.opt['network_vqgan']).to(self.device)
- self.hq_vqgan_fix.eval()
- self.generate_idx_gt = True
- for param in self.hq_vqgan_fix.parameters():
- param.requires_grad = False
- else:
- self.generate_idx_gt = False
-
- self.hq_feat_loss = train_opt.get('use_hq_feat_loss', True)
- self.feat_loss_weight = train_opt.get('feat_loss_weight', 1.0)
- self.cross_entropy_loss = train_opt.get('cross_entropy_loss', True)
- self.entropy_loss_weight = train_opt.get('entropy_loss_weight', 0.5)
- self.fidelity_weight = train_opt.get('fidelity_weight', 1.0)
- self.scale_adaptive_gan_weight = train_opt.get('scale_adaptive_gan_weight', 0.8)
-
-
- self.net_g.train()
- # define network net_d
- if self.fidelity_weight > 0:
- self.net_d = build_network(self.opt['network_d'])
- self.net_d = self.model_to_device(self.net_d)
- self.print_network(self.net_d)
-
- # load pretrained models
- load_path = self.opt['path'].get('pretrain_network_d', None)
- if load_path is not None:
- self.load_network(self.net_d, load_path, self.opt['path'].get('strict_load_d', True))
-
- self.net_d.train()
-
- # define losses
- if train_opt.get('pixel_opt'):
- self.cri_pix = build_loss(train_opt['pixel_opt']).to(self.device)
- else:
- self.cri_pix = None
-
- if train_opt.get('perceptual_opt'):
- self.cri_perceptual = build_loss(train_opt['perceptual_opt']).to(self.device)
- else:
- self.cri_perceptual = None
-
- if train_opt.get('gan_opt'):
- self.cri_gan = build_loss(train_opt['gan_opt']).to(self.device)
-
-
- self.fix_generator = train_opt.get('fix_generator', True)
- logger.info(f'fix_generator: {self.fix_generator}')
-
- self.net_g_start_iter = train_opt.get('net_g_start_iter', 0)
- self.net_d_iters = train_opt.get('net_d_iters', 1)
- self.net_d_start_iter = train_opt.get('net_d_start_iter', 0)
-
- # set up optimizers and schedulers
- self.setup_optimizers()
- self.setup_schedulers()
-
- def calculate_adaptive_weight(self, recon_loss, g_loss, last_layer, disc_weight_max):
- recon_grads = torch.autograd.grad(recon_loss, last_layer, retain_graph=True)[0]
- g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
-
- d_weight = torch.norm(recon_grads) / (torch.norm(g_grads) + 1e-4)
- d_weight = torch.clamp(d_weight, 0.0, disc_weight_max).detach()
- return d_weight
-
- def setup_optimizers(self):
- train_opt = self.opt['train']
- # optimizer g
- optim_params_g = []
- for k, v in self.net_g.named_parameters():
- if v.requires_grad:
- optim_params_g.append(v)
- else:
- logger = get_root_logger()
- logger.warning(f'Params {k} will not be optimized.')
- optim_type = train_opt['optim_g'].pop('type')
- self.optimizer_g = self.get_optimizer(optim_type, optim_params_g, **train_opt['optim_g'])
- self.optimizers.append(self.optimizer_g)
- # optimizer d
- if self.fidelity_weight > 0:
- optim_type = train_opt['optim_d'].pop('type')
- self.optimizer_d = self.get_optimizer(optim_type, self.net_d.parameters(), **train_opt['optim_d'])
- self.optimizers.append(self.optimizer_d)
-
- def gray_resize_for_identity(self, out, size=128):
- out_gray = (0.2989 * out[:, 0, :, :] + 0.5870 * out[:, 1, :, :] + 0.1140 * out[:, 2, :, :])
- out_gray = out_gray.unsqueeze(1)
- out_gray = F.interpolate(out_gray, (size, size), mode='bilinear', align_corners=False)
- return out_gray
-
- def optimize_parameters(self, current_iter):
- logger = get_root_logger()
- # optimize net_g
- for p in self.net_d.parameters():
- p.requires_grad = False
-
- self.optimizer_g.zero_grad()
-
- if self.generate_idx_gt:
- x = self.hq_vqgan_fix.encoder(self.gt)
- output, _, quant_stats = self.hq_vqgan_fix.quantize(x)
- min_encoding_indices = quant_stats['min_encoding_indices']
- self.idx_gt = min_encoding_indices.view(self.b, -1)
-
- if self.fidelity_weight > 0:
- self.output, logits, lq_feat = self.net_g(self.input, w=self.fidelity_weight, detach_16=True)
- else:
- logits, lq_feat = self.net_g(self.input, w=0, code_only=True)
-
- if self.hq_feat_loss:
- # quant_feats
- quant_feat_gt = self.net_g.module.quantize.get_codebook_feat(self.idx_gt, shape=[self.b,16,16,256])
-
- l_g_total = 0
- loss_dict = OrderedDict()
- if current_iter % self.net_d_iters == 0 and current_iter > self.net_g_start_iter:
- # hq_feat_loss
- if self.hq_feat_loss: # codebook loss
- l_feat_encoder = torch.mean((quant_feat_gt.detach()-lq_feat)**2) * self.feat_loss_weight
- l_g_total += l_feat_encoder
- loss_dict['l_feat_encoder'] = l_feat_encoder
-
- # cross_entropy_loss
- if self.cross_entropy_loss:
- # b(hw)n -> bn(hw)
- cross_entropy_loss = F.cross_entropy(logits.permute(0, 2, 1), self.idx_gt) * self.entropy_loss_weight
- l_g_total += cross_entropy_loss
- loss_dict['cross_entropy_loss'] = cross_entropy_loss
-
- if self.fidelity_weight > 0: # when fidelity_weight == 0 don't need image-level loss
- # pixel loss
- if self.cri_pix:
- l_g_pix = self.cri_pix(self.output, self.gt)
- l_g_total += l_g_pix
- loss_dict['l_g_pix'] = l_g_pix
-
- # perceptual loss
- if self.cri_perceptual:
- l_g_percep = self.cri_perceptual(self.output, self.gt)
- l_g_total += l_g_percep
- loss_dict['l_g_percep'] = l_g_percep
-
- # gan loss
- if current_iter > self.net_d_start_iter:
- fake_g_pred = self.net_d(self.output)
- l_g_gan = self.cri_gan(fake_g_pred, True, is_disc=False)
- recon_loss = l_g_pix + l_g_percep
- if not self.fix_generator:
- last_layer = self.net_g.module.generator.blocks[-1].weight
- d_weight = self.calculate_adaptive_weight(recon_loss, l_g_gan, last_layer, disc_weight_max=1.0)
- else:
- largest_fuse_size = self.opt['network_g']['connect_list'][-1]
- last_layer = self.net_g.module.fuse_convs_dict[largest_fuse_size].shift[-1].weight
- d_weight = self.calculate_adaptive_weight(recon_loss, l_g_gan, last_layer, disc_weight_max=1.0)
-
- d_weight *= self.scale_adaptive_gan_weight # 0.8
- loss_dict['d_weight'] = d_weight
- l_g_total += d_weight * l_g_gan
- loss_dict['l_g_gan'] = d_weight * l_g_gan
-
- l_g_total.backward()
- self.optimizer_g.step()
-
- if self.ema_decay > 0:
- self.model_ema(decay=self.ema_decay)
-
- # optimize net_d
- if current_iter > self.net_d_start_iter and self.fidelity_weight > 0:
- for p in self.net_d.parameters():
- p.requires_grad = True
-
- self.optimizer_d.zero_grad()
- # real
- real_d_pred = self.net_d(self.gt)
- l_d_real = self.cri_gan(real_d_pred, True, is_disc=True)
- loss_dict['l_d_real'] = l_d_real
- loss_dict['out_d_real'] = torch.mean(real_d_pred.detach())
- l_d_real.backward()
- # fake
- fake_d_pred = self.net_d(self.output.detach())
- l_d_fake = self.cri_gan(fake_d_pred, False, is_disc=True)
- loss_dict['l_d_fake'] = l_d_fake
- loss_dict['out_d_fake'] = torch.mean(fake_d_pred.detach())
- l_d_fake.backward()
-
- self.optimizer_d.step()
-
- self.log_dict = self.reduce_loss_dict(loss_dict)
-
-
- def test(self):
- with torch.no_grad():
- if hasattr(self, 'net_g_ema'):
- self.net_g_ema.eval()
- self.output, _, _ = self.net_g_ema(self.input, w=self.fidelity_weight)
- else:
- logger = get_root_logger()
- logger.warning('Do not have self.net_g_ema, use self.net_g.')
- self.net_g.eval()
- self.output, _, _ = self.net_g(self.input, w=self.fidelity_weight)
- self.net_g.train()
-
-
- def dist_validation(self, dataloader, current_iter, tb_logger, save_img):
- if self.opt['rank'] == 0:
- self.nondist_validation(dataloader, current_iter, tb_logger, save_img)
-
-
- def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
- dataset_name = dataloader.dataset.opt['name']
- with_metrics = self.opt['val'].get('metrics') is not None
- if with_metrics:
- self.metric_results = {metric: 0 for metric in self.opt['val']['metrics'].keys()}
- pbar = tqdm(total=len(dataloader), unit='image')
-
- for idx, val_data in enumerate(dataloader):
- img_name = osp.splitext(osp.basename(val_data['lq_path'][0]))[0]
- self.feed_data(val_data)
- self.test()
-
- visuals = self.get_current_visuals()
- sr_img = tensor2img([visuals['result']])
- if 'gt' in visuals:
- gt_img = tensor2img([visuals['gt']])
- del self.gt
-
- # tentative for out of GPU memory
- del self.lq
- del self.output
- torch.cuda.empty_cache()
-
- if save_img:
- if self.opt['is_train']:
- save_img_path = osp.join(self.opt['path']['visualization'], img_name,
- f'{img_name}_{current_iter}.png')
- else:
- if self.opt['val']['suffix']:
- save_img_path = osp.join(self.opt['path']['visualization'], dataset_name,
- f'{img_name}_{self.opt["val"]["suffix"]}.png')
- else:
- save_img_path = osp.join(self.opt['path']['visualization'], dataset_name,
- f'{img_name}_{self.opt["name"]}.png')
- imwrite(sr_img, save_img_path)
-
- if with_metrics:
- # calculate metrics
- for name, opt_ in self.opt['val']['metrics'].items():
- metric_data = dict(img1=sr_img, img2=gt_img)
- self.metric_results[name] += calculate_metric(metric_data, opt_)
- pbar.update(1)
- pbar.set_description(f'Test {img_name}')
- pbar.close()
-
- if with_metrics:
- for metric in self.metric_results.keys():
- self.metric_results[metric] /= (idx + 1)
-
- self._log_validation_metric_values(current_iter, dataset_name, tb_logger)
-
-
- def _log_validation_metric_values(self, current_iter, dataset_name, tb_logger):
- log_str = f'Validation {dataset_name}\n'
- for metric, value in self.metric_results.items():
- log_str += f'\t # {metric}: {value:.4f}\n'
- logger = get_root_logger()
- logger.info(log_str)
- if tb_logger:
- for metric, value in self.metric_results.items():
- tb_logger.add_scalar(f'metrics/{metric}', value, current_iter)
-
-
- def get_current_visuals(self):
- out_dict = OrderedDict()
- out_dict['gt'] = self.gt.detach().cpu()
- out_dict['result'] = self.output.detach().cpu()
- return out_dict
-
-
- def save(self, epoch, current_iter):
- if self.ema_decay > 0:
- self.save_network([self.net_g, self.net_g_ema], 'net_g', current_iter, param_key=['params', 'params_ema'])
- else:
- self.save_network(self.net_g, 'net_g', current_iter)
- if self.fidelity_weight > 0:
- self.save_network(self.net_d, 'net_d', current_iter)
- self.save_training_state(epoch, current_iter)
diff --git a/basicsr/models/codeformer_idx_model.py b/basicsr/models/codeformer_temporal_model.py
old mode 100644
new mode 100755
similarity index 90%
rename from basicsr/models/codeformer_idx_model.py
rename to basicsr/models/codeformer_temporal_model.py
index 005957dd..44ff2d96
--- a/basicsr/models/codeformer_idx_model.py
+++ b/basicsr/models/codeformer_temporal_model.py
@@ -2,6 +2,7 @@
from collections import OrderedDict
from os import path as osp
from tqdm import tqdm
+from einops import rearrange
from basicsr.archs import build_network
from basicsr.metrics import calculate_metric
@@ -10,13 +11,16 @@
import torch.nn.functional as F
from .sr_model import SRModel
+# from icecream import ic
@MODEL_REGISTRY.register()
-class CodeFormerIdxModel(SRModel):
+class CodeFormerTempModel(SRModel):
def feed_data(self, data):
self.gt = data['gt'].to(self.device)
self.input = data['in'].to(self.device)
self.b = self.gt.shape[0]
+ self.f = self.gt.shape[1]
+ self.bf = self.b * self.f
if 'latent_gt' in data:
self.idx_gt = data['latent_gt'].to(self.device)
@@ -62,6 +66,13 @@ def init_training_settings(self):
self.entropy_loss_weight = train_opt.get('entropy_loss_weight', 0.5)
self.net_g.train()
+ self.net_g.requires_grad_(False)
+
+ trainable_module = train_opt['trainable_para']
+ for name, module in self.net_g.named_modules():
+ if trainable_module in name :
+ for params in module.parameters():
+ params.requires_grad_(True)
# set up optimizers and schedulers
self.setup_optimizers()
@@ -72,16 +83,18 @@ def setup_optimizers(self):
train_opt = self.opt['train']
# optimizer g
optim_params_g = []
+ optim_name = []
for k, v in self.net_g.named_parameters():
if v.requires_grad:
optim_params_g.append(v)
- else:
- logger = get_root_logger()
- logger.warning(f'Params {k} will not be optimized.')
+ optim_name.append(k)
+
optim_type = train_opt['optim_g'].pop('type')
self.optimizer_g = self.get_optimizer(optim_type, optim_params_g, **train_opt['optim_g'])
self.optimizers.append(self.optimizer_g)
+ # ic(optim_name)
+
def optimize_parameters(self, current_iter):
logger = get_root_logger()
@@ -89,16 +102,21 @@ def optimize_parameters(self, current_iter):
self.optimizer_g.zero_grad()
if self.generate_idx_gt:
- x = self.hq_vqgan_fix.encoder(self.gt)
+ x = rearrange(self.gt, "b f c h w -> (b f) c h w")
+ x = self.hq_vqgan_fix.encoder(x)
_, _, quant_stats = self.hq_vqgan_fix.quantize(x)
min_encoding_indices = quant_stats['min_encoding_indices']
- self.idx_gt = min_encoding_indices.view(self.b, -1)
+ # ic(min_encoding_indices.shape)
+ self.idx_gt = min_encoding_indices.view(self.bf, -1)
+ # ic(self.idx_gt.shape)
if self.hq_feat_loss:
# quant_feats
- quant_feat_gt = self.net_g.module.quantize.get_codebook_feat(self.idx_gt, shape=[self.b,16,16,256])
+ quant_feat_gt = self.net_g.module.quantize.get_codebook_feat(self.idx_gt, shape=[self.bf,16,16,256])
logits, lq_feat = self.net_g(self.input, w=0, code_only=True)
+ # ic(logits.shape)
+ # ic(lq_feat.shape)
l_g_total = 0
loss_dict = OrderedDict()
diff --git a/basicsr/models/lr_scheduler.py b/basicsr/models/lr_scheduler.py
old mode 100644
new mode 100755
diff --git a/basicsr/models/sr_model.py b/basicsr/models/sr_model.py
old mode 100644
new mode 100755
diff --git a/basicsr/models/vqgan_model.py b/basicsr/models/vqgan_model.py
old mode 100644
new mode 100755
diff --git a/basicsr/ops/__init__.py b/basicsr/ops/__init__.py
old mode 100644
new mode 100755
diff --git a/basicsr/ops/dcn/__init__.py b/basicsr/ops/dcn/__init__.py
old mode 100644
new mode 100755
diff --git a/basicsr/ops/dcn/deform_conv.py b/basicsr/ops/dcn/deform_conv.py
old mode 100644
new mode 100755
diff --git a/basicsr/ops/dcn/src/deform_conv_cuda.cpp b/basicsr/ops/dcn/src/deform_conv_cuda.cpp
old mode 100644
new mode 100755
diff --git a/basicsr/ops/dcn/src/deform_conv_cuda_kernel.cu b/basicsr/ops/dcn/src/deform_conv_cuda_kernel.cu
old mode 100644
new mode 100755
diff --git a/basicsr/ops/dcn/src/deform_conv_ext.cpp b/basicsr/ops/dcn/src/deform_conv_ext.cpp
old mode 100644
new mode 100755
diff --git a/basicsr/ops/fused_act/__init__.py b/basicsr/ops/fused_act/__init__.py
old mode 100644
new mode 100755
diff --git a/basicsr/ops/fused_act/fused_act.py b/basicsr/ops/fused_act/fused_act.py
old mode 100644
new mode 100755
diff --git a/basicsr/ops/fused_act/src/fused_bias_act.cpp b/basicsr/ops/fused_act/src/fused_bias_act.cpp
old mode 100644
new mode 100755
diff --git a/basicsr/ops/fused_act/src/fused_bias_act_kernel.cu b/basicsr/ops/fused_act/src/fused_bias_act_kernel.cu
old mode 100644
new mode 100755
diff --git a/basicsr/ops/upfirdn2d/__init__.py b/basicsr/ops/upfirdn2d/__init__.py
old mode 100644
new mode 100755
diff --git a/basicsr/ops/upfirdn2d/src/upfirdn2d.cpp b/basicsr/ops/upfirdn2d/src/upfirdn2d.cpp
old mode 100644
new mode 100755
diff --git a/basicsr/ops/upfirdn2d/src/upfirdn2d_kernel.cu b/basicsr/ops/upfirdn2d/src/upfirdn2d_kernel.cu
old mode 100644
new mode 100755
diff --git a/basicsr/ops/upfirdn2d/upfirdn2d.py b/basicsr/ops/upfirdn2d/upfirdn2d.py
old mode 100644
new mode 100755
diff --git a/basicsr/setup.py b/basicsr/setup.py
old mode 100644
new mode 100755
diff --git a/basicsr/train.py b/basicsr/train.py
old mode 100644
new mode 100755
index a01c0dfc..84fe0a28
--- a/basicsr/train.py
+++ b/basicsr/train.py
@@ -25,7 +25,7 @@ def parse_options(root_path, is_train=True):
parser = argparse.ArgumentParser()
parser.add_argument('-opt', type=str, required=True, help='Path to option YAML file.')
parser.add_argument('--launcher', choices=['none', 'pytorch', 'slurm'], default='none', help='job launcher')
- parser.add_argument('--local_rank', type=int, default=0)
+ parser.add_argument('--local-rank','--local_rank', type=int, default=0)
args = parser.parse_args()
opt = parse(args.opt, root_path, is_train=is_train)
diff --git a/basicsr/utils/__init__.py b/basicsr/utils/__init__.py
old mode 100644
new mode 100755
diff --git a/basicsr/utils/dist_util.py b/basicsr/utils/dist_util.py
old mode 100644
new mode 100755
diff --git a/basicsr/utils/download_util.py b/basicsr/utils/download_util.py
old mode 100644
new mode 100755
diff --git a/basicsr/utils/file_client.py b/basicsr/utils/file_client.py
old mode 100644
new mode 100755
diff --git a/basicsr/utils/img_util.py b/basicsr/utils/img_util.py
old mode 100644
new mode 100755
diff --git a/basicsr/utils/lmdb_util.py b/basicsr/utils/lmdb_util.py
old mode 100644
new mode 100755
diff --git a/basicsr/utils/logger.py b/basicsr/utils/logger.py
old mode 100644
new mode 100755
diff --git a/basicsr/utils/matlab_functions.py b/basicsr/utils/matlab_functions.py
old mode 100644
new mode 100755
diff --git a/basicsr/utils/misc.py b/basicsr/utils/misc.py
old mode 100644
new mode 100755
diff --git a/basicsr/utils/options.py b/basicsr/utils/options.py
old mode 100644
new mode 100755
diff --git a/basicsr/utils/realesrgan_utils.py b/basicsr/utils/realesrgan_utils.py
old mode 100644
new mode 100755
diff --git a/basicsr/utils/registry.py b/basicsr/utils/registry.py
old mode 100644
new mode 100755
diff --git a/basicsr/utils/video_util.py b/basicsr/utils/video_util.py
old mode 100644
new mode 100755
diff --git a/options/CodeFormer_stage2.yml b/configs/train/video_sr.yaml
old mode 100755
new mode 100644
similarity index 83%
rename from options/CodeFormer_stage2.yml
rename to configs/train/video_sr.yaml
index 4dfe9c9d..28ba65db
--- a/options/CodeFormer_stage2.yml
+++ b/configs/train/video_sr.yaml
@@ -1,15 +1,15 @@
# general settings
-name: CodeFormer_stage2
-model_type: CodeFormerIdxModel
+name: CodeFormer_temp
+model_type: CodeFormerTempModel
num_gpu: 8
manual_seed: 0
# dataset and data loader settings
datasets:
train:
- name: FFHQ
- type: FFHQBlindDataset
- dataroot_gt: datasets/ffhq/ffhq_512
+ name: VFHQ
+ type: VFHQBlindDataset
+ dataroot_gt: ./VFHQ/image
filename_tmpl: '{}'
io_backend:
type: disk
@@ -20,6 +20,7 @@ datasets:
std: [0.5, 0.5, 0.5]
use_hflip: true
use_corrupt: true
+ video_length: 16
# large degradation in stageII
blur_kernel_size: 41
@@ -33,12 +34,11 @@ datasets:
jpeg_range: [30, 80]
latent_gt_path: ~ # without pre-calculated latent code
- # latent_gt_path: './experiments/pretrained_models/VQGAN/latent_gt_code1024.pth'
# data loader
- num_worker_per_gpu: 2
+ num_worker_per_gpu: 8
batch_size_per_gpu: 4
- dataset_enlarge_ratio: 100
+ dataset_enlarge_ratio: 1
prefetch_mode: ~
# val:
@@ -61,7 +61,7 @@ network_g:
codebook_size: 1024
connect_list: ['32', '64', '128', '256']
fix_modules: ['quantize','generator']
- vqgan_path: './experiments/pretrained_models/vqgan/vqgan_code1024.pth' # pretrained VQGAN
+ vqgan_path: './pretrained_models/CodeFormer/vqgan_code1024.pth' # pretrained VQGAN
network_vqgan: # this config is needed if no pre-calculated latent
type: VQAutoEncoder
@@ -70,10 +70,11 @@ network_vqgan: # this config is needed if no pre-calculated latent
ch_mult: [1, 2, 2, 4, 4, 8]
quantizer: 'nearest'
codebook_size: 1024
+ model_path: './pretrained_models/CodeFormer/vqgan_code1024.pth'
# path
path:
- pretrain_network_g: ~
+ pretrain_network_g: './pretrained_models/CodeFormer/codeformer.pth'
param_key_g: params_ema
strict_load_g: false
pretrain_network_d: ~
@@ -88,6 +89,8 @@ train:
entropy_loss_weight: 0.5
fidelity_weight: 0
+ trainable_para: temp
+
optim_g:
type: Adam
lr: !!float 1e-4
@@ -119,7 +122,7 @@ train:
# validation settings
val:
- val_freq: !!float 5e10 # no validation
+ val_freq: 1000
save_img: true
metrics:
@@ -130,8 +133,8 @@ val:
# logging settings
logger:
- print_freq: 100
- save_checkpoint_freq: !!float 1e4
+ print_freq: 1
+ save_checkpoint_freq: 1000
use_tb_logger: true
wandb:
project: ~
diff --git a/docs/history_changelog.md b/docs/history_changelog.md
deleted file mode 100644
index 6c35e34e..00000000
--- a/docs/history_changelog.md
+++ /dev/null
@@ -1,15 +0,0 @@
-# History of Changelog
-
-- **2023.04.19**: :whale: Training codes and config files are public available now.
-- **2023.04.09**: Add features of inpainting and colorization for cropped face images.
-- **2023.02.10**: Include `dlib` as a new face detector option, it produces more accurate face identity.
-- **2022.10.05**: Support video input `--input_path [YOUR_VIDEO.mp4]`. Try it to enhance your videos! :clapper:
-- **2022.09.14**: Integrated to :hugs: [Hugging Face](https://huggingface.co/spaces). Try out online demo! [](https://huggingface.co/spaces/sczhou/CodeFormer)
-- **2022.09.09**: Integrated to :rocket: [Replicate](https://replicate.com/explore). Try out online demo! [](https://replicate.com/sczhou/codeformer)
-- **2022.09.04**: Add face upsampling `--face_upsample` for high-resolution AI-created face enhancement.
-- **2022.08.23**: Some modifications on face detection and fusion for better AI-created face enhancement.
-- **2022.08.07**: Integrate [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN) to support background image enhancement.
-- **2022.07.29**: Integrate new face detectors of `['RetinaFace'(default), 'YOLOv5']`.
-- **2022.07.17**: Add Colab demo of CodeFormer.
-- **2022.07.16**: Release inference code for face restoration. :blush:
-- **2022.06.21**: This repo is created.
\ No newline at end of file
diff --git a/docs/train.md b/docs/train.md
deleted file mode 100644
index 873ee7cc..00000000
--- a/docs/train.md
+++ /dev/null
@@ -1,37 +0,0 @@
-# :milky_way: Training Procedures
-[English](train.md) **|** [įŽäŊä¸æ](train_CN.md)
-## Preparing Dataset
-
-- Download training dataset: [FFHQ](https://github.com/NVlabs/ffhq-dataset)
-
----
-
-## Training
-```
-For PyTorch versions >= 1.10, please replace `python -m torch.distributed.launch` in the commands below with `torchrun`.
-```
-
-### đž Stage I - VQGAN
-- Training VQGAN:
- > python -m torch.distributed.launch --nproc_per_node=gpu_num --master_port=4321 basicsr/train.py -opt options/VQGAN_512_ds32_nearest_stage1.yml --launcher pytorch
-
-- After VQGAN training, you can pre-calculate code sequence for the training dataset to speed up the later training stages:
- > python scripts/generate_latent_gt.py
-
-- If you don't require training your own VQGAN, you can find pre-trained VQGAN (`vqgan_code1024.pth`) and the corresponding code sequence (`latent_gt_code1024.pth`) in the folder of Releases v0.1.0: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
-
-### đ Stage II - CodeFormer (w=0)
-- Training Code Sequence Prediction Module:
- > python -m torch.distributed.launch --nproc_per_node=gpu_num --master_port=4322 basicsr/train.py -opt options/CodeFormer_stage2.yml --launcher pytorch
-
-- Pre-trained CodeFormer of stage II (`codeformer_stage2.pth`) can be found in the folder of Releases v0.1.0: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
-
-### đ¸ Stage III - CodeFormer (w=1)
-- Training Controllable Module:
- > python -m torch.distributed.launch --nproc_per_node=gpu_num --master_port=4323 basicsr/train.py -opt options/CodeFormer_stage3.yml --launcher pytorch
-
-- Pre-trained CodeFormer (`codeformer.pth`) can be found in the folder of Releases v0.1.0: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
-
----
-
-:whale: The project was built using the framework [BasicSR](https://github.com/XPixelGroup/BasicSR). For detailed information on training, resuming, and other related topics, please refer to the documentation: https://github.com/XPixelGroup/BasicSR/blob/master/docs/TrainTest.md
diff --git a/docs/train_CN.md b/docs/train_CN.md
deleted file mode 100644
index c1ac2cc4..00000000
--- a/docs/train_CN.md
+++ /dev/null
@@ -1,37 +0,0 @@
-# :milky_way: čŽįģææĄŖ
-[English](train.md) **|** [įŽäŊä¸æ](train_CN.md)
-
-## åå¤æ°æŽé
-- ä¸čŊŊčŽįģæ°æŽé: [FFHQ](https://github.com/NVlabs/ffhq-dataset)
-
----
-
-## čŽįģ
-```
-寚äēPyTorchįæŦ >= 1.10, č¯ˇå°ä¸éĸåŊäģ¤ä¸į`python -m torch.distributed.launch`æŋæĸä¸ē`torchrun`.
-```
-
-### đž éļæŽĩ I - VQGAN
-- čŽįģVQGAN:
- > python -m torch.distributed.launch --nproc_per_node=gpu_num --master_port=4321 basicsr/train.py -opt options/VQGAN_512_ds32_nearest_stage1.yml --launcher pytorch
-
-- čŽįģåŽVQGANåīŧå¯äģĨéčŋä¸éĸäģŖį éĸå
čˇåžčŽįģæ°æŽéįå¯į æŦåēåīŧäģčå éåéĸéļæŽĩįčŽįģčŋį¨:
- > python scripts/generate_latent_gt.py
-
-- åĻæäŊ ä¸éčĻčŽįģčĒåˇąįVQGANīŧå¯äģĨå¨Release v0.1.0ææĄŖä¸æžå°éĸčŽįģįVQGAN (`vqgan_code1024.pth`)å寚åēįå¯į æŦåēå (`latent_gt_code1024.pth`): https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
-
-### đ éļæŽĩ II - CodeFormer (w=0)
-- čŽįģå¯į æŦčŽįģéĸæĩæ¨Ąå:
- > python -m torch.distributed.launch --nproc_per_node=gpu_num --master_port=4322 basicsr/train.py -opt options/CodeFormer_stage2.yml --launcher pytorch
-
-- éĸčŽįģCodeFormerįŦŦäēéļæŽĩæ¨Ąå (`codeformer_stage2.pth`)å¯äģĨå¨Releases v0.1.0ææĄŖéä¸čŊŊ: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
-
-### đ¸ éļæŽĩ III - CodeFormer (w=1)
-- čŽįģå¯č°æ¨Ąå:
- > python -m torch.distributed.launch --nproc_per_node=gpu_num --master_port=4323 basicsr/train.py -opt options/CodeFormer_stage3.yml --launcher pytorch
-
-- éĸčŽįģCodeFormeræ¨Ąå (`codeformer.pth`)å¯äģĨå¨Releases v0.1.0ææĄŖéä¸čŊŊ: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
-
----
-
-:whale: č¯Ĩ饚įŽæ¯åēäē[BasicSR](https://github.com/XPixelGroup/BasicSR)æĄæļæåģēīŧæå
ŗčŽįģãResumeįč¯Ļįģäģįģå¯äģĨæĨįææĄŖ: https://github.com/XPixelGroup/BasicSR/blob/master/docs/TrainTest_CN.md
\ No newline at end of file
diff --git a/inference_colorization.py b/inference_colorization.py
deleted file mode 100644
index 0f1b763c..00000000
--- a/inference_colorization.py
+++ /dev/null
@@ -1,86 +0,0 @@
-import os
-import cv2
-import argparse
-import glob
-import torch
-from torchvision.transforms.functional import normalize
-from basicsr.utils import imwrite, img2tensor, tensor2img
-from basicsr.utils.download_util import load_file_from_url
-from basicsr.utils.misc import get_device
-from basicsr.utils.registry import ARCH_REGISTRY
-
-pretrain_model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer_colorization.pth'
-
-if __name__ == '__main__':
- # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
- device = get_device()
- parser = argparse.ArgumentParser()
-
- parser.add_argument('-i', '--input_path', type=str, default='./inputs/gray_faces',
- help='Input image or folder. Default: inputs/gray_faces')
- parser.add_argument('-o', '--output_path', type=str, default=None,
- help='Output folder. Default: results/
- 
-