mirror of
https://github.com/deepfakes/faceswap.git
synced 2026-10-20 16:45:00 +08:00
torzdf
915 lines
34 KiB
Python
915 lines
34 KiB
Python
#!/usr/bin/env python3
|
|
"""Facial landmarks extractor for faceswap.pnt_y
|
|
Code adapted and modified from:
|
|
https://github.com/1adrianb/face-alignment
|
|
"""
|
|
from __future__ import annotations
|
|
import logging
|
|
import typing as T
|
|
from dataclasses import dataclass
|
|
|
|
import numpy as np
|
|
|
|
import torch
|
|
from torch import nn
|
|
from torch.nn import functional as F
|
|
|
|
from lib.utils import get_module_objects, GetModel
|
|
from plugins.extract.base import ExtractPlugin
|
|
from . import hrnet_defaults as cfg
|
|
from .dark_decoder import Dark
|
|
|
|
if T.TYPE_CHECKING:
|
|
import numpy.typing as npt
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
# pylint:disable=duplicate-code
|
|
|
|
|
|
@dataclass
|
|
class HRNetStageConfig:
|
|
"""Configuration settings for each stage of HRNet"""
|
|
num_modules: int
|
|
num_branches: int
|
|
num_blocks: list[int]
|
|
num_channels: list[int]
|
|
block: T.Literal["ATTENTION", "BLOCK", "BOTTLENECK"]
|
|
|
|
|
|
class HRNet(ExtractPlugin):
|
|
"""HRNet Face alignment"""
|
|
def __init__(self) -> None:
|
|
super().__init__(input_size=256, # if cfg.weights() == "standard" else 512,
|
|
batch_size=cfg.batch_size(),
|
|
is_rgb=True,
|
|
dtype="float32",
|
|
scale=(0, 1))
|
|
self._stage_configs = self._get_stage_configs()
|
|
self._hm_size = 64 # if cfg.weights() == "standard" else 128
|
|
self.model: HighResolutionNet
|
|
self.realign_centering = "legacy"
|
|
|
|
self._mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
|
self._std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
|
self._dark = Dark(68, self._hm_size) if cfg.dark_decoder() else None
|
|
|
|
@classmethod
|
|
def _get_stage_configs(cls) -> tuple[HRNetStageConfig, HRNetStageConfig, HRNetStageConfig]:
|
|
"""Obtain the model configuration for the chosen weights
|
|
|
|
Returns
|
|
-------
|
|
The model configuration
|
|
"""
|
|
# weights = cfg.weights()
|
|
weights = "standard"
|
|
if weights == "standard":
|
|
retval = (HRNetStageConfig(num_modules=1,
|
|
num_branches=2,
|
|
num_blocks=[4, 4],
|
|
num_channels=[18, 36],
|
|
block="BLOCK"),
|
|
HRNetStageConfig(num_modules=4,
|
|
num_branches=3,
|
|
num_blocks=[4, 4, 4],
|
|
num_channels=[18, 36, 72],
|
|
block="ATTENTION"),
|
|
HRNetStageConfig(num_modules=3,
|
|
num_branches=4,
|
|
num_blocks=[4, 4, 4, 4],
|
|
num_channels=[18, 36, 72, 144],
|
|
block="ATTENTION"))
|
|
else:
|
|
retval = (HRNetStageConfig(num_modules=1,
|
|
num_branches=2,
|
|
num_blocks=[4, 4],
|
|
num_channels=[32, 64],
|
|
block="BLOCK"),
|
|
HRNetStageConfig(num_modules=4,
|
|
num_branches=3,
|
|
num_blocks=[4, 4, 4],
|
|
num_channels=[32, 64, 128],
|
|
block="ATTENTION"),
|
|
HRNetStageConfig(num_modules=3,
|
|
num_branches=4,
|
|
num_blocks=[4, 4, 4, 4],
|
|
num_channels=[32, 64, 128, 256],
|
|
block="ATTENTION"))
|
|
logger.debug("[HRNet] using config for weights '%s': %s", weights, retval)
|
|
return retval
|
|
|
|
def load_model(self) -> HighResolutionNet:
|
|
"""Load the HRNet model
|
|
|
|
Returns
|
|
-------
|
|
The loaded HRNet model
|
|
"""
|
|
# version = 2 if cfg.weights() == "standard" else 3
|
|
version = 2
|
|
weights = GetModel(f"hrnet_landmark_v{version}.pth", 34).model_path
|
|
assert isinstance(weights, str)
|
|
model = T.cast(HighResolutionNet, self.load_torch_model(
|
|
HighResolutionNet(num_joints=68,
|
|
final_conv_kernel=1,
|
|
stage_2_config=self._stage_configs[0],
|
|
stage_3_config=self._stage_configs[1],
|
|
stage_4_config=self._stage_configs[2]),
|
|
weights))
|
|
return model
|
|
|
|
def pre_process(self, batch: np.ndarray) -> np.ndarray:
|
|
"""Format the ROI faces detection boxes for prediction
|
|
|
|
Parameters
|
|
----------
|
|
batch
|
|
The batch of face detection bounding boxes as (bs, l, t, r, b)
|
|
|
|
Returns
|
|
-------
|
|
The face detection bounding boxes formatted to take an image patch for prediction
|
|
"""
|
|
heights = batch[:, 3] - batch[:, 1]
|
|
widths = batch[:, 2] - batch[:, 0]
|
|
ctr_x = np.rint((batch[:, 0] + batch[:, 2]) * 0.5).astype("int32")
|
|
ctr_y = np.rint((batch[:, 1] + batch[:, 3]) * 0.5).astype("int32")
|
|
size = np.maximum(widths, heights) * 1.25
|
|
half = np.rint(size * 0.5).astype("int32")
|
|
retval = np.empty((batch.shape[0], 4), dtype=np.int32)
|
|
retval[:, 0] = ctr_x - half
|
|
retval[:, 1] = ctr_y - half
|
|
retval[:, 2] = ctr_x + half
|
|
retval[:, 3] = ctr_y + half
|
|
return retval
|
|
|
|
def process(self, batch: np.ndarray) -> np.ndarray:
|
|
"""Predict the 68 point landmarks
|
|
|
|
Parameters
|
|
----------
|
|
batch
|
|
The batch to feed into the aligner
|
|
|
|
Returns
|
|
-------
|
|
The predictions from the aligner
|
|
"""
|
|
batch -= self._mean
|
|
batch /= self._std
|
|
return self.from_torch(batch.transpose(0, 3, 1, 2))
|
|
|
|
def _get_predictions(self, scores: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
|
|
"""Convert the score maps from the model into predictions
|
|
|
|
Parameters
|
|
----------
|
|
scores
|
|
The score maps from the model
|
|
|
|
Returns
|
|
-------
|
|
Predictions from the score maps
|
|
"""
|
|
batch_size, num_points, size = scores.shape[:3]
|
|
scores_r = scores.reshape((batch_size, num_points, -1))
|
|
|
|
max_val: npt.NDArray[np.float32] = scores_r.max(axis=2)
|
|
idx: npt.NDArray[np.int64] = scores_r.argmax(axis=2)
|
|
|
|
retval = np.empty((batch_size, num_points, 2), dtype=np.float32)
|
|
retval[:, :, 0] = (idx - 1) % size + 1
|
|
retval[:, :, 1] = np.floor((idx - 1) / size) + 1
|
|
|
|
mask = max_val[..., None] > 0.
|
|
retval *= mask
|
|
return retval
|
|
|
|
def post_process(self, batch: np.ndarray) -> np.ndarray: # pylint:disable=too-many-locals
|
|
"""Process the output from the model
|
|
|
|
Parameters
|
|
----------
|
|
batch
|
|
The predictions from the aligner
|
|
|
|
Returns
|
|
-------
|
|
The final landmarks in 0-1 space
|
|
"""
|
|
if self._dark is not None:
|
|
return self._dark(batch) / self._hm_size
|
|
batch_size, num_points, height, width = batch.shape
|
|
assert height == width, "Heatmaps must be square"
|
|
resolution = height
|
|
|
|
coords = self._get_predictions(batch)
|
|
pnt_x = coords[..., 0].astype(np.int32)
|
|
pnt_y = coords[..., 1].astype(np.int32)
|
|
mask = (pnt_x > 1) & (pnt_x < resolution) & (pnt_y > 1) & (pnt_y < resolution)
|
|
pnt_x = np.clip(pnt_x, 2, resolution - 1)
|
|
pnt_y = np.clip(pnt_x, 2, resolution - 1)
|
|
idx_batch = np.arange(batch_size)[:, None]
|
|
idx_pnt = np.arange(num_points)[None, :]
|
|
delta_x = (batch[idx_batch, idx_pnt, pnt_y - 1, pnt_x] -
|
|
batch[idx_batch, idx_pnt, pnt_y - 1, pnt_x - 2])
|
|
delta_y = (batch[idx_batch, idx_pnt, pnt_y, pnt_x - 1] -
|
|
batch[idx_batch, idx_pnt, pnt_y - 2, pnt_x - 1])
|
|
diff = np.stack([delta_x, delta_y], axis=-1)
|
|
coords += (np.sign(diff) * 0.25 * mask[..., None]) + 0.5
|
|
coords /= resolution
|
|
return coords
|
|
|
|
|
|
class BasicBlock(nn.Module):
|
|
""" Basic block for HRNet
|
|
|
|
Parameters
|
|
----------
|
|
in_channels
|
|
The number of in channels
|
|
out_channels
|
|
The number of out channels
|
|
stride
|
|
The stride for the first 3x3 conv block. Default: 1
|
|
downsample
|
|
The module to use for downsampling or ``None`` for no downsample. Default: ``None``
|
|
"""
|
|
expansion = 1
|
|
|
|
def __init__(self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
stride: int = 1,
|
|
downsample: nn.Module | None = None) -> None:
|
|
super().__init__()
|
|
self.conv1 = nn.Conv2d(in_channels, out_channels, 3, stride=stride, padding=1, bias=False)
|
|
self.bn1 = nn.BatchNorm2d(out_channels, momentum=0.01)
|
|
self.relu = nn.ReLU(inplace=True)
|
|
self.conv2 = nn.Conv2d(in_channels, out_channels, 3, stride=1, padding=1, bias=False)
|
|
self.bn2 = nn.BatchNorm2d(out_channels, momentum=0.01)
|
|
self.downsample = downsample
|
|
self.stride = stride
|
|
|
|
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
|
|
"""Forward pass through HRNet's basic block
|
|
|
|
Parameters
|
|
----------
|
|
inputs
|
|
Input to the conv block
|
|
|
|
Returns
|
|
-------
|
|
Output from the conv block
|
|
"""
|
|
residual = inputs
|
|
|
|
out = self.conv1(inputs)
|
|
out = self.bn1(out)
|
|
out = self.relu(out)
|
|
|
|
out = self.conv2(out)
|
|
out = self.bn2(out)
|
|
|
|
if self.downsample is not None:
|
|
residual = self.downsample(inputs)
|
|
|
|
out += residual
|
|
out = self.relu(out)
|
|
|
|
return out
|
|
|
|
|
|
class BasicBlockAttention(nn.Module): # pylint:disable=too-many-instance-attributes
|
|
""" Custom Basic block for HRNet with Attention
|
|
|
|
Parameters
|
|
----------
|
|
in_channels
|
|
The number of in channels
|
|
out_channels
|
|
The number of out channels
|
|
stride
|
|
The stride for the first 3x3 conv block. Default: 1
|
|
downsample
|
|
The module to use for downsampling or ``None`` for no downsample. Default: ``None``
|
|
"""
|
|
expansion = 1
|
|
|
|
def __init__(self, in_channels, out_channels, stride=1, downsample=None):
|
|
super().__init__()
|
|
self.conv1 = nn.Conv2d(in_channels, out_channels, 3, stride=stride, padding=1, bias=False)
|
|
self.bn1 = nn.BatchNorm2d(out_channels, momentum=0.01)
|
|
self.relu = nn.ReLU(inplace=True)
|
|
self.conv2 = nn.Conv2d(in_channels, out_channels, 3, stride=1, padding=1, bias=False)
|
|
self.bn2 = nn.BatchNorm2d(out_channels, momentum=0.01)
|
|
self.downsample = downsample
|
|
self.att1 = nn.Sequential(nn.Conv2d(in_channels,
|
|
max(1, out_channels // 4),
|
|
1,
|
|
stride=1,
|
|
bias=False),
|
|
nn.BatchNorm2d(out_channels // 4),
|
|
nn.ReLU(inplace=True))
|
|
self.att2 = nn.Conv2d(max(1, out_channels // 4),
|
|
max(1, out_channels // 4),
|
|
3,
|
|
stride=1,
|
|
padding=1,
|
|
bias=False)
|
|
self.att_bn1 = nn.BatchNorm2d(max(1, out_channels // 4), momentum=0.01)
|
|
self.att3 = nn.Conv2d(max(1, out_channels // 4), out_channels, 1, 1, bias=False)
|
|
self.att_bn2 = nn.BatchNorm2d(out_channels, momentum=0.01)
|
|
|
|
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
|
|
"""Forward pass through HRNet's basic block with attention
|
|
|
|
Parameters
|
|
----------
|
|
inputs
|
|
Input to the conv block
|
|
|
|
Returns
|
|
-------
|
|
Output from the conv block
|
|
"""
|
|
residual = inputs
|
|
|
|
out = self.conv1(inputs)
|
|
out = self.bn1(out)
|
|
out = self.relu(out)
|
|
|
|
out = self.conv2(out)
|
|
out = self.bn2(out)
|
|
|
|
att = self.att1(out)
|
|
att = self.att2(att)
|
|
att = self.att_bn1(att)
|
|
att = self.relu(att)
|
|
att = self.att3(att)
|
|
att = self.att_bn2(att)
|
|
att = torch.sigmoid(att)
|
|
out = out * att
|
|
|
|
if self.downsample is not None:
|
|
residual = self.downsample(inputs)
|
|
|
|
out += residual
|
|
out = self.relu(out)
|
|
|
|
return out
|
|
|
|
|
|
class Bottleneck(nn.Module):
|
|
""" Bottleneck for HRNet
|
|
|
|
Parameters
|
|
----------
|
|
in_channels
|
|
The number of in channels
|
|
out_channels
|
|
The number of out channels
|
|
stride
|
|
The stride for the first 3x3 conv block. Default: 1
|
|
downsample
|
|
The module to use for downsampling or ``None`` for no downsample. Default: ``None``
|
|
"""
|
|
expansion = 4
|
|
|
|
def __init__(self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
stride: int = 1,
|
|
downsample: nn.Module | None = None) -> None:
|
|
super().__init__()
|
|
self.conv1 = nn.Conv2d(in_channels, out_channels, 1, bias=False)
|
|
self.bn1 = nn.BatchNorm2d(out_channels, momentum=0.01)
|
|
self.conv2 = nn.Conv2d(out_channels, out_channels, 3, stride=stride, padding=1, bias=False)
|
|
self.bn2 = nn.BatchNorm2d(out_channels, momentum=0.01)
|
|
self.conv3 = nn.Conv2d(out_channels, out_channels * self.expansion, 1, bias=False)
|
|
self.bn3 = nn.BatchNorm2d(out_channels * self.expansion, momentum=0.01)
|
|
self.relu = nn.ReLU(inplace=True)
|
|
self.downsample = downsample
|
|
self.stride = stride
|
|
|
|
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
|
|
"""Forward pass through HRNet's basic block
|
|
|
|
Parameters
|
|
----------
|
|
inputs
|
|
Input to the conv block
|
|
|
|
Returns
|
|
-------
|
|
Output from the conv block
|
|
"""
|
|
residual = inputs
|
|
|
|
out = self.conv1(inputs)
|
|
out = self.bn1(out)
|
|
out = self.relu(out)
|
|
|
|
out = self.conv2(out)
|
|
out = self.bn2(out)
|
|
out = self.relu(out)
|
|
|
|
out = self.conv3(out)
|
|
out = self.bn3(out)
|
|
|
|
if self.downsample is not None:
|
|
residual = self.downsample(inputs)
|
|
|
|
out += residual
|
|
out = self.relu(out)
|
|
|
|
return out
|
|
|
|
|
|
class HighResolutionModule(nn.Module):
|
|
""" High Resolution Module for HRNet
|
|
|
|
Parameters
|
|
----------
|
|
num_branches
|
|
The number of branches to use
|
|
blocks
|
|
The block object to use
|
|
num_blocks
|
|
The number of blocks in each branch
|
|
num_in_channels
|
|
The number of input channels in each branch
|
|
num_channels
|
|
The number of channels in each branch
|
|
multi_scale_output
|
|
``True`` to output multi-scaled
|
|
"""
|
|
def __init__(self,
|
|
num_branches: int,
|
|
block: type[BasicBlockAttention] | type[BasicBlock] | type[Bottleneck],
|
|
num_blocks: list[int],
|
|
num_in_channels: list[int],
|
|
num_channels: list[int],
|
|
multi_scale_output: bool = True) -> None:
|
|
super().__init__()
|
|
self._check_branches(num_branches,
|
|
num_blocks,
|
|
num_in_channels,
|
|
num_channels)
|
|
|
|
self.num_in_channels = num_in_channels
|
|
self.num_branches = num_branches
|
|
self.multi_scale_output = multi_scale_output
|
|
self.branches = nn.ModuleList(self._make_one_branch(i,
|
|
block,
|
|
num_blocks[i],
|
|
num_channels[i])
|
|
for i in range(num_branches))
|
|
|
|
self.fuse_layers = self._make_fuse_layers()
|
|
self.relu = nn.ReLU(inplace=True)
|
|
|
|
def _check_branches(self,
|
|
num_branches: int,
|
|
num_blocks: list[int],
|
|
num_in_channels: list[int],
|
|
num_channels: list[int]) -> None:
|
|
"""Check that the branch configuration is valid
|
|
|
|
Parameters
|
|
----------
|
|
num_branches
|
|
The number of branches to use
|
|
num_blocks
|
|
The number of blocks in each branch
|
|
num_in_channels
|
|
The number of input channels in each branch
|
|
num_channels
|
|
The number of channels in each branch
|
|
|
|
Raises
|
|
------
|
|
ValueError
|
|
On an invalid configuration
|
|
"""
|
|
if num_branches != len(num_blocks):
|
|
raise ValueError(f"NUM_BRANCHES({num_branches}) <> NUM_BLOCKS({len(num_blocks)})")
|
|
if num_branches != len(num_channels):
|
|
raise ValueError(f"NUM_BRANCHES({num_branches}) <> NUM_CHANNELS({len(num_channels)})")
|
|
if num_branches != len(num_in_channels):
|
|
raise ValueError(f"NUM_BRANCHES({num_branches}) <> "
|
|
f"NUM_IN_CHANNELS({len(num_in_channels)})")
|
|
|
|
def _make_one_branch(self,
|
|
branch_index: int,
|
|
block: type[BasicBlockAttention | BasicBlock | Bottleneck],
|
|
num_blocks: int,
|
|
num_channels: int,
|
|
stride: int = 1) -> nn.Sequential:
|
|
""" Make a single branch
|
|
|
|
Parameters
|
|
----------
|
|
branch_index
|
|
The index of the branch to make
|
|
block
|
|
The block object to use
|
|
num_blocks
|
|
The number of blocks in each branch
|
|
num_in_channels
|
|
The number of input channels in each branch
|
|
num_channels
|
|
The number of channels in each branch
|
|
multi_scale_output
|
|
``True`` to output multi-scaled
|
|
|
|
Returns
|
|
-------
|
|
The sequential modules for the branch
|
|
"""
|
|
downsample = None
|
|
if stride != 1 or self.num_in_channels[branch_index] != num_channels * block.expansion:
|
|
downsample = nn.Sequential(nn.Conv2d(self.num_in_channels[branch_index],
|
|
num_channels * block.expansion,
|
|
1,
|
|
stride=stride,
|
|
bias=False),
|
|
nn.BatchNorm2d(num_channels * block.expansion,
|
|
momentum=0.01))
|
|
|
|
layers = []
|
|
layers.append(block(self.num_in_channels[branch_index],
|
|
num_channels,
|
|
stride,
|
|
downsample=downsample))
|
|
self.num_in_channels[branch_index] = num_channels * block.expansion
|
|
for _ in range(1, num_blocks):
|
|
layers.append(block(self.num_in_channels[branch_index], num_channels))
|
|
|
|
return nn.Sequential(*layers)
|
|
|
|
def _make_fuse_layers(self) -> nn.ModuleList | None:
|
|
"""Make the fuse layers for the HR Module
|
|
|
|
Returns
|
|
-------
|
|
The fuse layers module list or ``None`` if layers are not to be fused
|
|
"""
|
|
if self.num_branches == 1:
|
|
return None
|
|
|
|
num_branches = self.num_branches
|
|
num_in_channels = self.num_in_channels
|
|
fuse_layers = []
|
|
for i in range(num_branches if self.multi_scale_output else 1):
|
|
fuse_layer = []
|
|
for j in range(num_branches):
|
|
if j > i:
|
|
fuse_layer.append(nn.Sequential(nn.Conv2d(num_in_channels[j],
|
|
num_in_channels[i],
|
|
1,
|
|
stride=1,
|
|
padding=0,
|
|
bias=False),
|
|
nn.BatchNorm2d(num_in_channels[i],
|
|
momentum=0.01)))
|
|
elif j == i:
|
|
fuse_layer.append(None) # type:ignore[arg-type]
|
|
else:
|
|
conv3x3s = []
|
|
for k in range(i - j):
|
|
if k == i - j - 1:
|
|
num_out_channels_conv3x3 = num_in_channels[i]
|
|
conv3x3s.append(nn.Sequential(nn.Conv2d(num_in_channels[j],
|
|
num_out_channels_conv3x3,
|
|
3,
|
|
stride=2,
|
|
padding=1,
|
|
bias=False),
|
|
nn.BatchNorm2d(num_out_channels_conv3x3,
|
|
momentum=0.01)))
|
|
else:
|
|
num_out_channels_conv3x3 = num_in_channels[j]
|
|
conv3x3s.append(nn.Sequential(nn.Conv2d(num_in_channels[j],
|
|
num_out_channels_conv3x3,
|
|
3,
|
|
stride=2,
|
|
padding=1,
|
|
bias=False),
|
|
nn.BatchNorm2d(num_out_channels_conv3x3,
|
|
momentum=0.01),
|
|
nn.ReLU(inplace=True)))
|
|
fuse_layer.append(nn.Sequential(*conv3x3s))
|
|
fuse_layers.append(nn.ModuleList(fuse_layer))
|
|
|
|
return nn.ModuleList(fuse_layers)
|
|
|
|
def get_num_in_channels(self) -> list[int]:
|
|
"""Obtain the number of input channels to the module
|
|
|
|
Returns
|
|
-------
|
|
The number of input channels to the module
|
|
"""
|
|
return self.num_in_channels
|
|
|
|
def forward(self, inputs: list[torch.Tensor]) -> list[torch.Tensor]:
|
|
"""Forward pass through the HR Module
|
|
|
|
Parameters
|
|
----------
|
|
inputs
|
|
Input to the HR Module
|
|
|
|
Returns
|
|
-------
|
|
Output from the HR Module
|
|
"""
|
|
x = inputs
|
|
if self.num_branches == 1:
|
|
return [self.branches[0](x[0])]
|
|
assert self.fuse_layers is not None
|
|
|
|
for i in range(self.num_branches):
|
|
x[i] = self.branches[i](x[i])
|
|
|
|
x_fuse = []
|
|
for i, fuse_layer in enumerate(T.cast(list[nn.ModuleList], self.fuse_layers)):
|
|
y = x[0] if i == 0 else fuse_layer[0](x[0])
|
|
for j in range(1, self.num_branches):
|
|
if i == j:
|
|
y = y + x[j]
|
|
elif j > i:
|
|
y = y + F.interpolate(fuse_layer[j](x[j]),
|
|
size=[x[i].shape[2], x[i].shape[3]],
|
|
mode="bilinear")
|
|
else:
|
|
y = y + fuse_layer[j](x[j])
|
|
x_fuse.append(self.relu(y))
|
|
return x_fuse
|
|
|
|
|
|
class HighResolutionNet(nn.Module): # pylint:disable=too-many-instance-attributes
|
|
"""The HRNet Landmark Detection model
|
|
|
|
Parameters
|
|
----------
|
|
num_joints
|
|
The number of joints in the model
|
|
final_conv_kernel
|
|
Kernel size of the final convolution
|
|
stage_2_config
|
|
Configuration settings for stage 2 layers
|
|
stage_3_config
|
|
Configuration settings for stage 3 layers
|
|
stage_4_config
|
|
Configuration settings for stage 4 layers
|
|
"""
|
|
def __init__(self,
|
|
num_joints: int,
|
|
final_conv_kernel: int,
|
|
stage_2_config: HRNetStageConfig,
|
|
stage_3_config: HRNetStageConfig,
|
|
stage_4_config: HRNetStageConfig) -> None:
|
|
self.in_channels = 64
|
|
super().__init__()
|
|
self.stage_2_config = stage_2_config
|
|
self.stage_3_config = stage_3_config
|
|
self.stage_4_config = stage_4_config
|
|
|
|
# stem net
|
|
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1,
|
|
bias=False)
|
|
self.bn1 = nn.BatchNorm2d(64, momentum=0.01)
|
|
self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1,
|
|
bias=False)
|
|
self.bn2 = nn.BatchNorm2d(64, momentum=0.01)
|
|
self.relu = nn.ReLU(inplace=True)
|
|
self.sf = nn.Softmax(dim=1)
|
|
self.layer1 = self._make_layer(Bottleneck, 64, 64, 4)
|
|
|
|
_blocks: dict[str, type[BasicBlockAttention | BasicBlock | Bottleneck]] = {
|
|
"ATTENTION": BasicBlockAttention, "BLOCK": BasicBlock, "BOTTLENECK": Bottleneck}
|
|
|
|
num_channels = stage_2_config.num_channels
|
|
block = _blocks[stage_2_config.block]
|
|
num_channels = [num_channels[i] * block.expansion for i in range(len(num_channels))]
|
|
self.transition1 = self._make_transition_layer([256], num_channels)
|
|
self.stage2, pre_stage_channels = self._make_stage(block, stage_2_config, num_channels)
|
|
|
|
num_channels = stage_3_config.num_channels
|
|
block = _blocks[stage_3_config.block]
|
|
num_channels = [num_channels[i] * block.expansion for i in range(len(num_channels))]
|
|
self.transition2 = self._make_transition_layer(pre_stage_channels, num_channels)
|
|
self.stage3, pre_stage_channels = self._make_stage(block, stage_3_config, num_channels)
|
|
|
|
num_channels = stage_4_config.num_channels
|
|
block = _blocks[stage_4_config.block]
|
|
num_channels = [num_channels[i] * block.expansion for i in range(len(num_channels))]
|
|
self.transition3 = self._make_transition_layer(pre_stage_channels, num_channels)
|
|
self.stage4, pre_stage_channels = self._make_stage(block,
|
|
stage_4_config,
|
|
num_channels,
|
|
multi_scale_output=True)
|
|
|
|
final_inp_channels = sum(pre_stage_channels)
|
|
|
|
self.head = nn.Sequential(nn.Conv2d(final_inp_channels,
|
|
final_inp_channels,
|
|
1,
|
|
stride=1,
|
|
padding=1 if final_conv_kernel == 3 else 0),
|
|
nn.BatchNorm2d(final_inp_channels, momentum=0.01),
|
|
nn.ReLU(inplace=True),
|
|
nn.Conv2d(final_inp_channels,
|
|
num_joints,
|
|
final_conv_kernel,
|
|
stride=1,
|
|
padding=1 if final_conv_kernel == 3 else 0))
|
|
|
|
def _make_transition_layer(self,
|
|
num_channels_pre_layer: list[int],
|
|
num_channels_cur_layer: list[int]) -> nn.ModuleList:
|
|
"""Make an HRNet transition layer
|
|
|
|
Parameters
|
|
----------
|
|
num_channels_pre_layer
|
|
The number of channels from the previous layer
|
|
num_channels_cur_layer
|
|
The number of channels from the current layer
|
|
|
|
Returns
|
|
-------
|
|
The transition layer module list
|
|
"""
|
|
num_branches_cur = len(num_channels_cur_layer)
|
|
num_branches_pre = len(num_channels_pre_layer)
|
|
|
|
transition_layers = []
|
|
for i in range(num_branches_cur):
|
|
if i < num_branches_pre:
|
|
if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
|
|
transition_layers.append(nn.Sequential(
|
|
nn.Conv2d(num_channels_pre_layer[i],
|
|
num_channels_cur_layer[i],
|
|
3,
|
|
stride=1,
|
|
padding=1,
|
|
bias=False),
|
|
nn.BatchNorm2d(num_channels_cur_layer[i], momentum=0.01),
|
|
nn.ReLU(inplace=True)))
|
|
else:
|
|
transition_layers.append(None) # type:ignore[arg-type]
|
|
continue
|
|
conv3x3s = []
|
|
for j in range(i + 1 - num_branches_pre):
|
|
in_channels = num_channels_pre_layer[-1]
|
|
out_channels = (num_channels_cur_layer[i] if j == i - num_branches_pre
|
|
else in_channels)
|
|
conv3x3s.append(nn.Sequential(
|
|
nn.Conv2d(in_channels, out_channels, 3, stride=2, padding=1, bias=False),
|
|
nn.BatchNorm2d(out_channels, momentum=0.01),
|
|
nn.ReLU(inplace=True)))
|
|
transition_layers.append(nn.Sequential(*conv3x3s))
|
|
|
|
return nn.ModuleList(transition_layers)
|
|
|
|
def _make_layer(self,
|
|
block: type[BasicBlockAttention | BasicBlock | Bottleneck],
|
|
in_channels: int,
|
|
out_channels: int,
|
|
blocks: int,
|
|
stride: int = 1) -> nn.Sequential:
|
|
"""Make an HRNet layer
|
|
|
|
Parameters
|
|
----------
|
|
block
|
|
The type of block to use for the layer
|
|
in_channels
|
|
The number of input channels
|
|
out_channels
|
|
The number of output channels
|
|
blocks
|
|
The number of blocks
|
|
stride
|
|
The stride size. Default: 1
|
|
|
|
Returns
|
|
-------
|
|
The sequential layer
|
|
"""
|
|
downsample = None
|
|
if stride != 1 or in_channels != out_channels * block.expansion:
|
|
downsample = nn.Sequential(nn.Conv2d(in_channels,
|
|
out_channels * block.expansion,
|
|
1,
|
|
stride=stride,
|
|
bias=False),
|
|
nn.BatchNorm2d(out_channels * block.expansion,
|
|
momentum=0.01))
|
|
|
|
layers = []
|
|
layers.append(block(in_channels, out_channels, stride, downsample))
|
|
in_channels = out_channels * block.expansion
|
|
for _ in range(1, blocks):
|
|
layers.append(block(in_channels, out_channels))
|
|
|
|
return nn.Sequential(*layers)
|
|
|
|
def _make_stage(self,
|
|
block: type[BasicBlockAttention | BasicBlock | Bottleneck],
|
|
layer_config: HRNetStageConfig,
|
|
num_in_channels: list[int],
|
|
multi_scale_output: bool = True) -> tuple[nn.Sequential, list[int]]:
|
|
"""Make a stage for HRNet
|
|
|
|
Parameters
|
|
----------
|
|
block
|
|
The type of block to use for the layer
|
|
layer_config
|
|
The configuration for the stage
|
|
num_in_channels
|
|
The input channels for the stage
|
|
multi_scale_output
|
|
``True`` to output multi scale
|
|
|
|
Returns
|
|
-------
|
|
sequential
|
|
The stage Sequential Modules
|
|
num_in_channels
|
|
The number of input channels from the final module
|
|
"""
|
|
num_modules = layer_config.num_modules
|
|
num_branches = layer_config.num_branches
|
|
num_blocks = layer_config.num_blocks
|
|
num_channels = layer_config.num_channels
|
|
|
|
modules: list[HighResolutionModule] = []
|
|
for i in range(num_modules):
|
|
if not multi_scale_output and i == num_modules - 1:
|
|
reset_multi_scale_output = False
|
|
else:
|
|
reset_multi_scale_output = True
|
|
modules.append(HighResolutionModule(num_branches,
|
|
block,
|
|
num_blocks,
|
|
num_in_channels,
|
|
num_channels,
|
|
reset_multi_scale_output))
|
|
num_in_channels = modules[-1].get_num_in_channels()
|
|
|
|
return nn.Sequential(*modules), num_in_channels
|
|
|
|
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
|
|
"""Forward pass through HRNet
|
|
|
|
Parameters
|
|
----------
|
|
inputs
|
|
Input to HRNet
|
|
|
|
Returns
|
|
-------
|
|
Output from HRNet
|
|
"""
|
|
x = self.conv1(inputs)
|
|
x = self.bn1(x)
|
|
x = self.relu(x)
|
|
x = self.conv2(x)
|
|
x = self.bn2(x)
|
|
x = self.relu(x)
|
|
x = self.layer1(x)
|
|
|
|
x_list = [self.transition1[i](x) if self.transition1[i] is not None else x
|
|
for i in range(self.stage_2_config.num_branches)]
|
|
y_list = self.stage2(x_list)
|
|
|
|
x_list = [self.transition2[i](y_list[-1]) if self.transition2[i] is not None
|
|
else y_list[i]
|
|
for i in range(self.stage_3_config.num_branches)]
|
|
y_list = self.stage3(x_list)
|
|
|
|
x_list = [self.transition3[i](y_list[-1]) if self.transition3[i] is not None
|
|
else y_list[i]
|
|
for i in range(self.stage_4_config.num_branches)]
|
|
x = self.stage4(x_list)
|
|
|
|
# Head Part
|
|
height, width = x[0].size(2), x[0].size(3)
|
|
x1 = F.interpolate(x[1], size=(height, width), mode="bilinear", align_corners=False)
|
|
x2 = F.interpolate(x[2], size=(height, width), mode="bilinear", align_corners=False)
|
|
x3 = F.interpolate(x[3], size=(height, width), mode="bilinear", align_corners=False)
|
|
x = torch.cat([x[0], x1, x2, x3], 1)
|
|
x = self.head(x)
|
|
|
|
return x
|
|
|
|
|
|
__all__ = get_module_objects(__name__)
|