75 lines
2.8 KiB
Python
75 lines
2.8 KiB
Python
"""D-DFFNet Model Definition - PDnet_ResNest_4F_SL_LA_D
|
|
|
|
Refactored for inference-only, removed resnest.torch dependency.
|
|
Original source: D-DFFNet/model/ours/PDNet_resnest.py
|
|
"""
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from .RFB import RFB
|
|
from .aggregation import aggregation_4feat_side_loss_sa_d
|
|
|
|
|
|
class PDnet_ResNest_4F_SL_LA_D(nn.Module):
|
|
"""
|
|
D-DFFNet: Depth-aware Defocus Blur Detection Network.
|
|
|
|
Architecture:
|
|
- Backbone: ResNeSt-101 (dual-branch)
|
|
- Neck: 4x RFB modules (multi-scale feature extraction)
|
|
- Head: aggregation_4feat_side_loss_sa_d (4-level fusion + spatial attention)
|
|
"""
|
|
def __init__(self, channel=32):
|
|
super(PDnet_ResNest_4F_SL_LA_D, self).__init__()
|
|
from .backbone.resnet import ResNest101
|
|
|
|
self.resnest = ResNest101(
|
|
[3, 4, 23, 3],
|
|
radix=2, groups=1, bottleneck_width=64,
|
|
deep_stem=True, stem_width=64, avg_down=True,
|
|
avd=True, avd_first=False
|
|
)
|
|
self.rfb1_1 = RFB(256, channel)
|
|
self.rfb2_1 = RFB(512, channel)
|
|
self.rfb3_1 = RFB(1024, channel)
|
|
self.rfb4_1 = RFB(2048, channel)
|
|
self.agg1 = aggregation_4feat_side_loss_sa_d(channel)
|
|
self.upsample = nn.Upsample(scale_factor=4, mode='bilinear', align_corners=True)
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Args:
|
|
x: input image tensor, shape (B, 3, H, W), normalized with ImageNet stats
|
|
Returns:
|
|
out: defocus blur detection map, shape (B, 1, H*4, W*4) after upsample
|
|
feature_map: list of 4 backbone features (for visualization/distillation)
|
|
side: list of 4 side outputs (for training loss)
|
|
depth: list of 4 depth side outputs (for training loss)
|
|
"""
|
|
x = self.resnest.conv1(x)
|
|
x = self.resnest.bn1(x)
|
|
x = self.resnest.relu(x)
|
|
|
|
x = self.resnest.maxpool(x) # [B, 128, 80, 80]
|
|
x1 = self.resnest.layer1(x) # [B, 256, 80, 80]
|
|
x2 = self.resnest.layer2(x1) # [B, 512, 40, 40]
|
|
x3 = self.resnest.layer3_1(x2) # [B, 1024, 20, 20]
|
|
x4 = self.resnest.layer4_1(x3) # [B, 2048, 10, 10]
|
|
|
|
feature_map = []
|
|
feature_map.append(x4)
|
|
feature_map.append(x3)
|
|
feature_map.append(x2)
|
|
feature_map.append(x1)
|
|
|
|
# Apply RFB modules to compress channels and extract multi-scale features
|
|
x1 = self.rfb1_1(x1) # [B, 32, 80, 80]
|
|
x2 = self.rfb2_1(x2) # [B, 32, 40, 40]
|
|
x3 = self.rfb3_1(x3) # [B, 32, 20, 20]
|
|
x4 = self.rfb4_1(x4) # [B, 32, 10, 10]
|
|
|
|
out, side, depth = self.agg1(x4, x3, x2, x1)
|
|
out = self.upsample(out) # [B, 1, 320, 320] (for 320x320 input)
|
|
|
|
return out, feature_map, side, depth
|