From 8da677dbfea987e92ffc2eabf7136e5c4b277502 Mon Sep 17 00:00:00 2001 From: Sterling Taylor <166402033+staylorTT@users.noreply.github.com> Date: Wed, 4 Sep 2024 13:33:28 -0500 Subject: [PATCH] Delete forge/test/tvm/cnn/pytorch/alphapose directory cleanup --- .../alphapose/256x192_res50_lr1e-3_1x.yaml | 63 -- .../cnn/pytorch/alphapose/models/__init__.py | 14 - .../cnn/pytorch/alphapose/models/builder.py | 45 -- .../cnn/pytorch/alphapose/models/fastpose.py | 76 -- .../pytorch/alphapose/models/fastpose_duc.py | 110 --- .../alphapose/models/fastpose_duc_dense.py | 174 ----- .../pytorch/alphapose/models/layers/DUC.py | 30 - .../alphapose/models/layers/PixelUnshuffle.py | 33 - .../pytorch/alphapose/models/layers/Resnet.py | 251 ------- .../alphapose/models/layers/SE_Resnet.py | 275 ------- .../alphapose/models/layers/SE_module.py | 27 - .../alphapose/models/layers/ShuffleResnet.py | 276 ------- .../pytorch/alphapose/models/simplepose.py | 118 --- .../cnn/pytorch/alphapose/utils/__init__.py | 6 - .../tvm/cnn/pytorch/alphapose/utils/config.py | 12 - .../cnn/pytorch/alphapose/utils/registry.py | 79 -- .../cnn/pytorch/alphapose/utils/transforms.py | 698 ------------------ 17 files changed, 2287 deletions(-) delete mode 100644 forge/test/tvm/cnn/pytorch/alphapose/256x192_res50_lr1e-3_1x.yaml delete mode 100644 forge/test/tvm/cnn/pytorch/alphapose/models/__init__.py delete mode 100644 forge/test/tvm/cnn/pytorch/alphapose/models/builder.py delete mode 100644 forge/test/tvm/cnn/pytorch/alphapose/models/fastpose.py delete mode 100644 forge/test/tvm/cnn/pytorch/alphapose/models/fastpose_duc.py delete mode 100644 forge/test/tvm/cnn/pytorch/alphapose/models/fastpose_duc_dense.py delete mode 100644 forge/test/tvm/cnn/pytorch/alphapose/models/layers/DUC.py delete mode 100644 forge/test/tvm/cnn/pytorch/alphapose/models/layers/PixelUnshuffle.py delete mode 100644 forge/test/tvm/cnn/pytorch/alphapose/models/layers/Resnet.py delete mode 100644 forge/test/tvm/cnn/pytorch/alphapose/models/layers/SE_Resnet.py delete mode 100644 forge/test/tvm/cnn/pytorch/alphapose/models/layers/SE_module.py delete mode 100644 forge/test/tvm/cnn/pytorch/alphapose/models/layers/ShuffleResnet.py delete mode 100644 forge/test/tvm/cnn/pytorch/alphapose/models/simplepose.py delete mode 100644 forge/test/tvm/cnn/pytorch/alphapose/utils/__init__.py delete mode 100644 forge/test/tvm/cnn/pytorch/alphapose/utils/config.py delete mode 100644 forge/test/tvm/cnn/pytorch/alphapose/utils/registry.py delete mode 100644 forge/test/tvm/cnn/pytorch/alphapose/utils/transforms.py diff --git a/forge/test/tvm/cnn/pytorch/alphapose/256x192_res50_lr1e-3_1x.yaml b/forge/test/tvm/cnn/pytorch/alphapose/256x192_res50_lr1e-3_1x.yaml deleted file mode 100644 index dae407f4..00000000 --- a/forge/test/tvm/cnn/pytorch/alphapose/256x192_res50_lr1e-3_1x.yaml +++ /dev/null @@ -1,63 +0,0 @@ -DATASET: - TEST: - ANN: "annotations/person_keypoints_val2017.json" - DET_FILE: "./exp/json/test_det_yolo.json" - IMG_PREFIX: "val2017" - ROOT: "./data/coco/" - TYPE: "Mscoco_det" - TRAIN: - ANN: "annotations/person_keypoints_train2017.json" - AUG: - FLIP: true - NUM_JOINTS_HALF_BODY: 8 - PROB_HALF_BODY: -1 - ROT_FACTOR: 40 - SCALE_FACTOR: 0.3 - IMG_PREFIX: "train2017" - ROOT: "./data/coco/" - TYPE: "Mscoco" - VAL: - ANN: "annotations/person_keypoints_val2017.json" - IMG_PREFIX: "val2017" - ROOT: "./data/coco/" - TYPE: "Mscoco" -DATA_PRESET: - HEATMAP_SIZE: - - 64 - - 48 - IMAGE_SIZE: - - 256 - - 192 - NUM_JOINTS: 17 - SIGMA: 2 - TYPE: "simple" -DETECTOR: - CONFIDENCE: 0.05 - CONFIG: "detector/yolo/cfg/yolov3-spp.cfg" - NAME: "yolo" - NMS_THRES: 0.6 - WEIGHTS: "detector/yolo/data/yolov3-spp.weights" -MODEL: - NUM_DECONV_FILTERS: - - 256 - - 256 - - 256 - NUM_LAYERS: 50 - PRETRAINED: "" - TRY_LOAD: "" - TYPE: "FastPose" -TRAIN: - BATCH_SIZE: 32 - BEGIN_EPOCH: 0 - DPG_MILESTONE: 140 - DPG_STEP: - - 160 - - 190 - END_EPOCH: 200 - LR: 0.001 - LR_FACTOR: 0.1 - LR_STEP: - - 90 - - 120 - OPTIMIZER: "adam" - WORLD_SIZE: 4 diff --git a/forge/test/tvm/cnn/pytorch/alphapose/models/__init__.py b/forge/test/tvm/cnn/pytorch/alphapose/models/__init__.py deleted file mode 100644 index 75d7131e..00000000 --- a/forge/test/tvm/cnn/pytorch/alphapose/models/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC - -# SPDX-License-Identifier: Apache-2.0 -from .fastpose import FastPose -from .fastpose_duc import FastPose_DUC -from .simplepose import SimplePose -from .fastpose_duc_dense import FastPose_DUC_Dense - -__all__ = [ - "FastPose", - "SimplePose", - "FastPose_DUC", - "FastPose_DUC_Dense", -] diff --git a/forge/test/tvm/cnn/pytorch/alphapose/models/builder.py b/forge/test/tvm/cnn/pytorch/alphapose/models/builder.py deleted file mode 100644 index 65035d6b..00000000 --- a/forge/test/tvm/cnn/pytorch/alphapose/models/builder.py +++ /dev/null @@ -1,45 +0,0 @@ -# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC - -# SPDX-License-Identifier: Apache-2.0 -from torch import nn -import sys -import os - -sys.path.append(os.path.abspath("..")) -from ..utils import Registry, build_from_cfg - - -SPPE = Registry("sppe") -LOSS = Registry("loss") -DATASET = Registry("dataset") - - -def build(cfg, registry, default_args=None): - if isinstance(cfg, list): - modules = [build_from_cfg(cfg_, registry, default_args) for cfg_ in cfg] - return nn.Sequential(*modules) - else: - return build_from_cfg(cfg, registry, default_args) - - -def build_sppe(cfg, preset_cfg, **kwargs): - default_args = { - "PRESET": preset_cfg, - } - for key, value in kwargs.items(): - default_args[key] = value - return build(cfg, SPPE, default_args=default_args) - - -def build_loss(cfg): - return build(cfg, LOSS) - - -def build_dataset(cfg, preset_cfg, **kwargs): - exec(f"from ..datasets import {cfg.TYPE}") - default_args = { - "PRESET": preset_cfg, - } - for key, value in kwargs.items(): - default_args[key] = value - return build(cfg, DATASET, default_args=default_args) diff --git a/forge/test/tvm/cnn/pytorch/alphapose/models/fastpose.py b/forge/test/tvm/cnn/pytorch/alphapose/models/fastpose.py deleted file mode 100644 index 11b2f9e5..00000000 --- a/forge/test/tvm/cnn/pytorch/alphapose/models/fastpose.py +++ /dev/null @@ -1,76 +0,0 @@ -# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC - -# SPDX-License-Identifier: Apache-2.0 -# ----------------------------------------------------- -# Copyright (c) Shanghai Jiao Tong University. All rights reserved. -# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com) -# ----------------------------------------------------- - -import torch.nn as nn -import sys -import os -from .builder import SPPE -from .layers.DUC import DUC -from .layers.SE_Resnet import SEResnet - - -@SPPE.register_module -class FastPose(nn.Module): - conv_dim = 128 - - def __init__(self, norm_layer=nn.BatchNorm2d, **cfg): - super(FastPose, self).__init__() - self._preset_cfg = cfg["PRESET"] - - if "DCN" in cfg.keys(): - stage_with_dcn = cfg["STAGE_WITH_DCN"] - dcn = cfg["DCN"] - self.preact = SEResnet( - f"resnet{cfg['NUM_LAYERS']}", dcn=dcn, stage_with_dcn=stage_with_dcn - ) - else: - self.preact = SEResnet(f"resnet{cfg['NUM_LAYERS']}") - - import torchvision.models as tm - assert cfg["NUM_LAYERS"] in [18, 34, 50, 101, 152] - x = eval(f"tm.resnet{cfg['NUM_LAYERS']}(pretrained=True)") - - model_state = self.preact.state_dict() - state = { - k: v - for k, v in x.state_dict().items() - if k in self.preact.state_dict() - and v.size() == self.preact.state_dict()[k].size() - } - model_state.update(state) - self.preact.load_state_dict(model_state) - - self.suffle1 = nn.PixelShuffle(2) - self.duc1 = DUC(512, 1024, upscale_factor=2, norm_layer=norm_layer) - self.duc2 = DUC(256, 512, upscale_factor=2, norm_layer=norm_layer) - - self.conv_out = nn.Conv2d( - self.conv_dim, - self._preset_cfg["NUM_JOINTS"], - kernel_size=3, - stride=1, - padding=1, - ) - - def forward(self, x): - out = self.preact(x) - out = self.suffle1(out) - out = self.duc1(out) - out = self.duc2(out) - - out = self.conv_out(out) - return out - - def _initialize(self): - for m in self.conv_out.modules(): - if isinstance(m, nn.Conv2d): - # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - # logger.info('=> init {}.weight as normal(0, 0.001)'.format(name)) - # logger.info('=> init {}.bias as 0'.format(name)) - nn.init.normal_(m.weight, std=0.001) - nn.init.constant_(m.bias, 0) diff --git a/forge/test/tvm/cnn/pytorch/alphapose/models/fastpose_duc.py b/forge/test/tvm/cnn/pytorch/alphapose/models/fastpose_duc.py deleted file mode 100644 index 64e4f8b7..00000000 --- a/forge/test/tvm/cnn/pytorch/alphapose/models/fastpose_duc.py +++ /dev/null @@ -1,110 +0,0 @@ -# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC - -# SPDX-License-Identifier: Apache-2.0 -# ----------------------------------------------------- -# Copyright (c) Shanghai Jiao Tong University. All rights reserved. -# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com) -# ----------------------------------------------------- - -import torch.nn as nn - -from .builder import SPPE -from .layers.Resnet import ResNet -from .layers.SE_Resnet import SEResnet -from .layers.ShuffleResnet import ShuffleResnet - - -@SPPE.register_module -class FastPose_DUC(nn.Module): - conv_dim = 256 - - def __init__(self, norm_layer=nn.BatchNorm2d, **cfg): - super(FastPose_DUC, self).__init__() - self._preset_cfg = cfg["PRESET"] - if cfg["BACKBONE"] == "shuffle": - print("Load shuffle backbone...") - backbone = ShuffleResnet - elif cfg["BACKBONE"] == "se-resnet": - print("Load SE Resnet...") - backbone = SEResnet - else: - print("Load Resnet...") - backbone = ResNet - - if "DCN" in cfg.keys(): - stage_with_dcn = cfg["STAGE_WITH_DCN"] - dcn = cfg["DCN"] - self.preact = backbone( - f"resnet{cfg['NUM_LAYERS']}", dcn=dcn, stage_with_dcn=stage_with_dcn - ) - else: - self.preact = backbone(f"resnet{cfg['NUM_LAYERS']}") - - # Imagenet pretrain model - sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../../") - import torchvision as tm # noqa: F401,F403 - - assert cfg["NUM_LAYERS"] in [18, 34, 50, 101, 152] - x = eval(f"tm.resnet{cfg['NUM_LAYERS']}(pretrained=True)") - - model_state = self.preact.state_dict() - state = { - k: v - for k, v in x.state_dict().items() - if k in self.preact.state_dict() - and v.size() == self.preact.state_dict()[k].size() - } - model_state.update(state) - self.preact.load_state_dict(model_state) - self.norm_layer = norm_layer - - stage1_cfg = cfg["STAGE1"] - stage2_cfg = cfg["STAGE2"] - stage3_cfg = cfg["STAGE3"] - - self.duc1 = self._make_duc_stage(stage1_cfg, 2048, 1024) - self.duc2 = self._make_duc_stage(stage2_cfg, 1024, 512) - self.duc3 = self._make_duc_stage(stage3_cfg, 512, self.conv_dim) - - self.conv_out = nn.Conv2d( - self.conv_dim, - self._preset_cfg["NUM_JOINTS"], - kernel_size=3, - stride=1, - padding=1, - ) - - def forward(self, x): - out = self.preact(x) - out = self.duc1(out) - out = self.duc2(out) - out = self.duc3(out) - - out = self.conv_out(out) - return out - - def _make_duc_stage(self, layer_config, inplanes, outplanes): - layers = [] - - shuffle = nn.PixelShuffle(2) - inplanes //= 4 - layers.append(shuffle) - for i in range(layer_config.NUM_CONV - 1): - conv = nn.Conv2d(inplanes, inplanes, kernel_size=3, padding=1, bias=False) - norm_layer = self.norm_layer(inplanes, momentum=0.1) - relu = nn.ReLU(inplace=True) - layers += [conv, norm_layer, relu] - conv = nn.Conv2d(inplanes, outplanes, kernel_size=3, padding=1, bias=False) - norm_layer = self.norm_layer(outplanes, momentum=0.1) - relu = nn.ReLU(inplace=True) - layers += [conv, norm_layer, relu] - return nn.Sequential(*layers) - - def _initialize(self): - for m in self.conv_out.modules(): - if isinstance(m, nn.Conv2d): - # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - # logger.info('=> init {}.weight as normal(0, 0.001)'.format(name)) - # logger.info('=> init {}.bias as 0'.format(name)) - nn.init.normal_(m.weight, std=0.001) - nn.init.constant_(m.bias, 0) diff --git a/forge/test/tvm/cnn/pytorch/alphapose/models/fastpose_duc_dense.py b/forge/test/tvm/cnn/pytorch/alphapose/models/fastpose_duc_dense.py deleted file mode 100644 index d7358020..00000000 --- a/forge/test/tvm/cnn/pytorch/alphapose/models/fastpose_duc_dense.py +++ /dev/null @@ -1,174 +0,0 @@ -# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC - -# SPDX-License-Identifier: Apache-2.0 -# ----------------------------------------------------- -# Copyright (c) Shanghai Jiao Tong University. All rights reserved. -# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com) -# ----------------------------------------------------- - -import torch -import torch.nn as nn -from .builder import SPPE -from .layers.Resnet import ResNet -from .layers.SE_Resnet import SEResnet -from .layers.ShuffleResnet import ShuffleResnet - - -@SPPE.register_module -class FastPose_DUC_Dense(nn.Module): - conv_dim = 256 - - def __init__(self, norm_layer=nn.BatchNorm2d, **cfg): - super(FastPose_DUC_Dense, self).__init__() - self._preset_cfg = cfg["PRESET"] - if cfg["BACKBONE"] == "shuffle": - print("Load shuffle backbone...") - backbone = ShuffleResnet - elif cfg["BACKBONE"] == "se-resnet": - print("Load SE Resnet...") - backbone = SEResnet - else: - print("Load Resnet...") - backbone = ResNet - - if "DCN" in cfg.keys(): - stage_with_dcn = cfg["STAGE_WITH_DCN"] - dcn = cfg["DCN"] - self.preact = backbone( - f"resnet{cfg['NUM_LAYERS']}", dcn=dcn, stage_with_dcn=stage_with_dcn - ) - else: - self.preact = backbone(f"resnet{cfg['NUM_LAYERS']}") - - # Init Backbone - for m in self.preact.modules(): - if isinstance(m, nn.Conv2d): - nn.init.normal_(m.weight, std=0.001) - for name, _ in m.named_parameters(): - if name in ["bias"]: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.BatchNorm2d): - # nn.init.constant_(m.weight, 1) - nn.init.uniform_(m.weight, 0, 1) - nn.init.constant_(m.bias, 0) - - # Imagenet pretrain model - sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../../") - import torchvision as tm # noqa: F401,F403 - - if cfg["NUM_LAYERS"] == 152: - """ Load pretrained model """ - x = tm.resnet152(pretrained=True) - elif cfg["NUM_LAYERS"] == 101: - """ Load pretrained model """ - x = tm.resnet101(pretrained=True) - elif cfg["NUM_LAYERS"] == 50: - x = tm.resnet50(pretrained=True) - elif cfg["NUM_LAYERS"] == 18: - x = tm.resnet18(pretrained=True) - else: - raise NotImplementedError - model_state = self.preact.state_dict() - state = { - k: v - for k, v in x.state_dict().items() - if k in self.preact.state_dict() - and v.size() == self.preact.state_dict()[k].size() - } - model_state.update(state) - self.preact.load_state_dict(model_state) - self.norm_layer = norm_layer - - stage1_cfg = cfg["STAGE1"] - stage2_cfg = cfg["STAGE2"] - stage3_cfg = cfg["STAGE3"] - - duc1 = self._make_duc_stage(stage1_cfg, 2048, 1024) - duc2 = self._make_duc_stage(stage2_cfg, 1024, 512) - duc3 = self._make_duc_stage(stage3_cfg, 512, self.conv_dim) - - self.duc = nn.Sequential(duc1, duc2, duc3) - - duc1_dense = self._make_duc_stage(stage1_cfg, 2048, 1024) - duc2_dense = self._make_duc_stage(stage2_cfg, 1024, 512) - duc3_dense = self._make_duc_stage(stage3_cfg, 512, self.conv_dim) - - self.duc_dense = nn.Sequential(duc1_dense, duc2_dense, duc3_dense) - - self.conv_out = nn.Conv2d( - self.conv_dim, - self._preset_cfg["NUM_JOINTS"], - kernel_size=3, - stride=1, - padding=1, - ) - - self.conv_out_dense = nn.Conv2d( - self.conv_dim, - (self._preset_cfg["NUM_JOINTS_DENSE"] - self._preset_cfg["NUM_JOINTS"]), - kernel_size=3, - stride=1, - padding=1, - ) - for params in self.preact.parameters(): - params.requires_grad = False - for params in self.duc.parameters(): - params.requires_grad = False - - def forward(self, x): - bk_out = self.preact(x) - out = self.duc(bk_out) - out_dense = self.duc_dense(bk_out) - out = self.conv_out(out) - out_dense = self.conv_out_dense(out_dense) - out = torch.cat((out, out_dense), 1) - return out - - def _make_duc_stage(self, layer_config, inplanes, outplanes): - layers = [] - - shuffle = nn.PixelShuffle(2) - inplanes //= 4 - layers.append(shuffle) - for i in range(layer_config.NUM_CONV - 1): - conv = nn.Conv2d(inplanes, inplanes, kernel_size=3, padding=1, bias=False) - norm_layer = self.norm_layer(inplanes, momentum=0.1) - relu = nn.ReLU(inplace=True) - layers += [conv, norm_layer, relu] - conv = nn.Conv2d(inplanes, outplanes, kernel_size=3, padding=1, bias=False) - norm_layer = self.norm_layer(outplanes, momentum=0.1) - relu = nn.ReLU(inplace=True) - layers += [conv, norm_layer, relu] - return nn.Sequential(*layers) - - def _initialize(self): - for m in self.duc.modules(): - if isinstance(m, nn.Conv2d): - nn.init.normal_(m.weight, std=0.001) - for name, _ in m.named_parameters(): - if name in ["bias"]: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.BatchNorm2d): - # nn.init.constant_(m.weight, 1) - nn.init.uniform_(m.weight, 0, 1) - nn.init.constant_(m.bias, 0) - for m in self.conv_out.modules(): - if isinstance(m, nn.Conv2d): - nn.init.normal_(m.weight, std=0.001) - nn.init.constant_(m.bias, 0) - - # init dense-branch - for m in self.duc_dense.modules(): - if isinstance(m, nn.Conv2d): - nn.init.normal_(m.weight, std=0.001) - for name, _ in m.named_parameters(): - if name in ["bias"]: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.BatchNorm2d): - # nn.init.constant_(m.weight, 1) - nn.init.uniform_(m.weight, 0, 1) - nn.init.constant_(m.bias, 0) - for m in self.conv_out_dense.modules(): - if isinstance(m, nn.Conv2d): - nn.init.normal_(m.weight, std=0.001) - nn.init.constant_(m.bias, 0) diff --git a/forge/test/tvm/cnn/pytorch/alphapose/models/layers/DUC.py b/forge/test/tvm/cnn/pytorch/alphapose/models/layers/DUC.py deleted file mode 100644 index 0de18ff0..00000000 --- a/forge/test/tvm/cnn/pytorch/alphapose/models/layers/DUC.py +++ /dev/null @@ -1,30 +0,0 @@ -# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC - -# SPDX-License-Identifier: Apache-2.0 -# ----------------------------------------------------- -# Copyright (c) Shanghai Jiao Tong University. All rights reserved. -# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com) -# ----------------------------------------------------- - -import torch.nn as nn - - -class DUC(nn.Module): - """ - Initialize: inplanes, planes, upscale_factor - OUTPUT: (planes // upscale_factor^2) * ht * wd - """ - - def __init__(self, inplanes, planes, upscale_factor=2, norm_layer=nn.BatchNorm2d): - super(DUC, self).__init__() - self.conv = nn.Conv2d(inplanes, planes, kernel_size=3, padding=1, bias=False) - self.bn = norm_layer(planes, momentum=0.1) - self.relu = nn.ReLU(inplace=True) - self.pixel_shuffle = nn.PixelShuffle(upscale_factor) - - def forward(self, x): - x = self.conv(x) - x = self.bn(x) - x = self.relu(x) - x = self.pixel_shuffle(x) - return x diff --git a/forge/test/tvm/cnn/pytorch/alphapose/models/layers/PixelUnshuffle.py b/forge/test/tvm/cnn/pytorch/alphapose/models/layers/PixelUnshuffle.py deleted file mode 100644 index 346470ee..00000000 --- a/forge/test/tvm/cnn/pytorch/alphapose/models/layers/PixelUnshuffle.py +++ /dev/null @@ -1,33 +0,0 @@ -# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC - -# SPDX-License-Identifier: Apache-2.0 -# ----------------------------------------------------- -# Copyright (c) Shanghai Jiao Tong University. All rights reserved. -# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com) -# ----------------------------------------------------- - -import torch.nn as nn - - -class PixelUnshuffle(nn.Module): - """ - Initialize: inplanes, planes, upscale_factor - OUTPUT: (planes // upscale_factor^2) * ht * wd - """ - - def __init__(self, downscale_factor=2): - super(PixelUnshuffle, self).__init__() - self._r = downscale_factor - - def forward(self, x): - b, c, h, w = x.shape - out_c = c * (self._r * self._r) - out_h = h // self._r - out_w = w // self._r - - x_view = x.contiguous().view(b, c, out_h, self._r, out_w, self._r) - x_prime = ( - x_view.permute(0, 1, 3, 5, 2, 4).contiguous().view(b, out_c, out_h, out_w) - ) - - return x_prime diff --git a/forge/test/tvm/cnn/pytorch/alphapose/models/layers/Resnet.py b/forge/test/tvm/cnn/pytorch/alphapose/models/layers/Resnet.py deleted file mode 100644 index 61d6596b..00000000 --- a/forge/test/tvm/cnn/pytorch/alphapose/models/layers/Resnet.py +++ /dev/null @@ -1,251 +0,0 @@ -# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC - -# SPDX-License-Identifier: Apache-2.0 -# ----------------------------------------------------- -# Copyright (c) Shanghai Jiao Tong University. All rights reserved. -# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com) -# ----------------------------------------------------- - -import torch.nn as nn -import torch.nn.functional as F - -# from .dcn import DeformConv, ModulatedDeformConv - - -def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): - """3x3 convolution with padding""" - return nn.Conv2d( - in_planes, - out_planes, - kernel_size=3, - stride=stride, - padding=dilation, - groups=groups, - bias=False, - dilation=dilation, - ) - - -class BasicBlock(nn.Module): - expansion = 1 - - def __init__( - self, - inplanes, - planes, - stride=1, - downsample=None, - groups=1, - base_width=64, - dilation=1, - norm_layer=None, - ): - super(BasicBlock, self).__init__() - if norm_layer is None: - norm_layer = nn.BatchNorm2d - if groups != 1 or base_width != 64: - raise ValueError("BasicBlock only supports groups=1 and base_width=64") - if dilation > 1: - raise NotImplementedError("Dilation > 1 not supported in BasicBlock") - # Both self.conv1 and self.downsample layers downsample the input when stride != 1 - self.conv1 = conv3x3(inplanes, planes, stride) - self.bn1 = norm_layer(planes) - self.relu = nn.ReLU(inplace=True) - self.conv2 = conv3x3(planes, planes) - self.bn2 = norm_layer(planes) - self.downsample = downsample - self.stride = stride - - def forward(self, x): - identity = x - - out = self.conv1(x) - out = self.bn1(out) - out = self.relu(out) - - out = self.conv2(out) - out = self.bn2(out) - - if self.downsample is not None: - identity = self.downsample(x) - - out += identity - out = self.relu(out) - - return out - - -class Bottleneck(nn.Module): - expansion = 4 - - def __init__( - self, - inplanes, - planes, - stride=1, - downsample=None, - norm_layer=nn.BatchNorm2d, - dcn=None, - ): - super(Bottleneck, self).__init__() - self.dcn = dcn - self.with_dcn = dcn is not None - - self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) - self.bn1 = norm_layer(planes, momentum=0.1) - if self.with_dcn: - fallback_on_stride = dcn.get("FALLBACK_ON_STRIDE", False) - self.with_modulated_dcn = dcn.get("MODULATED", False) - if not self.with_dcn or fallback_on_stride: - self.conv2 = nn.Conv2d( - planes, planes, kernel_size=3, stride=stride, padding=1, bias=False - ) - # else: - # self.deformable_groups = dcn.get("DEFORM_GROUP", 1) - # if not self.with_modulated_dcn: - # conv_op = DeformConv - # offset_channels = 18 - # else: - # conv_op = ModulatedDeformConv - # offset_channels = 27 - - # self.conv2_offset = nn.Conv2d( - # planes, - # self.deformable_groups * offset_channels, - # kernel_size=3, - # stride=stride, - # padding=1, - # ) - # self.conv2 = conv_op( - # planes, - # planes, - # kernel_size=3, - # stride=stride, - # padding=1, - # deformable_groups=self.deformable_groups, - # bias=False, - # ) - - self.bn2 = norm_layer(planes, momentum=0.1) - self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) - self.bn3 = norm_layer(planes * 4, momentum=0.1) - self.downsample = downsample - self.stride = stride - - def forward(self, x): - residual = x - - out = F.relu(self.bn1(self.conv1(x)), inplace=True) - if not self.with_dcn: - out = F.relu(self.bn2(self.conv2(out)), inplace=True) - # elif self.with_modulated_dcn: - # offset_mask = self.conv2_offset(out) - # offset = offset_mask[:, : 18 * self.deformable_groups, :, :] - # mask = offset_mask[:, -9 * self.deformable_groups :, :, :] - # mask = mask.sigmoid() - # out = F.relu(self.bn2(self.conv2(out, offset, mask))) - # else: - # offset = self.conv2_offset(out) - # out = F.relu(self.bn2(self.conv2(out, offset)), inplace=True) - - out = self.conv3(out) - out = self.bn3(out) - - if self.downsample is not None: - residual = self.downsample(x) - - out += residual - out = F.relu(out) - - return out - - -class ResNet(nn.Module): - """ ResNet """ - - def __init__( - self, - architecture, - norm_layer=nn.BatchNorm2d, - dcn=None, - stage_with_dcn=(False, False, False, False), - ): - super(ResNet, self).__init__() - self._norm_layer = norm_layer - assert architecture in ["resnet18", "resnet50", "resnet101", "resnet152"] - layers = { - "resnet18": [2, 2, 2, 2], - "resnet34": [3, 4, 6, 3], - "resnet50": [3, 4, 6, 3], - "resnet101": [3, 4, 23, 3], - "resnet152": [3, 8, 36, 3], - } - self.inplanes = 64 - if architecture == "resnet18" or architecture == "resnet34": - self.block = BasicBlock - else: - self.block = Bottleneck - self.layers = layers[architecture] - - self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False) - self.bn1 = norm_layer(64, eps=1e-5, momentum=0.1, affine=True) - self.relu = nn.ReLU(inplace=True) - self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) - - stage_dcn = [dcn if with_dcn else None for with_dcn in stage_with_dcn] - - self.layer1 = self.make_layer(self.block, 64, self.layers[0], dcn=stage_dcn[0]) - self.layer2 = self.make_layer( - self.block, 128, self.layers[1], stride=2, dcn=stage_dcn[1] - ) - self.layer3 = self.make_layer( - self.block, 256, self.layers[2], stride=2, dcn=stage_dcn[2] - ) - - self.layer4 = self.make_layer( - self.block, 512, self.layers[3], stride=2, dcn=stage_dcn[3] - ) - - def forward(self, x): - x = self.maxpool(self.relu(self.bn1(self.conv1(x)))) # 64 * h/4 * w/4 - x = self.layer1(x) # 256 * h/4 * w/4 - x = self.layer2(x) # 512 * h/8 * w/8 - x = self.layer3(x) # 1024 * h/16 * w/16 - x = self.layer4(x) # 2048 * h/32 * w/32 - return x - - def stages(self): - return [self.layer1, self.layer2, self.layer3, self.layer4] - - def make_layer(self, block, planes, blocks, stride=1, dcn=None): - downsample = None - if stride != 1 or self.inplanes != planes * block.expansion: - downsample = nn.Sequential( - nn.Conv2d( - self.inplanes, - planes * block.expansion, - kernel_size=1, - stride=stride, - bias=False, - ), - self._norm_layer(planes * block.expansion), - ) - - layers = [] - layers.append( - block( - self.inplanes, - planes, - stride, - downsample, - norm_layer=self._norm_layer, - dcn=dcn, - ) - ) - self.inplanes = planes * block.expansion - for i in range(1, blocks): - layers.append( - block(self.inplanes, planes, norm_layer=self._norm_layer, dcn=dcn) - ) - - return nn.Sequential(*layers) diff --git a/forge/test/tvm/cnn/pytorch/alphapose/models/layers/SE_Resnet.py b/forge/test/tvm/cnn/pytorch/alphapose/models/layers/SE_Resnet.py deleted file mode 100644 index ac20cd03..00000000 --- a/forge/test/tvm/cnn/pytorch/alphapose/models/layers/SE_Resnet.py +++ /dev/null @@ -1,275 +0,0 @@ -# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC - -# SPDX-License-Identifier: Apache-2.0 -# ----------------------------------------------------- -# Copyright (c) Shanghai Jiao Tong University. All rights reserved. -# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com) -# ----------------------------------------------------- - -import torch.nn as nn -import torch.nn.functional as F - -# from .dcn import DeformConv, ModulatedDeformConv -from .SE_module import SELayer - - -def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): - """3x3 convolution with padding""" - return nn.Conv2d( - in_planes, - out_planes, - kernel_size=3, - stride=stride, - padding=dilation, - groups=groups, - bias=False, - dilation=dilation, - ) - - -class BasicBlock(nn.Module): - expansion = 1 - - def __init__( - self, - inplanes, - planes, - stride=1, - downsample=None, - reduction=False, - norm_layer=nn.BatchNorm2d, - ): - super(BasicBlock, self).__init__() - - # Both self.conv1 and self.downsample layers downsample the input when stride != 1 - self.conv1 = conv3x3(inplanes, planes, stride) - self.bn1 = norm_layer(planes) - self.relu = nn.ReLU(inplace=True) - self.conv2 = conv3x3(planes, planes) - self.bn2 = norm_layer(planes) - self.downsample = downsample - self.stride = stride - if reduction: - self.se = SELayer(planes) - self.reduc = reduction - - def forward(self, x): - identity = x - - out = self.conv1(x) - out = self.bn1(out) - out = self.relu(out) - - out = self.conv2(out) - out = self.bn2(out) - - if self.reduc: - out = self.se(out) - - if self.downsample is not None: - identity = self.downsample(x) - - out += identity - out = self.relu(out) - - return out - - -class Bottleneck(nn.Module): - expansion = 4 - - def __init__( - self, - inplanes, - planes, - stride=1, - downsample=None, - reduction=False, - norm_layer=nn.BatchNorm2d, - dcn=None, - ): - super(Bottleneck, self).__init__() - - self.dcn = dcn - self.with_dcn = dcn is not None - - self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) - self.bn1 = norm_layer(planes, momentum=0.1) - if self.with_dcn: - fallback_on_stride = dcn.get("FALLBACK_ON_STRIDE", False) - self.with_modulated_dcn = dcn.get("MODULATED", False) - - assert (not self.with_dcn) or (fallback_on_stride) - - if not self.with_dcn or fallback_on_stride: - self.conv2 = nn.Conv2d( - planes, planes, kernel_size=3, stride=stride, padding=1, bias=False - ) - # else: - # self.deformable_groups = dcn.get("DEFORM_GROUP", 1) - # if not self.with_modulated_dcn: - # conv_op = DeformConv - # offset_channels = 18 - # else: - # conv_op = ModulatedDeformConv - # offset_channels = 27 - - # self.conv2_offset = nn.Conv2d( - # planes, - # self.deformable_groups * offset_channels, - # kernel_size=3, - # stride=stride, - # padding=1, - # ) - # self.conv2 = conv_op( - # planes, - # planes, - # kernel_size=3, - # stride=stride, - # padding=1, - # deformable_groups=self.deformable_groups, - # bias=False, - # ) - - self.bn2 = norm_layer(planes, momentum=0.1) - self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) - self.bn3 = norm_layer(planes * 4, momentum=0.1) - if reduction: - self.se = SELayer(planes * 4) - - self.reduc = reduction - self.downsample = downsample - self.stride = stride - - def forward(self, x): - residual = x - - out = F.relu(self.bn1(self.conv1(x)), inplace=True) - if not self.with_dcn: - out = F.relu(self.bn2(self.conv2(out)), inplace=True) - # elif self.with_modulated_dcn: - # offset_mask = self.conv2_offset(out) - # offset = offset_mask[:, : 18 * self.deformable_groups, :, :] - # mask = offset_mask[:, -9 * self.deformable_groups :, :, :] - # mask = mask.sigmoid() - # out = F.relu(self.bn2(self.conv2(out, offset, mask))) - # else: - # offset = self.conv2_offset(out) - # out = F.relu(self.bn2(self.conv2(out, offset)), inplace=True) - - out = self.conv3(out) - out = self.bn3(out) - if self.reduc: - out = self.se(out) - - if self.downsample is not None: - residual = self.downsample(x) - - out += residual - out = F.relu(out) - - return out - - -class SEResnet(nn.Module): - """ SEResnet """ - - def __init__( - self, - architecture, - norm_layer=nn.BatchNorm2d, - dcn=None, - stage_with_dcn=(False, False, False, False), - ): - super(SEResnet, self).__init__() - self._norm_layer = norm_layer - assert architecture in ["resnet18", "resnet50", "resnet101", "resnet152"] - layers = { - "resnet18": [2, 2, 2, 2], - "resnet34": [3, 4, 6, 3], - "resnet50": [3, 4, 6, 3], - "resnet101": [3, 4, 23, 3], - "resnet152": [3, 8, 36, 3], - } - self.inplanes = 64 - if architecture == "resnet18" or architecture == "resnet34": - self.block = BasicBlock - else: - self.block = Bottleneck - self.layers = layers[architecture] - - self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False) - self.bn1 = norm_layer(64, eps=1e-5, momentum=0.1, affine=True) - self.relu = nn.ReLU(inplace=True) - self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) - - stage_dcn = [dcn if with_dcn else None for with_dcn in stage_with_dcn] - - self.layer1 = self.make_layer(self.block, 64, self.layers[0], dcn=stage_dcn[0]) - self.layer2 = self.make_layer( - self.block, 128, self.layers[1], stride=2, dcn=stage_dcn[1] - ) - self.layer3 = self.make_layer( - self.block, 256, self.layers[2], stride=2, dcn=stage_dcn[2] - ) - - self.layer4 = self.make_layer( - self.block, 512, self.layers[3], stride=2, dcn=stage_dcn[3] - ) - - def forward(self, x): - x = self.maxpool(self.relu(self.bn1(self.conv1(x)))) # 64 * h/4 * w/4 - x = self.layer1(x) # 256 * h/4 * w/4 - x = self.layer2(x) # 512 * h/8 * w/8 - x = self.layer3(x) # 1024 * h/16 * w/16 - x = self.layer4(x) # 2048 * h/32 * w/32 - return x - - def stages(self): - return [self.layer1, self.layer2, self.layer3, self.layer4] - - def make_layer(self, block, planes, blocks, stride=1, dcn=None): - downsample = None - if stride != 1 or self.inplanes != planes * block.expansion: - downsample = nn.Sequential( - nn.Conv2d( - self.inplanes, - planes * block.expansion, - kernel_size=1, - stride=stride, - bias=False, - ), - self._norm_layer(planes * block.expansion, momentum=0.1), - ) - - layers = [] - if downsample is not None: - layers.append( - block( - self.inplanes, - planes, - stride, - downsample, - reduction=True, - norm_layer=self._norm_layer, - dcn=dcn, - ) - ) - else: - layers.append( - block( - self.inplanes, - planes, - stride, - downsample, - norm_layer=self._norm_layer, - dcn=dcn, - ) - ) - self.inplanes = planes * block.expansion - for i in range(1, blocks): - layers.append( - block(self.inplanes, planes, norm_layer=self._norm_layer, dcn=dcn) - ) - - return nn.Sequential(*layers) diff --git a/forge/test/tvm/cnn/pytorch/alphapose/models/layers/SE_module.py b/forge/test/tvm/cnn/pytorch/alphapose/models/layers/SE_module.py deleted file mode 100644 index 0fd9e21e..00000000 --- a/forge/test/tvm/cnn/pytorch/alphapose/models/layers/SE_module.py +++ /dev/null @@ -1,27 +0,0 @@ -# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC - -# SPDX-License-Identifier: Apache-2.0 -# ----------------------------------------------------- -# Copyright (c) Shanghai Jiao Tong University. All rights reserved. -# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com) -# ----------------------------------------------------- - -from torch import nn - - -class SELayer(nn.Module): - def __init__(self, channel, reduction=1): - super(SELayer, self).__init__() - self.avg_pool = nn.AdaptiveAvgPool2d(1) - self.fc = nn.Sequential( - nn.Linear(channel, channel // reduction), - nn.ReLU(inplace=True), - nn.Linear(channel // reduction, channel), - nn.Sigmoid(), - ) - - def forward(self, x): - b, c, _, _ = x.size() - y = self.avg_pool(x).view(b, c) - y = self.fc(y).view(b, c, 1, 1) - return x * y diff --git a/forge/test/tvm/cnn/pytorch/alphapose/models/layers/ShuffleResnet.py b/forge/test/tvm/cnn/pytorch/alphapose/models/layers/ShuffleResnet.py deleted file mode 100644 index 6f6b9ff8..00000000 --- a/forge/test/tvm/cnn/pytorch/alphapose/models/layers/ShuffleResnet.py +++ /dev/null @@ -1,276 +0,0 @@ -# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC - -# SPDX-License-Identifier: Apache-2.0 -# ----------------------------------------------------- -# Copyright (c) Shanghai Jiao Tong University. All rights reserved. -# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com) -# ----------------------------------------------------- - -import torch.nn as nn -import torch.nn.functional as F - -# from .dcn import DCN -from .PixelUnshuffle import PixelUnshuffle -from .SE_module import SELayer - - -def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): - """3x3 convolution with padding""" - return nn.Conv2d( - in_planes, - out_planes, - kernel_size=3, - stride=stride, - padding=dilation, - groups=groups, - bias=False, - dilation=dilation, - ) - - -class BasicBlock(nn.Module): - expansion = 1 - - def __init__( - self, - inplanes, - planes, - stride=1, - downsample=None, - reduction=False, - norm_layer=nn.BatchNorm2d, - ): - super(BasicBlock, self).__init__() - - # Both self.conv1 and self.downsample layers downsample the input when stride != 1 - self.conv1 = conv3x3(inplanes, planes, stride) - self.bn1 = norm_layer(planes) - self.relu = nn.ReLU(inplace=True) - self.conv2 = conv3x3(planes, planes) - self.bn2 = norm_layer(planes) - self.downsample = downsample - self.stride = stride - if reduction: - self.se = SELayer(planes) - self.reduc = reduction - - def forward(self, x): - identity = x - - out = self.conv1(x) - out = self.bn1(out) - out = self.relu(out) - - out = self.conv2(out) - out = self.bn2(out) - - if self.reduc: - out = self.se(out) - - if self.downsample is not None: - identity = self.downsample(x) - - out += identity - out = self.relu(out) - - return out - - -class Bottleneck(nn.Module): - expansion = 4 - - def __init__( - self, - inplanes, - planes, - stride=1, - downsample=None, - reduction=False, - norm_layer=nn.BatchNorm2d, - dcn=None, - ): - super(Bottleneck, self).__init__() - self.dcn = dcn - self.with_dcn = dcn is not None - - self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) - self.bn1 = norm_layer(planes, momentum=0.1) - if self.with_dcn: - fallback_on_stride = dcn.get("FALLBACK_ON_STRIDE", False) - self.with_modulated_dcn = dcn.get("MODULATED", False) - - if stride > 1: - conv_layers = [] - conv_layers.append(PixelUnshuffle(stride)) - if not self.with_dcn or fallback_on_stride: - conv_layers.append( - nn.Conv2d( - planes * 4, - planes, - kernel_size=3, - stride=1, - padding=1, - bias=False, - ) - ) - # else: - # conv_layers.append( - # DCN( - # planes * 4, - # planes, - # dcn, - # kernel_size=3, - # stride=1, - # padding=1, - # bias=False, - # ) - # ) - self.conv2 = nn.Sequential(*conv_layers) - else: - if not self.with_dcn or fallback_on_stride: - self.conv2 = nn.Conv2d( - planes, planes, kernel_size=3, stride=stride, padding=1, bias=False - ) - # else: - # self.conv2 = DCN( - # planes, - # planes, - # dcn, - # kernel_size=3, - # stride=stride, - # padding=1, - # bias=False, - # ) - - self.bn2 = norm_layer(planes, momentum=0.1) - self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) - self.bn3 = norm_layer(planes * 4, momentum=0.1) - if reduction: - self.se = SELayer(planes * 4) - - self.reduc = reduction - self.downsample = downsample - self.stride = stride - - def forward(self, x): - residual = x - - out = F.relu(self.bn1(self.conv1(x)), inplace=True) - out = F.relu(self.bn2(self.conv2(out)), inplace=True) - - out = self.conv3(out) - out = self.bn3(out) - if self.reduc: - out = self.se(out) - - if self.downsample is not None: - residual = self.downsample(x) - - out += residual - out = F.relu(out, inplace=True) - - return out - - -class ShuffleResnet(nn.Module): - """ ShuffleResnet """ - - def __init__( - self, - architecture, - norm_layer=nn.BatchNorm2d, - dcn=None, - stage_with_dcn=(False, False, False, False), - ): - super(ShuffleResnet, self).__init__() - self._norm_layer = norm_layer - assert architecture in ["resnet18", "resnet50", "resnet101", "resnet152"] - layers = { - "resnet18": [2, 2, 2, 2], - "resnet34": [3, 4, 6, 3], - "resnet50": [3, 4, 6, 3], - "resnet101": [3, 4, 23, 3], - "resnet152": [3, 8, 36, 3], - } - self.inplanes = 64 - if architecture == "resnet18" or architecture == "resnet34": - self.block = BasicBlock - else: - self.block = Bottleneck - self.layers = layers[architecture] - - self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False) - self.bn1 = norm_layer(64, eps=1e-5, momentum=0.1, affine=True) - self.relu = nn.ReLU(inplace=True) - self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) - - stage_dcn = [dcn if with_dcn else None for with_dcn in stage_with_dcn] - - self.layer1 = self.make_layer(self.block, 64, self.layers[0], dcn=stage_dcn[0]) - self.layer2 = self.make_layer( - self.block, 128, self.layers[1], stride=2, dcn=stage_dcn[1] - ) - self.layer3 = self.make_layer( - self.block, 256, self.layers[2], stride=2, dcn=stage_dcn[2] - ) - - self.layer4 = self.make_layer( - self.block, 512, self.layers[3], stride=2, dcn=stage_dcn[3] - ) - - def forward(self, x): - x = self.maxpool(self.relu(self.bn1(self.conv1(x)))) # 64 * h/4 * w/4 - x = self.layer1(x) # 256 * h/4 * w/4 - x = self.layer2(x) # 512 * h/8 * w/8 - x = self.layer3(x) # 1024 * h/16 * w/16 - x = self.layer4(x) # 2048 * h/32 * w/32 - return x - - def stages(self): - return [self.layer1, self.layer2, self.layer3, self.layer4] - - def make_layer(self, block, planes, blocks, stride=1, dcn=None): - downsample = None - if stride != 1 or self.inplanes != planes * block.expansion: - downsample = nn.Sequential( - nn.Conv2d( - self.inplanes, - planes * block.expansion, - kernel_size=1, - stride=stride, - bias=False, - ), - self._norm_layer(planes * block.expansion, momentum=0.1), - ) - - layers = [] - if downsample is not None: - layers.append( - block( - self.inplanes, - planes, - stride, - downsample, - reduction=True, - norm_layer=self._norm_layer, - dcn=dcn, - ) - ) - else: - layers.append( - block( - self.inplanes, - planes, - stride, - downsample, - norm_layer=self._norm_layer, - dcn=dcn, - ) - ) - self.inplanes = planes * block.expansion - for i in range(1, blocks): - layers.append( - block(self.inplanes, planes, norm_layer=self._norm_layer, dcn=dcn) - ) - - return nn.Sequential(*layers) diff --git a/forge/test/tvm/cnn/pytorch/alphapose/models/simplepose.py b/forge/test/tvm/cnn/pytorch/alphapose/models/simplepose.py deleted file mode 100644 index b8d59e49..00000000 --- a/forge/test/tvm/cnn/pytorch/alphapose/models/simplepose.py +++ /dev/null @@ -1,118 +0,0 @@ -# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC - -# SPDX-License-Identifier: Apache-2.0 -# ----------------------------------------------------- -# Copyright (c) Shanghai Jiao Tong University. All rights reserved. -# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com) -# ----------------------------------------------------- - -import torch.nn as nn - -from .builder import SPPE -from .layers.Resnet import ResNet - - -@SPPE.register_module -class SimplePose(nn.Module): - def __init__(self, norm_layer=nn.BatchNorm2d, **cfg): - super(SimplePose, self).__init__() - self._preset_cfg = cfg["PRESET"] - self.deconv_dim = cfg["NUM_DECONV_FILTERS"] - self._norm_layer = norm_layer - - self.preact = ResNet(f"resnet{cfg['NUM_LAYERS']}") - - # Imagenet pretrain model - sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../../") - import torchvision as tm # noqa: F401,F403 - - assert cfg["NUM_LAYERS"] in [18, 34, 50, 101, 152] - x = eval(f"tm.resnet{cfg['NUM_LAYERS']}(pretrained=True)") - - model_state = self.preact.state_dict() - state = { - k: v - for k, v in x.state_dict().items() - if k in self.preact.state_dict() - and v.size() == self.preact.state_dict()[k].size() - } - model_state.update(state) - self.preact.load_state_dict(model_state) - - self.deconv_layers = self._make_deconv_layer() - self.final_layer = nn.Conv2d( - self.deconv_dim[2], - self._preset_cfg["NUM_JOINTS"], - kernel_size=1, - stride=1, - padding=0, - ) - - def _make_deconv_layer(self): - deconv_layers = [] - deconv1 = nn.ConvTranspose2d( - 2048, - self.deconv_dim[0], - kernel_size=4, - stride=2, - padding=int(4 / 2) - 1, - bias=False, - ) - bn1 = self._norm_layer(self.deconv_dim[0]) - deconv2 = nn.ConvTranspose2d( - self.deconv_dim[0], - self.deconv_dim[1], - kernel_size=4, - stride=2, - padding=int(4 / 2) - 1, - bias=False, - ) - bn2 = self._norm_layer(self.deconv_dim[1]) - deconv3 = nn.ConvTranspose2d( - self.deconv_dim[1], - self.deconv_dim[2], - kernel_size=4, - stride=2, - padding=int(4 / 2) - 1, - bias=False, - ) - bn3 = self._norm_layer(self.deconv_dim[2]) - - deconv_layers.append(deconv1) - deconv_layers.append(bn1) - deconv_layers.append(nn.ReLU(inplace=True)) - deconv_layers.append(deconv2) - deconv_layers.append(bn2) - deconv_layers.append(nn.ReLU(inplace=True)) - deconv_layers.append(deconv3) - deconv_layers.append(bn3) - deconv_layers.append(nn.ReLU(inplace=True)) - - return nn.Sequential(*deconv_layers) - - def _initialize(self): - for name, m in self.deconv_layers.named_modules(): - if isinstance(m, nn.ConvTranspose2d): - # logger.info('=> init {}.weight as normal(0, 0.001)'.format(name)) - # logger.info('=> init {}.bias as 0'.format(name)) - nn.init.normal_(m.weight, std=0.001) - # if self.deconv_with_bias: - # nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.BatchNorm2d): - # logger.info('=> init {}.weight as 1'.format(name)) - # logger.info('=> init {}.bias as 0'.format(name)) - nn.init.constant_(m.weight, 1) - nn.init.constant_(m.bias, 0) - for m in self.final_layer.modules(): - if isinstance(m, nn.Conv2d): - # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - # logger.info('=> init {}.weight as normal(0, 0.001)'.format(name)) - # logger.info('=> init {}.bias as 0'.format(name)) - nn.init.normal_(m.weight, std=0.001) - nn.init.constant_(m.bias, 0) - - def forward(self, x): - out = self.preact(x) - out = self.deconv_layers(out) - out = self.final_layer(out) - return out diff --git a/forge/test/tvm/cnn/pytorch/alphapose/utils/__init__.py b/forge/test/tvm/cnn/pytorch/alphapose/utils/__init__.py deleted file mode 100644 index e9bb522e..00000000 --- a/forge/test/tvm/cnn/pytorch/alphapose/utils/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC - -# SPDX-License-Identifier: Apache-2.0 -from .registry import Registry, build_from_cfg - -__all__ = ["Registry", "build_from_cfg"] diff --git a/forge/test/tvm/cnn/pytorch/alphapose/utils/config.py b/forge/test/tvm/cnn/pytorch/alphapose/utils/config.py deleted file mode 100644 index 5f5a97a4..00000000 --- a/forge/test/tvm/cnn/pytorch/alphapose/utils/config.py +++ /dev/null @@ -1,12 +0,0 @@ -# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC - -# SPDX-License-Identifier: Apache-2.0 -import yaml - -# from easydict import EasyDict as edict - - -def update_config(config_file): - with open(config_file) as f: - config = dict(yaml.load(f, Loader=yaml.FullLoader)) - return config diff --git a/forge/test/tvm/cnn/pytorch/alphapose/utils/registry.py b/forge/test/tvm/cnn/pytorch/alphapose/utils/registry.py deleted file mode 100644 index abd4d551..00000000 --- a/forge/test/tvm/cnn/pytorch/alphapose/utils/registry.py +++ /dev/null @@ -1,79 +0,0 @@ -# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC - -# SPDX-License-Identifier: Apache-2.0 -import inspect - - -class Registry(object): - def __init__(self, name): - self._name = name - self._module_dict = dict() - - def __repr__(self): - format_str = self.__class__.__name__ + "(name={}, items={})".format( - self._name, list(self._module_dict.keys()) - ) - return format_str - - @property - def name(self): - return self._name - - @property - def module_dict(self): - return self._module_dict - - def get(self, key): - return self._module_dict.get(key, None) - - def _register_module(self, module_class): - """Register a module. - Args: - module (:obj:`nn.Module`): Module to be registered. - """ - if not inspect.isclass(module_class): - raise TypeError( - "module must be a class, but got {}".format(type(module_class)) - ) - module_name = module_class.__name__ - if module_name in self._module_dict: - raise KeyError( - "{} is already registered in {}".format(module_name, self.name) - ) - self._module_dict[module_name] = module_class - - def register_module(self, cls): - self._register_module(cls) - return cls - - -def build_from_cfg(cfg, registry, default_args=None): - """Build a module from config dict. - Args: - cfg (dict): Config dict. It should at least contain the key "type". - registry (:obj:`Registry`): The registry to search the type from. - default_args (dict, optional): Default initialization arguments. - Returns: - obj: The constructed object. - """ - assert isinstance(cfg, dict) and "TYPE" in cfg - assert isinstance(default_args, dict) or default_args is None - args = cfg.copy() - obj_type = args.pop("TYPE") - - if isinstance(obj_type, str): - obj_cls = registry.get(obj_type) - if obj_cls is None: - raise KeyError( - "{} is not in the {} registry".format(obj_type, registry.name) - ) - elif inspect.isclass(obj_type): - obj_cls = obj_type - else: - raise TypeError( - "type must be a str or valid type, but got {}".format(type(obj_type)) - ) - if default_args is not None: - for name, value in default_args.items(): - args.setdefault(name, value) - return obj_cls(**args) diff --git a/forge/test/tvm/cnn/pytorch/alphapose/utils/transforms.py b/forge/test/tvm/cnn/pytorch/alphapose/utils/transforms.py deleted file mode 100644 index 2cb1f6c4..00000000 --- a/forge/test/tvm/cnn/pytorch/alphapose/utils/transforms.py +++ /dev/null @@ -1,698 +0,0 @@ -# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC - -# SPDX-License-Identifier: Apache-2.0 -# ----------------------------------------------------- -# Copyright (c) Shanghai Jiao Tong University. All rights reserved. -# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com) -# ----------------------------------------------------- - -"""Pose related transforrmation functions.""" - -import torch -import cv2 -import scipy.misc -import numpy as np -import random - - -def rnd(x): - return max(-2 * x, min(2 * x, np.random.randn(1)[0] * x)) - - -def box_transform(bbox, sf, imgwidth, imght, train): - """Random scaling.""" - width = bbox[2] - bbox[0] - ht = bbox[3] - bbox[1] - if train: - scaleRate = 0.25 * np.clip(np.random.randn() * sf, -sf, sf) - - bbox[0] = max(0, bbox[0] - width * scaleRate / 2) - bbox[1] = max(0, bbox[1] - ht * scaleRate / 2) - bbox[2] = min(imgwidth, bbox[2] + width * scaleRate / 2) - bbox[3] = min(imght, bbox[3] + ht * scaleRate / 2) - else: - scaleRate = 0.25 - - bbox[0] = max(0, bbox[0] - width * scaleRate / 2) - bbox[1] = max(0, bbox[1] - ht * scaleRate / 2) - bbox[2] = min(imgwidth, max(bbox[2] + width * scaleRate / 2, bbox[0] + 5)) - bbox[3] = min(imght, max(bbox[3] + ht * scaleRate / 2, bbox[1] + 5)) - - return bbox - - -def addDPG(bbox, imgwidth, imght): - """Add dpg for data augmentation, including random crop and random sample.""" - PatchScale = random.uniform(0, 1) - width = bbox[2] - bbox[0] - ht = bbox[3] - bbox[1] - - if PatchScale > 0.85: - ratio = ht / width - if width < ht: - patchWidth = PatchScale * width - patchHt = patchWidth * ratio - else: - patchHt = PatchScale * ht - patchWidth = patchHt / ratio - - xmin = bbox[0] + random.uniform(0, 1) * (width - patchWidth) - ymin = bbox[1] + random.uniform(0, 1) * (ht - patchHt) - xmax = xmin + patchWidth + 1 - ymax = ymin + patchHt + 1 - else: - xmin = max( - 1, min(bbox[0] + np.random.normal(-0.0142, 0.1158) * width, imgwidth - 3) - ) - ymin = max(1, min(bbox[1] + np.random.normal(0.0043, 0.068) * ht, imght - 3)) - xmax = min( - max(xmin + 2, bbox[2] + np.random.normal(0.0154, 0.1337) * width), - imgwidth - 3, - ) - ymax = min( - max(ymin + 2, bbox[3] + np.random.normal(-0.0013, 0.0711) * ht), imght - 3 - ) - - bbox[0] = xmin - bbox[1] = ymin - bbox[2] = xmax - bbox[3] = ymax - - return bbox - - -def im_to_torch(img): - """Transform ndarray image to torch tensor. - Parameters - ---------- - img: numpy.ndarray - An ndarray with shape: `(H, W, 3)`. - Returns - ------- - torch.Tensor - A tensor with shape: `(3, H, W)`. - """ - img = np.transpose(img, (2, 0, 1)) # C*H*W - img = to_torch(img).float() - if img.max() > 1: - img /= 255 - return img - - -def torch_to_im(img): - """Transform torch tensor to ndarray image. - Parameters - ---------- - img: torch.Tensor - A tensor with shape: `(3, H, W)`. - Returns - ------- - numpy.ndarray - An ndarray with shape: `(H, W, 3)`. - """ - img = to_numpy(img) - img = np.transpose(img, (1, 2, 0)) # C*H*W - return img - - -def load_image(img_path): - # H x W x C => C x H x W - return im_to_torch(scipy.misc.imread(img_path, mode="RGB")) - - -def to_numpy(tensor): - # torch.Tensor => numpy.ndarray - if torch.is_tensor(tensor): - return tensor.cpu().numpy() - elif type(tensor).__module__ != "numpy": - raise ValueError("Cannot convert {} to numpy array".format(type(tensor))) - return tensor - - -def to_torch(ndarray): - # numpy.ndarray => torch.Tensor - if type(ndarray).__module__ == "numpy": - return torch.from_numpy(ndarray) - elif not torch.is_tensor(ndarray): - raise ValueError("Cannot convert {} to torch tensor".format(type(ndarray))) - return ndarray - - -def cv_cropBox(img, bbox, input_size): - """Crop bbox from image by Affinetransform. - Parameters - ---------- - img: torch.Tensor - A tensor with shape: `(3, H, W)`. - bbox: list or tuple - [xmin, ymin, xmax, ymax]. - input_size: tuple - Resulting image size, as (height, width). - Returns - ------- - torch.Tensor - A tensor with shape: `(3, height, width)`. - """ - xmin, ymin, xmax, ymax = bbox - xmax -= 1 - ymax -= 1 - resH, resW = input_size - - lenH = max((ymax - ymin), (xmax - xmin) * resH / resW) - lenW = lenH * resW / resH - if img.dim() == 2: - img = img[np.newaxis, :, :] - - box_shape = [ymax - ymin, xmax - xmin] - pad_size = [(lenH - box_shape[0]) // 2, (lenW - box_shape[1]) // 2] - # Padding Zeros - img[:, :ymin, :], img[:, :, :xmin] = 0, 0 - img[:, ymax + 1 :, :], img[:, :, xmax + 1 :] = 0, 0 - - src = np.zeros((3, 2), dtype=np.float32) - dst = np.zeros((3, 2), dtype=np.float32) - - src[0, :] = np.array([xmin - pad_size[1], ymin - pad_size[0]], np.float32) - src[1, :] = np.array([xmax + pad_size[1], ymax + pad_size[0]], np.float32) - dst[0, :] = 0 - dst[1, :] = np.array([resW - 1, resH - 1], np.float32) - - src[2:, :] = get_3rd_point(src[0, :], src[1, :]) - dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :]) - - trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) - dst_img = cv2.warpAffine( - torch_to_im(img), trans, (resW, resH), flags=cv2.INTER_LINEAR - ) - if dst_img.ndim == 2: - dst_img = dst_img[:, :, np.newaxis] - - return im_to_torch(torch.Tensor(dst_img)) - - -def cv_cropBox_rot(img, bbox, input_size, rot): - """Crop bbox from image by Affinetransform. - Parameters - ---------- - img: torch.Tensor - A tensor with shape: `(3, H, W)`. - bbox: list or tuple - [xmin, ymin, xmax, ymax]. - input_size: tuple - Resulting image size, as (height, width). - Returns - ------- - torch.Tensor - A tensor with shape: `(3, height, width)`. - """ - xmin, ymin, xmax, ymax = bbox - xmax -= 1 - ymax -= 1 - resH, resW = input_size - rot_rad = np.pi * rot / 180 - - if img.dim() == 2: - img = img[np.newaxis, :, :] - - src = np.zeros((3, 2), dtype=np.float32) - dst = np.zeros((3, 2), dtype=np.float32) - center = np.array([(xmax + xmin) / 2, (ymax + ymin) / 2]) - - src_dir = get_dir([0, (ymax - ymin) * -0.5], rot_rad) - dst_dir = np.array([0, (resH - 1) * -0.5], np.float32) - - src = np.zeros((3, 2), dtype=np.float32) - dst = np.zeros((3, 2), dtype=np.float32) - - src[0, :] = center - src[1, :] = center + src_dir - dst[0, :] = [(resW - 1) * 0.5, (resH - 1) * 0.5] - dst[1, :] = np.array([(resW - 1) * 0.5, (resH - 1) * 0.5]) + dst_dir - - src[2:, :] = get_3rd_point(src[0, :], src[1, :]) - dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :]) - - trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) - dst_img = cv2.warpAffine( - torch_to_im(img), trans, (resW, resH), flags=cv2.INTER_LINEAR - ) - if dst_img.ndim == 2: - dst_img = dst_img[:, :, np.newaxis] - - return im_to_torch(torch.Tensor(dst_img)) - - -def fix_cropBox(img, bbox, input_size): - """Crop bbox from image by Affinetransform. - Parameters - ---------- - img: torch.Tensor - A tensor with shape: `(3, H, W)`. - bbox: list or tuple - [xmin, ymin, xmax, ymax]. - input_size: tuple - Resulting image size, as (height, width). - Returns - ------- - torch.Tensor - A tensor with shape: `(3, height, width)`. - """ - xmin, ymin, xmax, ymax = bbox - input_ratio = input_size[0] / input_size[1] - bbox_ratio = (ymax - ymin) / (xmax - xmin) - if bbox_ratio > input_ratio: - # expand width - cx = (xmax + xmin) / 2 - h = ymax - ymin - w = h / input_ratio - xmin = cx - w / 2 - xmax = cx + w / 2 - elif bbox_ratio < input_ratio: - # expand height - cy = (ymax + ymin) / 2 - w = xmax - xmin - h = w * input_ratio - ymin = cy - h / 2 - ymax = cy + h / 2 - bbox = [int(x) for x in [xmin, ymin, xmax, ymax]] - - return cv_cropBox(img, bbox, input_size), bbox - - -def fix_cropBox_rot(img, bbox, input_size, rot): - """Crop bbox from image by Affinetransform. - Parameters - ---------- - img: torch.Tensor - A tensor with shape: `(3, H, W)`. - bbox: list or tuple - [xmin, ymin, xmax, ymax]. - input_size: tuple - Resulting image size, as (height, width). - Returns - ------- - torch.Tensor - A tensor with shape: `(3, height, width)`. - """ - xmin, ymin, xmax, ymax = bbox - input_ratio = input_size[0] / input_size[1] - bbox_ratio = (ymax - ymin) / (xmax - xmin) - if bbox_ratio > input_ratio: - # expand width - cx = (xmax + xmin) / 2 - h = ymax - ymin - w = h / input_ratio - xmin = cx - w / 2 - xmax = cx + w / 2 - elif bbox_ratio < input_ratio: - # expand height - cy = (ymax + ymin) / 2 - w = xmax - xmin - h = w * input_ratio - ymin = cy - h / 2 - ymax = cy + h / 2 - bbox = [int(x) for x in [xmin, ymin, xmax, ymax]] - - return cv_cropBox_rot(img, bbox, input_size, rot), bbox - - -def get_3rd_point(a, b): - """Return vector c that perpendicular to (a - b).""" - direct = a - b - return b + np.array([-direct[1], direct[0]], dtype=np.float32) - - -def get_dir(src_point, rot_rad): - """Rotate the point by `rot_rad` degree.""" - sn, cs = np.sin(rot_rad), np.cos(rot_rad) - - src_result = [0, 0] - src_result[0] = src_point[0] * cs - src_point[1] * sn - src_result[1] = src_point[0] * sn + src_point[1] * cs - - return src_result - - -def cv_cropBoxInverse(inp, bbox, img_size, output_size): - """Paste the cropped bbox to the original image. - Parameters - ---------- - inp: torch.Tensor - A tensor with shape: `(3, height, width)`. - bbox: list or tuple - [xmin, ymin, xmax, ymax]. - img_size: tuple - Original image size, as (img_H, img_W). - output_size: tuple - Cropped input size, as (height, width). - Returns - ------- - torch.Tensor - A tensor with shape: `(3, img_H, img_W)`. - """ - xmin, ymin, xmax, ymax = bbox - xmax -= 1 - ymax -= 1 - resH, resW = output_size - imgH, imgW = img_size - - lenH = max((ymax - ymin), (xmax - xmin) * resH / resW) - lenW = lenH * resW / resH - if inp.dim() == 2: - inp = inp[np.newaxis, :, :] - - box_shape = [ymax - ymin, xmax - xmin] - pad_size = [(lenH - box_shape[0]) // 2, (lenW - box_shape[1]) // 2] - - src = np.zeros((3, 2), dtype=np.float32) - dst = np.zeros((3, 2), dtype=np.float32) - - src[0, :] = 0 - src[1, :] = np.array([resW - 1, resH - 1], np.float32) - dst[0, :] = np.array([xmin - pad_size[1], ymin - pad_size[0]], np.float32) - dst[1, :] = np.array([xmax + pad_size[1], ymax + pad_size[0]], np.float32) - - src[2:, :] = get_3rd_point(src[0, :], src[1, :]) - dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :]) - - trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) - dst_img = cv2.warpAffine( - torch_to_im(inp), trans, (imgW, imgH), flags=cv2.INTER_LINEAR - ) - if dst_img.ndim == 3 and dst_img.shape[2] == 1: - dst_img = dst_img[:, :, 0] - return dst_img - elif dst_img.ndim == 2: - return dst_img - else: - return im_to_torch(torch.Tensor(dst_img)) - - -def cv_rotate(img, rot, input_size): - """Rotate image by Affinetransform. - Parameters - ---------- - img: torch.Tensor - A tensor with shape: `(3, H, W)`. - rot: int - Rotation degree. - input_size: tuple - Resulting image size, as (height, width). - Returns - ------- - torch.Tensor - A tensor with shape: `(3, height, width)`. - """ - resH, resW = input_size - center = np.array((resW - 1, resH - 1)) / 2 - rot_rad = np.pi * rot / 180 - - src_dir = get_dir([0, (resH - 1) * -0.5], rot_rad) - dst_dir = np.array([0, (resH - 1) * -0.5], np.float32) - - src = np.zeros((3, 2), dtype=np.float32) - dst = np.zeros((3, 2), dtype=np.float32) - - src[0, :] = center - src[1, :] = center + src_dir - dst[0, :] = [(resW - 1) * 0.5, (resH - 1) * 0.5] - dst[1, :] = np.array([(resW - 1) * 0.5, (resH - 1) * 0.5]) + dst_dir - - src[2:, :] = get_3rd_point(src[0, :], src[1, :]) - dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :]) - - trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) - - dst_img = cv2.warpAffine( - torch_to_im(img), trans, (resW, resH), flags=cv2.INTER_LINEAR - ) - if dst_img.ndim == 2: - dst_img = dst_img[:, :, np.newaxis] - - return im_to_torch(torch.Tensor(dst_img)) - - -def count_visible(bbox, joints_3d): - """Count number of visible joints given bound box.""" - vis = np.logical_and.reduce( - ( - joints_3d[:, 0, 0] > 0, - joints_3d[:, 0, 0] > bbox[0], - joints_3d[:, 0, 0] < bbox[2], - joints_3d[:, 1, 0] > 0, - joints_3d[:, 1, 0] > bbox[1], - joints_3d[:, 1, 0] < bbox[3], - joints_3d[:, 0, 1] > 0, - joints_3d[:, 1, 1] > 0, - ) - ) - return np.sum(vis), vis - - -def drawGaussian(img, pt, sigma): - """Draw 2d gaussian on input image. - Parameters - ---------- - img: torch.Tensor - A tensor with shape: `(3, H, W)`. - pt: list or tuple - A point: (x, y). - sigma: int - Sigma of gaussian distribution. - Returns - ------- - torch.Tensor - A tensor with shape: `(3, H, W)`. - """ - img = to_numpy(img) - tmpSize = 3 * sigma - # Check that any part of the gaussian is in-bounds - ul = [int(pt[0] - tmpSize), int(pt[1] - tmpSize)] - br = [int(pt[0] + tmpSize + 1), int(pt[1] + tmpSize + 1)] - - if ul[0] >= img.shape[1] or ul[1] >= img.shape[0] or br[0] < 0 or br[1] < 0: - # If not, just return the image as is - return to_torch(img) - - # Generate gaussian - size = 2 * tmpSize + 1 - x = np.arange(0, size, 1, float) - y = x[:, np.newaxis] - x0 = y0 = size // 2 - # The gaussian is not normalized, we want the center value to equal 1 - g = np.exp(-((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2)) - - # Usable gaussian range - g_x = max(0, -ul[0]), min(br[0], img.shape[1]) - ul[0] - g_y = max(0, -ul[1]), min(br[1], img.shape[0]) - ul[1] - # Image range - img_x = max(0, ul[0]), min(br[0], img.shape[1]) - img_y = max(0, ul[1]), min(br[1], img.shape[0]) - - img[img_y[0] : img_y[1], img_x[0] : img_x[1]] = g[g_y[0] : g_y[1], g_x[0] : g_x[1]] - return to_torch(img) - - -def flip(x): - assert x.dim() == 3 or x.dim() == 4 - dim = x.dim() - 1 - - return x.flip(dims=(dim,)) - - -def flip_heatmap(heatmap, joint_pairs, shift=False): - """Flip pose heatmap according to joint pairs. - Parameters - ---------- - heatmap : numpy.ndarray - Heatmap of joints. - joint_pairs : list - List of joint pairs. - shift : bool - Whether to shift the output. - Returns - ------- - numpy.ndarray - Flipped heatmap. - """ - assert heatmap.dim() == 3 or heatmap.dim() == 4 - out = flip(heatmap) - - for pair in joint_pairs: - dim0, dim1 = pair - idx = torch.Tensor((dim0, dim1)).long() - inv_idx = torch.Tensor((dim1, dim0)).long() - if out.dim() == 4: - out[:, idx] = out[:, inv_idx] - else: - out[idx] = out[inv_idx] - - if shift: - if out.dim() == 3: - out[:, :, 1:] = out[:, :, 0:-1] - else: - out[:, :, :, 1:] = out[:, :, :, 0:-1] - return out - - -def flip_joints_3d(joints_3d, width, joint_pairs): - """Flip 3d joints. - Parameters - ---------- - joints_3d : numpy.ndarray - Joints in shape (num_joints, 3, 2) - width : int - Image width. - joint_pairs : list - List of joint pairs. - Returns - ------- - numpy.ndarray - Flipped 3d joints with shape (num_joints, 3, 2) - """ - joints = joints_3d.copy() - # flip horizontally - joints[:, 0, 0] = width - joints[:, 0, 0] - 1 - # change left-right parts - for pair in joint_pairs: - joints[pair[0], :, 0], joints[pair[1], :, 0] = ( - joints[pair[1], :, 0], - joints[pair[0], :, 0].copy(), - ) - joints[pair[0], :, 1], joints[pair[1], :, 1] = ( - joints[pair[1], :, 1], - joints[pair[0], :, 1].copy(), - ) - - joints[:, :, 0] *= joints[:, :, 1] - return joints - - -def heatmap_to_coord_simple(hms, bbox): - coords, maxvals = get_max_pred(hms) - - hm_h = hms.shape[1] - hm_w = hms.shape[2] - - # post-processing - for p in range(coords.shape[0]): - hm = hms[p] - px = int(round(float(coords[p][0]))) - py = int(round(float(coords[p][1]))) - if 1 < px < hm_w - 1 and 1 < py < hm_h - 1: - diff = np.array( - (hm[py][px + 1] - hm[py][px - 1], hm[py + 1][px] - hm[py - 1][px]) - ) - coords[p] += np.sign(diff) * 0.25 - - preds = np.zeros_like(coords) - - # transform bbox to scale - xmin, ymin, xmax, ymax = bbox - w = xmax - xmin - h = ymax - ymin - center = np.array([xmin + w * 0.5, ymin + h * 0.5]) - scale = np.array([w, h]) - # Transform back - for i in range(coords.shape[0]): - preds[i] = transform_preds(coords[i], center, scale, [hm_w, hm_h]) - - return preds, maxvals - - -def transform_preds(coords, center, scale, output_size): - target_coords = np.zeros(coords.shape) - trans = get_affine_transform(center, scale, 0, output_size, inv=1) - target_coords[0:2] = affine_transform(coords[0:2], trans) - return target_coords - - -def get_max_pred(heatmaps): - num_joints = heatmaps.shape[0] - width = heatmaps.shape[2] - heatmaps_reshaped = heatmaps.reshape((num_joints, -1)) - idx = np.argmax(heatmaps_reshaped, 1) - maxvals = np.max(heatmaps_reshaped, 1) - - maxvals = maxvals.reshape((num_joints, 1)) - idx = idx.reshape((num_joints, 1)) - - preds = np.tile(idx, (1, 2)).astype(np.float32) - - preds[:, 0] = (preds[:, 0]) % width - preds[:, 1] = np.floor((preds[:, 1]) / width) - - pred_mask = np.tile(np.greater(maxvals, 0.0), (1, 2)) - pred_mask = pred_mask.astype(np.float32) - - preds *= pred_mask - return preds, maxvals - - -def get_max_pred_batch(batch_heatmaps): - batch_size = batch_heatmaps.shape[0] - num_joints = batch_heatmaps.shape[1] - width = batch_heatmaps.shape[3] - heatmaps_reshaped = batch_heatmaps.reshape((batch_size, num_joints, -1)) - idx = np.argmax(heatmaps_reshaped, 2) - maxvals = np.max(heatmaps_reshaped, 2) - - maxvals = maxvals.reshape((batch_size, num_joints, 1)) - idx = idx.reshape((batch_size, num_joints, 1)) - - preds = np.tile(idx, (1, 1, 2)).astype(np.float32) - - preds[:, :, 0] = (preds[:, :, 0]) % width - preds[:, :, 1] = np.floor((preds[:, :, 1]) / width) - - pred_mask = np.tile(np.greater(maxvals, 0.0), (1, 1, 2)) - pred_mask = pred_mask.astype(np.float32) - - preds *= pred_mask - return preds, maxvals - - -def get_affine_transform( - center, scale, rot, output_size, shift=np.array([0, 0], dtype=np.float32), inv=0 -): - if not isinstance(scale, np.ndarray) and not isinstance(scale, list): - scale = np.array([scale, scale]) - - scale_tmp = scale - src_w = scale_tmp[0] - dst_w = output_size[0] - dst_h = output_size[1] - - rot_rad = np.pi * rot / 180 - src_dir = get_dir([0, src_w * -0.5], rot_rad) - dst_dir = np.array([0, dst_w * -0.5], np.float32) - - src = np.zeros((3, 2), dtype=np.float32) - dst = np.zeros((3, 2), dtype=np.float32) - src[0, :] = center + scale_tmp * shift - src[1, :] = center + src_dir + scale_tmp * shift - dst[0, :] = [dst_w * 0.5, dst_h * 0.5] - dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir - - src[2:, :] = get_3rd_point(src[0, :], src[1, :]) - dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :]) - - if inv: - trans = cv2.getAffineTransform(np.float32(dst), np.float32(src)) - else: - trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) - - return trans - - -def affine_transform(pt, t): - new_pt = np.array([pt[0], pt[1], 1.0]).T - new_pt = np.dot(t, new_pt) - return new_pt[:2] - - -def get_func_heatmap_to_coord(cfg): - if cfg.DATA_PRESET.TYPE == "simple": - return heatmap_to_coord_simple - else: - raise NotImplementedError