"""图像级缺陷分类头。

输出维度由 `label_mode` 决定，见 `models.defect_net.DefectNet`：
    multilabel: num_classes      —— 每一维独立 sigmoid，全 0 表示 OK 片
    multiclass: num_classes + 1  —— softmax 互斥，最后一维为 "无缺陷"
"""

from __future__ import annotations

import torch
from torch import nn

from models.registry import HEADS


@HEADS.register("classification")
class ClassificationHead(nn.Module):
    def __init__(self, in_channels: int, num_outputs: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.drop = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
        self.fc = nn.Linear(in_channels, num_outputs)
        self.num_outputs = num_outputs

    def forward(self, feat: torch.Tensor) -> torch.Tensor:
        x = self.pool(feat).flatten(1)
        return self.fc(self.drop(x))