517 lines
21 KiB
Python
517 lines
21 KiB
Python
"""
|
||
Enhanced Transformer Segmentation Head for BEVFusion + RMT-PPAD Integration
|
||
|
||
融合RMT-PPAD的Transformer分割解码器，提升分割性能：
|
||
- TransformerSegmentationDecoder: 自适应多尺度融合
|
||
- TaskAdapterLite: 轻量级任务适配
|
||
- LiteDynamicGate: 动态门控机制 (可选)
|
||
- 保留BEVFusion的ASPP和注意力机制
|
||
"""
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
from typing import Dict, List, Optional, Any, Union
|
||
|
||
from mmdet3d.models import HEADS
|
||
|
||
# 导入RMT-PPAD的核心组件
|
||
from mmdet3d.models.modules.rmtppad_integration import (
|
||
TransformerSegmentationDecoder,
|
||
TaskAdapterLite,
|
||
LiteDynamicGate
|
||
)
|
||
|
||
|
||
@HEADS.register_module()
|
||
class EnhancedTransformerSegmentationHead(nn.Module):
|
||
"""
|
||
Enhanced Transformer Segmentation Head with RMT-PPAD Integration
|
||
|
||
架构设计：
|
||
单尺度: Input BEV → BEVGridTransform → ASPP → Attention → TaskAdapter → TransformerDecoder → Output
|
||
多尺度: Input Multi-Scale → TransformerDecoder → Output
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
in_channels: int,
|
||
grid_transform: Dict[str, Any],
|
||
classes: List[str],
|
||
loss: str = "focal",
|
||
loss_weight: Optional[Dict[str, float]] = None,
|
||
deep_supervision: bool = True,
|
||
use_dice_loss: bool = True,
|
||
dice_weight: float = 0.5,
|
||
focal_alpha: float = 0.25,
|
||
focal_gamma: float = 2.0,
|
||
# RMT-PPAD集成参数
|
||
transformer_hidden_dim: int = 256,
|
||
transformer_C: int = 64,
|
||
transformer_num_layers: int = 2,
|
||
use_task_adapter: bool = True,
|
||
use_dynamic_gate: bool = False,
|
||
gate_reduction: int = 8,
|
||
adapter_reduction: int = 4,
|
||
# 兼容BEVFusion参数
|
||
use_internal_gca: bool = False,
|
||
internal_gca_reduction: int = 4,
|
||
debug_mode: bool = False, # 🔧 调试模式
|
||
) -> None:
|
||
super().__init__()
|
||
|
||
# Phase 4B: 设置标志，跳过checkpoint加载 (使用随机初始化)
|
||
self._phase4b_skip_checkpoint = True
|
||
|
||
# 🔧 调试模式设置
|
||
self.debug_mode = debug_mode
|
||
self.debug_print = lambda *args, **kwargs: print(*args, **kwargs) if self.debug_mode else None
|
||
|
||
self.in_channels = in_channels
|
||
self.classes = classes
|
||
self.loss_type = loss
|
||
self.deep_supervision = deep_supervision
|
||
self.use_dice_loss = use_dice_loss
|
||
self.dice_weight = dice_weight
|
||
self.focal_alpha = focal_alpha
|
||
self.focal_gamma = focal_gamma
|
||
|
||
# RMT-PPAD参数
|
||
self.transformer_hidden_dim = transformer_hidden_dim
|
||
self.transformer_C = transformer_C
|
||
self.transformer_num_layers = transformer_num_layers
|
||
self.use_task_adapter = use_task_adapter
|
||
self.use_dynamic_gate = use_dynamic_gate
|
||
self.gate_reduction = gate_reduction
|
||
self.adapter_reduction = adapter_reduction
|
||
|
||
# BEVFusion兼容参数
|
||
self.use_internal_gca = use_internal_gca
|
||
self.internal_gca_reduction = internal_gca_reduction
|
||
|
||
# Default class weights (针对nuScenes的类别不平衡)
|
||
if loss_weight is None:
|
||
self.loss_weight = {
|
||
'drivable_area': 1.0, # 大类别，基础权重
|
||
'ped_crossing': 3.0, # 小类别，增加权重
|
||
'walkway': 1.5, # 中等类别
|
||
'stop_line': 4.0, # 最小类别，最高权重
|
||
'carpark_area': 2.0, # 小类别
|
||
'divider': 5.0, # ✨ 增强divider权重（线性特征最难分割）
|
||
}
|
||
else:
|
||
self.loss_weight = loss_weight
|
||
|
||
# BEV Grid Transform (保留BEVFusion的坐标变换)
|
||
from mmdet3d.models.heads.segm.vanilla import BEVGridTransform
|
||
self.transform = BEVGridTransform(**grid_transform)
|
||
|
||
# ASPP for multi-scale features (保留BEVFusion的ASPP)
|
||
self.aspp = self._build_aspp(in_channels, transformer_hidden_dim)
|
||
|
||
# Channel and Spatial Attention (保留BEVFusion的注意力)
|
||
self.channel_attn = self._build_channel_attention(transformer_hidden_dim)
|
||
self.spatial_attn = self._build_spatial_attention()
|
||
|
||
# ✨ RMT-PPAD集成：TaskAdapterLite (轻量级任务适配)
|
||
if self.use_task_adapter:
|
||
self.task_adapter = TaskAdapterLite(transformer_hidden_dim)
|
||
print(f"[EnhancedTransformerSegmentationHead] ✨ TaskAdapterLite enabled")
|
||
else:
|
||
self.task_adapter = None
|
||
|
||
# ✨ RMT-PPAD集成：LiteDynamicGate (可选的动态门控)
|
||
if self.use_dynamic_gate:
|
||
self.dynamic_gate = LiteDynamicGate(transformer_hidden_dim, reduction=self.gate_reduction)
|
||
print(f"[EnhancedTransformerSegmentationHead] ✨ LiteDynamicGate enabled (reduction={self.gate_reduction})")
|
||
else:
|
||
self.dynamic_gate = None
|
||
|
||
# ✨ RMT-PPAD集成：TransformerSegmentationDecoder (核心)
|
||
self.transformer_decoder = TransformerSegmentationDecoder(
|
||
hidden_dim=transformer_hidden_dim,
|
||
nc=len(classes),
|
||
C=transformer_C,
|
||
nhead=8, # 默认8头注意力
|
||
num_layers=transformer_num_layers
|
||
)
|
||
print(f"[EnhancedTransformerSegmentationHead] ✨ TransformerSegmentationDecoder enabled")
|
||
print(f" - hidden_dim: {transformer_hidden_dim}")
|
||
print(f" - C: {transformer_C}")
|
||
print(f" - num_layers: {transformer_num_layers}")
|
||
|
||
# BEVFusion兼容：可选的内部GCA (为了向后兼容)
|
||
if self.use_internal_gca:
|
||
from mmdet3d.models.modules.gca import GCA
|
||
self.gca = GCA(
|
||
in_channels=transformer_hidden_dim,
|
||
reduction=self.internal_gca_reduction
|
||
)
|
||
print(f"[EnhancedTransformerSegmentationHead] ⚪ Internal GCA enabled (reduction={self.internal_gca_reduction})")
|
||
else:
|
||
self.gca = None
|
||
|
||
print(f"[EnhancedTransformerSegmentationHead] 🚀 Phase 4B: RMT-PPAD Segmentation Integration Complete")
|
||
|
||
def _build_aspp(self, in_channels: int, out_channels: int) -> nn.Module:
|
||
"""Build ASPP module (保留BEVFusion设计)"""
|
||
class ASPP(nn.Module):
|
||
def __init__(self, in_ch, out_ch, dilation_rates=[6, 12, 18]):
|
||
super().__init__()
|
||
self.convs = nn.ModuleList()
|
||
self.bns = nn.ModuleList()
|
||
|
||
# 1x1 conv
|
||
self.convs.append(nn.Conv2d(in_ch, out_ch, 1, bias=False))
|
||
self.bns.append(nn.GroupNorm(min(32, out_ch), out_ch))
|
||
|
||
# Dilated convs
|
||
for rate in dilation_rates:
|
||
self.convs.append(nn.Conv2d(in_ch, out_ch, 3, padding=rate, dilation=rate, bias=False))
|
||
self.bns.append(nn.GroupNorm(min(32, out_ch), out_ch))
|
||
|
||
# Global pooling
|
||
self.global_avg_pool = nn.AdaptiveAvgPool2d(1)
|
||
self.global_conv = nn.Conv2d(in_ch, out_ch, 1, bias=False)
|
||
self.global_bn = nn.GroupNorm(min(32, out_ch), out_ch)
|
||
|
||
# Project
|
||
num_branches = len(dilation_rates) + 2
|
||
self.project = nn.Sequential(
|
||
nn.Conv2d(out_ch * num_branches, out_ch, 1, bias=False),
|
||
nn.GroupNorm(min(32, out_ch), out_ch),
|
||
nn.ReLU(True),
|
||
nn.Dropout2d(0.1),
|
||
)
|
||
|
||
def forward(self, x):
|
||
res = []
|
||
for conv, bn in zip(self.convs, self.bns):
|
||
res.append(F.relu(bn(conv(x))))
|
||
|
||
# Global context
|
||
global_feat = self.global_avg_pool(x)
|
||
global_feat = F.relu(self.global_bn(self.global_conv(global_feat)))
|
||
global_feat = F.interpolate(global_feat, size=x.shape[-2:], mode='bilinear', align_corners=False)
|
||
res.append(global_feat)
|
||
|
||
# Concatenate and project
|
||
res = torch.cat(res, dim=1)
|
||
return self.project(res)
|
||
|
||
return ASPP(in_channels, out_channels)
|
||
|
||
def _build_channel_attention(self, channels: int) -> nn.Module:
|
||
"""Build Channel Attention module"""
|
||
class ChannelAttention(nn.Module):
|
||
def __init__(self, ch):
|
||
super().__init__()
|
||
self.avg_pool = nn.AdaptiveAvgPool2d(1)
|
||
self.max_pool = nn.AdaptiveMaxPool2d(1)
|
||
self.fc = nn.Sequential(
|
||
nn.Conv2d(ch, ch // 16, 1, bias=False),
|
||
nn.ReLU(True),
|
||
nn.Conv2d(ch // 16, ch, 1, bias=False),
|
||
)
|
||
|
||
def forward(self, x):
|
||
avg_out = self.fc(self.avg_pool(x))
|
||
max_out = self.fc(self.max_pool(x))
|
||
attention = torch.sigmoid(avg_out + max_out)
|
||
return x * attention
|
||
|
||
return ChannelAttention(channels)
|
||
|
||
def _build_spatial_attention(self) -> nn.Module:
|
||
"""Build Spatial Attention module"""
|
||
class SpatialAttention(nn.Module):
|
||
def __init__(self, kernel_size=7):
|
||
super().__init__()
|
||
self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
|
||
|
||
def forward(self, x):
|
||
avg_out = torch.mean(x, dim=1, keepdim=True)
|
||
max_out, _ = torch.max(x, dim=1, keepdim=True)
|
||
attention = torch.cat([avg_out, max_out], dim=1)
|
||
attention = torch.sigmoid(self.conv(attention))
|
||
return x * attention
|
||
|
||
return SpatialAttention()
|
||
|
||
def forward(
|
||
self,
|
||
x: Union[torch.Tensor, List[torch.Tensor]],
|
||
target: Optional[torch.Tensor] = None,
|
||
) -> Union[torch.Tensor, Dict[str, Any]]:
|
||
"""
|
||
Forward pass with RMT-PPAD integration
|
||
|
||
Args:
|
||
x: Input BEV features, either:
|
||
- Single tensor: shape (B, C, H, W)
|
||
- Multi-scale list: [S3_180, S4_360, S5_600]
|
||
target: Ground truth masks, shape (B, num_classes, H_out, W_out)
|
||
|
||
Returns:
|
||
If training: Dict of losses
|
||
If testing: Predicted masks, shape (B, num_classes, H_out, W_out)
|
||
"""
|
||
# DEBUG: 输入特征日志
|
||
# print(f"[分割头] 🎯 BEV特征输入分析:")
|
||
if isinstance(x, (list, tuple)):
|
||
# print(f"[分割头] 输入类型: 多尺度列表，包含 {len(x)} 个尺度")
|
||
for i, scale_feat in enumerate(x):
|
||
# print(f"[分割头] 尺度{i}: {scale_feat.shape} (B={scale_feat.shape[0]}, C={scale_feat.shape[1]}, H={scale_feat.shape[2]}, W={scale_feat.shape[3]})")
|
||
if i == 0: # 只分析第一个尺度的详细信息
|
||
# print(f"[分割头] 设备: {scale_feat.device}, 数据类型: {scale_feat.dtype}")
|
||
# print(f"[分割头] 值范围: [{scale_feat.min():.4f}, {scale_feat.max():.4f}], 均值: {scale_feat.mean():.4f}")
|
||
# print(f"[分割头] 空间尺寸: {scale_feat.shape[2]}×{scale_feat.shape[3]} 像素")
|
||
# 使用第一个尺度进行特征结构分析
|
||
if scale_feat.shape[1] == 512:
|
||
# print(f"[分割头] 特征结构: 512通道 = 256(相机) + 256(LiDAR)")
|
||
pass
|
||
elif scale_feat.shape[1] == 256:
|
||
# print(f"[分割头] 特征结构: 256通道 (单模态)")
|
||
pass
|
||
else:
|
||
# print(f"[分割头] 特征结构: {scale_feat.shape[1]}通道 (未知结构)")
|
||
pass
|
||
else:
|
||
# 单尺度输入
|
||
# print(f"[分割头] 输入类型: 单尺度张量")
|
||
# print(f"[分割头] 形状: {x.shape} (B={x.shape[0]}, C={x.shape[1]}, H={x.shape[2]}, W={x.shape[3]})")
|
||
# print(f"[分割头] 设备: {x.device}, 数据类型: {x.dtype}")
|
||
# print(f"[分割头] 值范围: [{x.min():.4f}, {x.max():.4f}], 均值: {x.mean():.4f}")
|
||
# print(f"[分割头] 空间尺寸: {x.shape[2]}×{x.shape[3]} 像素")
|
||
|
||
# 分析BEV特征的结构
|
||
if x.shape[1] == 512:
|
||
# print(f"[分割头] 特征结构: 512通道 = 256(相机) + 256(LiDAR)")
|
||
pass
|
||
elif x.shape[1] == 256:
|
||
# print(f"[分割头] 特征结构: 256通道 (单模态)")
|
||
pass
|
||
else:
|
||
# print(f"[分割头] 特征结构: {x.shape[1]}通道 (未知结构)")
|
||
pass
|
||
|
||
# Phase 4B 设计：decoder输出 + ASPP + Attention + RMT-PPAD组件
|
||
|
||
# 1. BEVFusion 基础处理
|
||
x = self.transform(x) # BEV Grid Transform
|
||
# print(f"[分割头] BEV变换后: shape={x.shape}")
|
||
|
||
x = self.aspp(x) # ASPP Multi-scale Features
|
||
# print(f"[分割头] ASPP后: shape={x.shape}")
|
||
|
||
x = self.channel_attn(x) # Channel Attention
|
||
# print(f"[分割头] 通道注意力后: shape={x.shape}")
|
||
|
||
x = self.spatial_attn(x) # Spatial Attention
|
||
# print(f"[分割头] 空间注意力后: shape={x.shape}, range=[{x.min():.4f}, {x.max():.4f}]")
|
||
|
||
# 2. RMT-PPAD 增强组件
|
||
if self.gca is not None:
|
||
x = self.gca(x) # Internal GCA
|
||
|
||
if self.dynamic_gate is not None:
|
||
gate_weights = self.dynamic_gate(x)
|
||
x = x * gate_weights # Dynamic Gate
|
||
|
||
if self.task_adapter is not None:
|
||
x = self.task_adapter(x) # TaskAdapterLite
|
||
|
||
# 3. Transformer 分割解码器
|
||
x_multi_scale = self._prepare_multi_scale_features(x)
|
||
# print(f"[分割头] 多尺度特征: type={type(x_multi_scale)}, len={len(x_multi_scale) if isinstance(x_multi_scale, (list, tuple)) else 'N/A'}")
|
||
if isinstance(x_multi_scale, (list, tuple)):
|
||
for i, feat in enumerate(x_multi_scale):
|
||
# print(f"[分割头] 尺度{i}: shape={feat.shape}, device={feat.device}")
|
||
pass
|
||
else:
|
||
pass
|
||
|
||
# 使用第一个尺度确定目标图像尺寸
|
||
first_scale = x_multi_scale[0] if isinstance(x_multi_scale, (list, tuple)) else x_multi_scale
|
||
target_img_size = self._get_target_image_size(first_scale)
|
||
# print(f"[分割头] 目标图像尺寸: {target_img_size}")
|
||
|
||
# Transformer解码 - 多尺度输入
|
||
# print(f"[分割头] 进入Transformer解码器...")
|
||
seg_masks, aux_list = self.transformer_decoder(x_multi_scale, target_img_size)
|
||
# print(f"[分割头] Transformer输出: seg_masks.shape={seg_masks.shape}, aux_list长度={len(aux_list) if aux_list else 0}")
|
||
# print(f"[分割头] 预测范围: min={seg_masks.min():.4f}, max={seg_masks.max():.4f}")
|
||
|
||
# 处理输出格式
|
||
if self.training:
|
||
if target is None:
|
||
raise ValueError("target (gt_masks_bev) is None during training!")
|
||
|
||
# print(f"[分割头] 训练模式 - GT目标: shape={target.shape}, device={target.device}")
|
||
# print(f"[分割头] GT范围: min={target.min():.4f}, max={target.max():.4f}")
|
||
|
||
# 确保输出尺寸与target匹配
|
||
if seg_masks.shape[-2:] != target.shape[-2:]:
|
||
# print(f"[分割头] 尺寸不匹配: 预测{seg_masks.shape[-2:]} vs 目标{target.shape[-2:]}，进行插值")
|
||
seg_masks = F.interpolate(seg_masks, size=target.shape[-2:], mode='bilinear', align_corners=False)
|
||
# print(f"[分割头] 插值后: shape={seg_masks.shape}")
|
||
|
||
# print(f"[分割头] 开始计算loss...")
|
||
losses = self._compute_loss(seg_masks, target, aux_list)
|
||
# print(f"[分割头] Loss计算完成: {list(losses.keys())}")
|
||
return losses
|
||
else:
|
||
# print(f"[分割头] 推理模式 - 应用sigmoid")
|
||
final_output = torch.sigmoid(seg_masks)
|
||
# print(f"[分割头] 最终输出: shape={final_output.shape}, range=[{final_output.min():.4f}, {final_output.max():.4f}]")
|
||
return final_output
|
||
|
||
def _prepare_multi_scale_features(self, x: torch.Tensor, bev_180_features: List[torch.Tensor] = None) -> List[torch.Tensor]:
|
||
"""
|
||
Phase 4B 多尺度方案：基于输入尺寸动态生成多尺度BEV特征
|
||
|
||
核心设计：以360×360为基准，生成[180×180, 360×360, 600×600]三尺度
|
||
这样与epoch2的BEV输入尺寸保持一致
|
||
|
||
Args:
|
||
x: BEV特征 (经过BEVFusion处理 + RMT-PPAD组件)，尺寸可能为360×360
|
||
bev_180_features: 未使用
|
||
|
||
Returns:
|
||
多尺度BEV特征列表用于Transformer解码器 [180, 360, 600]
|
||
"""
|
||
h, w = x.shape[2], x.shape[3]
|
||
print(f"[多尺度] 输入BEV: shape={x.shape} (基准尺寸: {h}×{w})")
|
||
|
||
# Phase 4B: 基于360×360基准生成三尺度 [180, 360, 600]
|
||
# 这样与epoch2的BEV输入尺寸保持一致
|
||
multi_scale_features = []
|
||
|
||
# 尺度1: 180×180 (基准的0.5倍)
|
||
scale_180 = F.interpolate(x, scale_factor=0.5, mode='bilinear', align_corners=False)
|
||
multi_scale_features.append(scale_180)
|
||
print(f"[多尺度] 180×180: shape={scale_180.shape}")
|
||
|
||
# 尺度2: 360×360 (基准尺寸)
|
||
scale_360 = x
|
||
multi_scale_features.append(scale_360)
|
||
print(f"[多尺度] 360×360: shape={scale_360.shape}")
|
||
|
||
# 尺度3: 600×600 (基准的1.667倍)
|
||
scale_600 = F.interpolate(x, scale_factor=600/360, mode='bilinear', align_corners=False)
|
||
multi_scale_features.append(scale_600)
|
||
print(f"[多尺度] 600×600: shape={scale_600.shape}")
|
||
|
||
# Phase 4B: 返回三尺度特征列表 [180, 360, 600]
|
||
print("[多尺度] 返回三尺度特征列表 [180, 360, 600]")
|
||
return [scale_180, scale_360, scale_600]
|
||
|
||
def _get_target_image_size(self, x: torch.Tensor) -> int:
|
||
"""
|
||
获取目标图像尺寸 (用于Transformer解码器的上采样)
|
||
|
||
BEVFusion中BEV从180x180上采样到360x360作为最终尺寸
|
||
"""
|
||
# BEVFusion的标准BEV尺寸是360x360
|
||
target_size = 360
|
||
return target_size
|
||
|
||
def _compute_loss(
|
||
self,
|
||
pred: torch.Tensor,
|
||
target: torch.Tensor,
|
||
aux_pred: Optional[List[torch.Tensor]] = None,
|
||
) -> Dict[str, torch.Tensor]:
|
||
"""Compute losses with class weighting and optional dice loss"""
|
||
losses = {}
|
||
|
||
for idx, name in enumerate(self.classes):
|
||
pred_cls = pred[:, idx]
|
||
target_cls = target[:, idx]
|
||
|
||
# Main Focal Loss (with alpha for class balance)
|
||
focal_loss = sigmoid_focal_loss(
|
||
pred_cls,
|
||
target_cls,
|
||
alpha=self.focal_alpha,
|
||
gamma=self.focal_gamma,
|
||
)
|
||
|
||
# Dice Loss (better for small objects)
|
||
total_loss = focal_loss
|
||
if self.use_dice_loss:
|
||
dice = dice_loss(pred_cls, target_cls)
|
||
total_loss = focal_loss + self.dice_weight * dice
|
||
# 存储dice系数而不是dice loss，用于监控
|
||
dice_coeff = 1 - dice
|
||
losses[f"{name}/dice"] = dice_coeff
|
||
|
||
# Apply class-specific weight
|
||
class_weight = self.loss_weight.get(name, 1.0)
|
||
losses[f"{name}/focal"] = focal_loss * class_weight
|
||
|
||
# Auxiliary Loss from Transformer Decoder (deep supervision)
|
||
if aux_pred is not None and len(aux_pred) > idx:
|
||
# Resize aux prediction to match target
|
||
aux_pred_cls = aux_pred[idx]
|
||
if aux_pred_cls.shape[-2:] != target_cls.shape[-2:]:
|
||
aux_pred_cls = F.interpolate(aux_pred_cls, size=target_cls.shape[-2:], mode='nearest')
|
||
|
||
aux_focal = sigmoid_focal_loss(
|
||
aux_pred_cls,
|
||
target_cls,
|
||
alpha=self.focal_alpha,
|
||
gamma=self.focal_gamma,
|
||
)
|
||
losses[f"{name}/aux_focal"] = aux_focal * class_weight * 0.3 # 较低权重
|
||
|
||
return losses
|
||
|
||
|
||
# 保留原有的损失函数
|
||
def sigmoid_focal_loss(
|
||
inputs: torch.Tensor,
|
||
targets: torch.Tensor,
|
||
alpha: float = 0.25,
|
||
gamma: float = 2.0,
|
||
reduction: str = "mean",
|
||
) -> torch.Tensor:
|
||
"""Focal Loss with class balancing"""
|
||
inputs = inputs.float()
|
||
targets = targets.float()
|
||
|
||
p = torch.sigmoid(inputs)
|
||
ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
|
||
|
||
p_t = p * targets + (1 - p) * (1 - targets)
|
||
focal_weight = (1 - p_t) ** gamma
|
||
loss = ce_loss * focal_weight
|
||
|
||
alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
|
||
loss = alpha_t * loss
|
||
|
||
if reduction == "mean":
|
||
loss = loss.mean()
|
||
elif reduction == "sum":
|
||
loss = loss.sum()
|
||
|
||
return loss
|
||
|
||
|
||
def dice_loss(
|
||
pred: torch.Tensor,
|
||
target: torch.Tensor,
|
||
smooth: float = 1.0,
|
||
) -> torch.Tensor:
|
||
"""Dice Loss for better small object segmentation"""
|
||
pred = torch.sigmoid(pred)
|
||
|
||
pred_flat = pred.reshape(pred.shape[0], -1)
|
||
target_flat = target.reshape(target.shape[0], -1)
|
||
|
||
intersection = (pred_flat * target_flat).sum(dim=1)
|
||
union = pred_flat.sum(dim=1) + target_flat.sum(dim=1)
|
||
dice_coeff = (2.0 * intersection + smooth) / (union + smooth)
|
||
|
||
return 1 - dice_coeff.mean()
|