ewrfcas / ZITS-PlusPlus

ZITS++: Image Inpainting by Improving the Incremental Transformer on Structural Priors (TPAMI2023)
Apache License 2.0
58 stars 5 forks source link

Missing keys in the provided checkpoint #6

Open LoukasOrtyx opened 8 months ago

LoukasOrtyx commented 8 months ago

First of all, I would like to say I loved your project.

I tried to replicate it on my machine using the weights you provided model_512 and LSM-HAWP:, but when I execute the test script I get the following output (shortened for better visualization):

root@457540d74c7a:/app# sh test.sh VANFFCResNetGenerator initialized use swish use pure FFC load HAWP Loading checkpoint: ./last.ckpt ... missing keys:['wf.tspan', 'wf.backbone.conv1.weight', 'wf.backbone.conv1.bias', 'wf.backbone.bn1.weight', 'wf.backbone.bn1.bias', 'wf.backbone.bn1.running_mean', 'wf.backbone.bn1.running_var', 'wf.backbone.layer1.0.bn1.weight', 'wf.backbone.layer1.0.bn1.bias', 'wf.backbone.layer1.0.bn1.running_mean', 'wf.backbone.layer1.0.bn1.running_var', 'wf.backbone.layer1.0.conv1.weight', 'wf.backbone.layer1.0.conv1.bias', 'wf.backbone.layer1.0.bn2.weight', 'wf.backbone.layer1.0.bn2.bias', 'wf.backbone.layer1.0.bn2.running_mean', 'wf.backbone.layer1.0.bn2.running_var', 'wf.backbone.layer1.0.conv2.weight', .......... 'loss_resnet_pl.impl.layer4.2.bn2.running_mean', 'loss_resnet_pl.impl.layer4.2.bn2.running_var', 'loss_resnet_pl.impl.layer4.2.bn2.num_batches_tracked', 'loss_resnet_pl.impl.layer4.2.conv3.weight', 'loss_resnet_pl.impl.layer4.2.bn3.weight', 'loss_resnet_pl.impl.layer4.2.bn3.bias', 'loss_resnet_pl.impl.layer4.2.bn3.running_mean', 'loss_resnet_pl.impl.layer4.2.bn3.running_var', 'loss_resnet_pl.impl.layer4.2.bn3.num_batches_tracked']

Despite this, the code outputs the inpainted results, although they seem to lack some quality. I believe the missing keys are the reason the result falls short.

Here is the command I use:

CUDA_VISIBLE_DEVICES=0 python3 test.py --config config.yml \ --exp_name teste \ --ckpt_resume ./last.ckpt \ --save_path ./outputs \ --img_dir ./images \ --mask_dir ./masks \ --use_ema \ --wf_ckpt ./best_lsm_hawp.pth \ --test_size 512

The config.yaml and the last.ckpt comes from the model_512 and the best_lsm_hawp.pth comes from LSM-HAWP:.

And here are the results: img5 img10

We can see in the third column that the network successfully infers the contours of the image, but surprisingly fails to fill the missing spots correctly although there is little to no background complexity around the mask, so I don't understand why the result is so poor.

Am I doing something wrong or the results are like that, is the provided weights missing some keys as stated by the code?

ewrfcas commented 8 months ago

Hi, wf is the line detection model. Since the line results are correct, wf should be already loaded successfully. Besides, loss_resnet_pl is only used for training. I guess that the failure might be caused by the E-NMS. Could you try the testing in 256*256 with model_256 again?

ewrfcas commented 8 months ago

Moreover, you could try to raise binary_threshold for get_np_nms to filter useless edges.