yuki-11 / cssr

Crack Segmentation for Low-Resolution Images using Joint Learning with Super-Resolution (CSSR) was accepted to international conference on MVA2021 (oral), and selected for the Best Practical Paper Award.

Hi 👋, I'm Yuki Kondo.

A Research sientist at Toyota Motor Corporation.

A Researcher at Toyota Technological Institute.

• 📫 How to reach me [email protected]

• 📄 Know about my experiences https://yuki-11.github.io/ (NOTE: The English version may be out of date, as my latest information has been updated with priority given to the Japanese version)

Connect with me:

Languages and Tools:

 

Recommended repositories

cssr's Issues

How can we get result for new images, not from test datasets?

About saving results

Hi.
I'm trying to retrain the SS_SR, but during training and after it I won't get the results saved nowhere.
I'm using the cfg parameter in the yaml file "OUTPUT_DIR", and the --output_dirname parameter.

There is a bug somewhere on this branch or I'm doing something wrong?

On Colab:
!python train.py --config_file "./config/exp2/config_CSSR_SR-SS.yaml" --num_workers 2 --num_gpus 1 --output_dirname "/content/CSSR/output/" --save_step 250

Test Runtime error: Missing key in state_dict

Description

When performing the test with standard configurazione I got an error of missing keys in the model configuration

Error

Configration file is loaded from output/CSSR_SS-SR/config.yaml
Running with config:
BASE_NET: weights/vgg16_reducedfc.pth
DATASET:
  TEST_IMAGE_DIR: datasets/crack_segmentation_dataset/test/images
  TEST_MASK_DIR: datasets/crack_segmentation_dataset/test/masks
  TRAIN_IMAGE_DIR: datasets/crack_segmentation_dataset/train/images
  TRAIN_MASK_DIR: datasets/crack_segmentation_dataset/train/masks
DEVICE: cuda
INPUT:
  IMAGE_SIZE: [448, 448]
  MEAN: [0.4741, 0.4937, 0.5048]
  STD: [0.1621, 0.1532, 0.1523]
MODEL:
  DETECTOR_DBPN_NUM_STAGES: 4
  DETECTOR_TYPE: u-net16
  JOINT_LEARNING: True
  NUM_CLASSES: 1
  NUM_STAGES: 6
  OPTIMIZER: Adam
  SCALE_FACTOR: 4
  SR: DBPN
  SR_SEG_INV: True
  UP_SAMPLE_METHOD: pixel_shuffle
OUTPUT_DIR: output/CSSR_SS-SR/eval_AIU/iter_100000
SEED: 123
SOLVER:
  ALPHA_MIN: 0.01
  BATCH_SIZE: 6
  BCELOSS_WEIGHT: [1, 1]
  BOUNDARY_DEC_RATIO: 1.0
  DECREASE_RATIO: 1.0
  GAMMA: 0.1
  LR: 1e-05
  LR_STEPS: []
  MAX_ITER: 100000
  NORM_SR_OUTPUT: instance
  SEG_LOSS_FUNC: BoundaryCombo
  SEG_PRETRAIN_ITER: 0
  SR_LOSS_FUNC: L1
  SR_PRETRAIN_ITER: 0
  SYNC_BATCHNORM: True
  TASK_LOSS_WEIGHT: 0.3
  TRAIN_DATASET_RATIO: 0.95
  WARMUP_FACTOR: 1.0
  WARMUP_ITERS: 5000
  WB_AND_D_WEIGHT: [1, 1]
Building model...
/usr/local/lib/python3.7/dist-packages/torchvision/models/_utils.py:209: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and will be removed in 0.15, please use 'weights' instead.
  f"The parameter '{pretrained_param}' is deprecated since 0.13 and will be removed in 0.15, "
/usr/local/lib/python3.7/dist-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and will be removed in 0.15. The current behavior is equivalent to passing `weights=VGG16_Weights.IMAGENET1K_V1`. You can also use `weights=VGG16_Weights.DEFAULT` to get the most up-to-date weights.
  warnings.warn(msg)
SR pretrained model was loaded from weights/sr_pretrain_x4_stage6.pth
------------Model Architecture-------------

<Network SS>
UNet16(
  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (encoder): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (relu): ReLU(inplace=True)
  (conv1): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
  )
  (conv2): Sequential(
    (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
  )
  (conv3): Sequential(
    (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (5): ReLU(inplace=True)
  )
  (conv4): Sequential(
    (0): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (5): ReLU(inplace=True)
  )
  (conv5): Sequential(
    (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (5): ReLU(inplace=True)
  )
  (center): DecoderBlockV2(
    (block): Sequential(
      (0): ConvRelu(
        (conv): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (activation): ReLU(inplace=True)
      )
      (1): PixelShuffle(upscale_factor=2)
      (2): ReLU(inplace=True)
    )
  )
  (dec4): DecoderBlockV2(
    (block): Sequential(
      (0): ConvRelu(
        (conv): Conv2d(768, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (activation): ReLU(inplace=True)
      )
      (1): PixelShuffle(upscale_factor=2)
      (2): ReLU(inplace=True)
    )
  )
  (dec3): DecoderBlockV2(
    (block): Sequential(
      (0): ConvRelu(
        (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (activation): ReLU(inplace=True)
      )
      (1): PixelShuffle(upscale_factor=2)
      (2): ReLU(inplace=True)
    )
  )
  (dec2): DecoderBlockV2(
    (block): Sequential(
      (0): ConvRelu(
        (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (activation): ReLU(inplace=True)
      )
      (1): PixelShuffle(upscale_factor=2)
      (2): ReLU(inplace=True)
    )
  )
  (dec1): ConvRelu(
    (conv): Conv2d(128, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (activation): ReLU(inplace=True)
  )
  (final): Conv2d(32, 1, kernel_size=(1, 1), stride=(1, 1))
  (sigmoid): Sigmoid()
)

<Network SR>
Net_6(
  (feat0): ConvBlock(
    (conv): Conv2d(1, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (act): PReLU(num_parameters=1)
  )
  (feat1): ConvBlock(
    (conv): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1))
    (act): PReLU(num_parameters=1)
  )
  (up1): UpBlock(
    (up_conv1): DeconvBlock(
      (deconv): ConvTranspose2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (up_conv2): ConvBlock(
      (conv): Conv2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (up_conv3): DeconvBlock(
      (deconv): ConvTranspose2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
  )
  (down1): DownBlock(
    (down_conv1): ConvBlock(
      (conv): Conv2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (down_conv2): DeconvBlock(
      (deconv): ConvTranspose2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (down_conv3): ConvBlock(
      (conv): Conv2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
  )
  (up2): UpBlock(
    (up_conv1): DeconvBlock(
      (deconv): ConvTranspose2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (up_conv2): ConvBlock(
      (conv): Conv2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (up_conv3): DeconvBlock(
      (deconv): ConvTranspose2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
  )
  (down2): D_DownBlock(
    (conv): ConvBlock(
      (conv): Conv2d(128, 64, kernel_size=(1, 1), stride=(1, 1))
      (act): PReLU(num_parameters=1)
    )
    (down_conv1): ConvBlock(
      (conv): Conv2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (down_conv2): DeconvBlock(
      (deconv): ConvTranspose2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (down_conv3): ConvBlock(
      (conv): Conv2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
  )
  (up3): D_UpBlock(
    (conv): ConvBlock(
      (conv): Conv2d(128, 64, kernel_size=(1, 1), stride=(1, 1))
      (act): PReLU(num_parameters=1)
    )
    (up_conv1): DeconvBlock(
      (deconv): ConvTranspose2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (up_conv2): ConvBlock(
      (conv): Conv2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (up_conv3): DeconvBlock(
      (deconv): ConvTranspose2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
  )
  (down3): D_DownBlock(
    (conv): ConvBlock(
      (conv): Conv2d(192, 64, kernel_size=(1, 1), stride=(1, 1))
      (act): PReLU(num_parameters=1)
    )
    (down_conv1): ConvBlock(
      (conv): Conv2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (down_conv2): DeconvBlock(
      (deconv): ConvTranspose2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (down_conv3): ConvBlock(
      (conv): Conv2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
  )
  (up4): D_UpBlock(
    (conv): ConvBlock(
      (conv): Conv2d(192, 64, kernel_size=(1, 1), stride=(1, 1))
      (act): PReLU(num_parameters=1)
    )
    (up_conv1): DeconvBlock(
      (deconv): ConvTranspose2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (up_conv2): ConvBlock(
      (conv): Conv2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (up_conv3): DeconvBlock(
      (deconv): ConvTranspose2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
  )
  (down4): D_DownBlock(
    (conv): ConvBlock(
      (conv): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1))
      (act): PReLU(num_parameters=1)
    )
    (down_conv1): ConvBlock(
      (conv): Conv2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (down_conv2): DeconvBlock(
      (deconv): ConvTranspose2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (down_conv3): ConvBlock(
      (conv): Conv2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
  )
  (up5): D_UpBlock(
    (conv): ConvBlock(
      (conv): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1))
      (act): PReLU(num_parameters=1)
    )
    (up_conv1): DeconvBlock(
      (deconv): ConvTranspose2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (up_conv2): ConvBlock(
      (conv): Conv2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (up_conv3): DeconvBlock(
      (deconv): ConvTranspose2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
  )
  (down5): D_DownBlock(
    (conv): ConvBlock(
      (conv): Conv2d(320, 64, kernel_size=(1, 1), stride=(1, 1))
      (act): PReLU(num_parameters=1)
    )
    (down_conv1): ConvBlock(
      (conv): Conv2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (down_conv2): DeconvBlock(
      (deconv): ConvTranspose2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (down_conv3): ConvBlock(
      (conv): Conv2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
  )
  (up6): D_UpBlock(
    (conv): ConvBlock(
      (conv): Conv2d(320, 64, kernel_size=(1, 1), stride=(1, 1))
      (act): PReLU(num_parameters=1)
    )
    (up_conv1): DeconvBlock(
      (deconv): ConvTranspose2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (up_conv2): ConvBlock(
      (conv): Conv2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
    (up_conv3): DeconvBlock(
      (deconv): ConvTranspose2d(64, 64, kernel_size=(8, 8), stride=(4, 4), padding=(2, 2))
      (act): PReLU(num_parameters=1)
    )
  )
  (output_conv): ConvBlock(
    (conv): Conv2d(384, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  )
)
Traceback (most recent call last):
  File "test.py", line 96, in <module>
    main()
  File "test.py", line 91, in main
    test(args, cfg)
  File "test.py", line 28, in test
    model.load_state_dict(fix_model_state_dict(torch.load(args.trained_model, map_location=lambda storage, loc:storage)))
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1605, in load_state_dict
    self.__class__.__name__, "\n\t".join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for InvModel:
	Missing key(s) in state_dict: "sr_model.feat0.conv.weight", "sr_model.feat0.conv.bias", "sr_model.feat0.act.weight", "sr_model.output_conv.conv.weight", "sr_model.output_conv.conv.bias". 
	Unexpected key(s) in state_dict: "sr_model.feat0_1ch.conv.weight", "sr_model.feat0_1ch.conv.bias", "sr_model.feat0_1ch.act.weight", "sr_model.output_conv_1ch.conv.weight", "sr_model.output_conv_1ch.conv.bias". 

To Reproduce

!python test.py "output/CSSR_SS-SR/" 100000

Enviromet

Google Colab GPU