mmagic.models.editors.nafnet.nafbaseline_net

Module Contents

Classes

NAFBaseline	The original version of Baseline model in "Simple Baseline for Image
NAFBaselineLocal	The original version of Baseline model in "Simple Baseline for Image
BaselineBlock	Baseline's Block in paper.

class mmagic.models.editors.nafnet.nafbaseline_net.NAFBaseline(img_channel=3, mid_channels=16, middle_blk_num=1, enc_blk_nums=[1, 1, 1, 28], dec_blk_nums=[1, 1, 1, 1], dw_expand=1, ffn_expand=2)

Bases: mmengine.model.BaseModule

The original version of Baseline model in “Simple Baseline for Image Restoration”.

Parameters

• img_channels (int) – Channel number of inputs.

• mid_channels (int) – Channel number of intermediate features.

• middle_blk_num (int) – Number of middle blocks.

• enc_blk_nums (List of int) – Number of blocks for each encoder.

• dec_blk_nums (List of int) – Number of blocks for each decoder.

forward(inp)

Forward function.

Parameters

inp – input tensor image with (B, C, H, W) shape

check_image_size(x)

Check image size and pad images so that it has enough dimension do downsample.

Parameters

x – input tensor image with (B, C, H, W) shape.

class mmagic.models.editors.nafnet.nafbaseline_net.NAFBaselineLocal(*args, train_size=(1, 3, 256, 256), fast_imp=False, **kwargs)

Bases: mmagic.models.editors.nafnet.naf_avgpool2d.Local_Base, NAFBaseline

The original version of Baseline model in “Simple Baseline for Image Restoration”.

Parameters

• img_channels (int) – Channel number of inputs.

• mid_channels (int) – Channel number of intermediate features.

• middle_blk_num (int) – Number of middle blocks.

• enc_blk_nums (List of int) – Number of blocks for each encoder.

• dec_blk_nums (L`ist of int) – Number of blocks for each decoder.

class mmagic.models.editors.nafnet.nafbaseline_net.BaselineBlock(in_channels, DW_Expand=1, FFN_Expand=2, drop_out_rate=0.0)

Bases: mmengine.model.BaseModule

Baseline’s Block in paper.

Parameters

• in_channels (int) – number of channels

• DW_Expand (int) – channel expansion factor for part 1

• FFN_Expand (int) – channel expansion factor for part 2

• drop_out_rate (float) – drop out ratio

forward(inp)

Forward Function.

Parameters

inp – input tensor image