text: # Copyright (c) OpenMMLab. All rights reserved.
  import mmcv
  import numpy as np
  import torch
  
  from mmpose.models import build_posenet
  
  
  def _create_inputs(joint_num_in,
                     joint_channel_in,
                     joint_num_out,
                     joint_channel_out,
                     seq_len,
                     batch_size,
                     semi=False):
      rng = np.random.RandomState(0)
      pose_in = rng.rand(batch_size, joint_num_in * joint_channel_in, seq_len)
      target = np.zeros((batch_size, joint_num_out, joint_channel_out),
                        dtype=np.float32)
      target_weight = np.ones((batch_size, joint_num_out, joint_channel_out),
                              dtype=np.float32)
  
      meta_info = {
          'root_position': np.zeros((1, joint_channel_out), np.float32),
          'root_position_index': 0,
          'target_mean': np.zeros((joint_num_out, joint_channel_out),
                                  np.float32),
          'target_std': np.ones((joint_num_out, joint_channel_out), np.float32)
      }
      metas = [meta_info.copy() for _ in range(batch_size)]
      inputs = {
          'input': torch.FloatTensor(pose_in).requires_grad_(True),
          'target': torch.FloatTensor(target),
          'target_weight': torch.FloatTensor(target_weight),
          'metas': metas,
      }
  
      if semi:
          traj_target = np.zeros((batch_size, 1, joint_channel_out), np.float32)
          unlabeled_pose_in = rng.rand(batch_size,
                                       joint_num_in * joint_channel_in, seq_len)
          unlabeled_target_2d = np.zeros(
              (batch_size, joint_num_out, joint_channel_in), dtype=np.float32)
          intrinsics = np.ones((batch_size, 4))
  
          inputs['traj_target'] = torch.FloatTensor(traj_target)
          inputs['unlabeled_input'] = torch.FloatTensor(
              unlabeled_pose_in).requires_grad_(True)
          inputs['unlabeled_target_2d'] = torch.FloatTensor(unlabeled_target_2d)
          inputs['intrinsics'] = torch.FloatTensor(intrinsics)
  
      return inputs
  
  
  def test_pose_lifter_forward():
      # Test forward train for supervised learning with pose model only
      model_cfg = dict(
          type='PoseLifter',
          pretrained=None,
          backbone=dict(type='TCN', in_channels=2 * 17),
          keypoint_head=dict(
              type='TemporalRegressionHead',
              in_channels=1024,
              num_joints=16,
              max_norm=1.0,
              loss_keypoint=dict(type='MPJPELoss'),
              test_cfg=dict(restore_global_position=True)),
          train_cfg=dict(),
          test_cfg=dict())
  
      cfg = mmcv.Config({'model': model_cfg})
      detector = build_posenet(cfg.model)
  
      detector.init_weights()
  
      inputs = _create_inputs(
          joint_num_in=17,
          joint_channel_in=2,
          joint_num_out=16,
          joint_channel_out=3,
          seq_len=27,
          batch_size=8)
  
      losses = detector.forward(
          inputs['input'],
          inputs['target'],
          inputs['target_weight'],
          inputs['metas'],
          return_loss=True)
  
      assert isinstance(losses, dict)
  
      # Test forward test for supervised learning wiht pose model only
      with torch.no_grad():
          _ = detector.forward(
              inputs['input'],
              inputs['target'],
              inputs['target_weight'],
              inputs['metas'],
              return_loss=False)
          _ = detector.forward_dummy(inputs['input'])
  
      # Test forward train for semi-supervised learning
      model_cfg = dict(
          type='PoseLifter',
          pretrained=None,
          backbone=dict(type='TCN', in_channels=2 * 17),
          keypoint_head=dict(
              type='TemporalRegressionHead',
              in_channels=1024,
              num_joints=17,
              loss_keypoint=dict(type='MPJPELoss'),
              test_cfg=dict(restore_global_position=True)),
          traj_backbone=dict(type='TCN', in_channels=2 * 17),
          traj_head=dict(
              type='TemporalRegressionHead',
              in_channels=1024,
              num_joints=1,
              loss_keypoint=dict(type='MPJPELoss'),
              is_trajectory=True),
          loss_semi=dict(
              type='SemiSupervisionLoss',
              joint_parents=[
                  0, 0, 1, 2, 0, 4, 5, 0, 7, 8, 9, 8, 11, 12, 8, 14, 15
              ]),
          train_cfg=dict(),
          test_cfg=dict())
  
      cfg = mmcv.Config({'model': model_cfg})
      detector = build_posenet(cfg.model)
  
      detector.init_weights()
  
      inputs = _create_inputs(
          joint_num_in=17,
          joint_channel_in=2,
          joint_num_out=17,
          joint_channel_out=3,
          seq_len=27,
          batch_size=8,
          semi=True)
  
      losses = detector.forward(**inputs, return_loss=True)
  
      assert isinstance(losses, dict)
      assert 'proj_loss' in losses
  
      # Test forward test for semi-supervised learning
      with torch.no_grad():
          _ = detector.forward(**inputs, return_loss=False)
          _ = detector.forward_dummy(inputs['input'])
  
      # Test forward train for supervised learning with pose model and trajectory
      # model sharing one backbone
      model_cfg = dict(
          type='PoseLifter',
          pretrained=None,
          backbone=dict(type='TCN', in_channels=2 * 17),
          keypoint_head=dict(
              type='TemporalRegressionHead',
              in_channels=1024,
              num_joints=17,
              loss_keypoint=dict(type='MPJPELoss'),
              test_cfg=dict(restore_global_position=True)),
          traj_head=dict(
              type='TemporalRegressionHead',
              in_channels=1024,
              num_joints=1,
              loss_keypoint=dict(type='MPJPELoss'),
              is_trajectory=True),
          train_cfg=dict(),
          test_cfg=dict())
  
      cfg = mmcv.Config({'model': model_cfg})
      detector = build_posenet(cfg.model)
  
      detector.init_weights()
  
      inputs = _create_inputs(
          joint_num_in=17,
          joint_channel_in=2,
          joint_num_out=17,
          joint_channel_out=3,
          seq_len=27,
          batch_size=8,
          semi=True)
  
      losses = detector.forward(**inputs, return_loss=True)
  
      assert isinstance(losses, dict)
      assert 'traj_loss' in losses
  
      # Test forward test for semi-supervised learning with pose model and
      # trajectory model sharing one backbone
      with torch.no_grad():
          _ = detector.forward(**inputs, return_loss=False)
          _ = detector.forward_dummy(inputs['input'])
  