Add atr and pascal model of human part segmentation

[kyakuno]

current implmentation (lip only) https://github.com/axinc-ai/ailia-models/tree/master/image_segmentation/human_part_segmentation

source https://github.com/mayankgrwl97/human-part-segmentation

atr and pascal model https://github.com/mayankgrwl97/human-part-segmentation/blob/master/human_part_segmentation/api.py

[kyakuno]

You can select the model via dataset argument. https://github.com/mayankgrwl97/human-part-segmentation/blob/master/human_part_segmentation_demo.ipynb

hps = HumanPartSegmentation(dataset='lip', gpu='0')

[kyakuno]

Add arch option to current implementation.

[kyakuno]

I already checked to working atr and pascal model on colaboratory. https://colab.research.google.com/github/mayankgrwl97/human-part-segmentation/blob/master/human_part_segmentation_demo.ipynb

[kyakuno]

@ooe1123 animalposeの後に検討いただければと思います。

[ooe1123]

○ ONNXエクスポートできるように、outputの形式を修正

• human_part_segmentation/networks/AugmentCE2P.py

  class ResNet(nn.Module):
  def forward(self, x):
    return [[parsing_result, fusion_result], [edge_result]]

  ↓

  class ResNet(nn.Module):
  def forward(self, x):
    return parsing_result, fusion_result, edge_result

○ CUDAでない実装を使うようにする

• human_part_segmentation/modules/bn.py

  class InPlaceABNSync(ABN):
  def forward(self, x):
      x, _, _ =  inplace_abn_sync(x, self.weight, self.bias, self.running_mean, self.running_var,
                                 self.training, self.momentum, self.eps, self.activation, self.slope)
      return x

  ↓

  class InPlaceABNSync(ABN):
  def forward(self, x):
      x = super().forward(x)
      return x

○ 起動に失敗する場合はCUDAのコードをコメントアウトする

• human_part_segmentation/modules/functions.py

  _backend = load(name="inplace_abn",
              extra_cflags=["-O3"],
              sources=[path.join(_src_path, f) for f in [
                  ...
              ]],
              extra_cuda_cflags=["--expt-extended-lambda"])

  ↓

  # _backend = load(name="inplace_abn",
              # extra_cflags=["-O3"],
              # sources=[path.join(_src_path, f) for f in [
                  ...
              # ]],
              # extra_cuda_cflags=["--expt-extended-lambda"])

[ooe1123]

ONNXエクスポート時、AdaptiveAvgPool2dでサイズがアンマッチのエラーがでるため、独自の実装に置き換える。

# 32x32→3x3に畳み込み
class MyModule(nn.Module):
    def forward(self, x):
        xx = torch.zeros(1,2048,3,3)
        for i,(y0,y1) in enumerate(((0,11),(10,22),(21,33))):
            for j, (x0,x1) in enumerate(((0,11),(10,22),(21,32))):
                xx[:,:,i:i+1,j:j+1] = torch.mean(x[:,:,y0:y1,x0:x1], dim=(2,3), keepdim=True)
        return xx

# 32x32→6x6に畳み込み
class MyModule2(nn.Module):
    def forward(self, x):
        xx = torch.zeros(1,2048,6,6)
        for i,(y0,y1) in enumerate(((0,6),(5,11),(10,16),(16,22),(21,27),(26,32))):
            for j, (x0,x1) in enumerate(((0,6),(5,11),(10,16),(16,22),(21,27),(26,32))):
                xx[:,:,i:i+1,j:j+1] = torch.mean(x[:,:,y0:y1,x0:x1], dim=(2,3), keepdim=True)
        return xx

class HumanPartSegmentation:
    def __init__(self, dataset='lip', gpu='0'):
        ...
        self.model.load_state_dict(new_state_dict)

        # 重みロード後にモジュールを置き換え
        self.model.context_encoding.stages[2][0] = MyModule()
        self.model.context_encoding.stages[3][0] = MyModule2()

[参考] https://qiita.com/kaeruair/items/17d5e0372963a3d3f50b