def generate_motion()

in long_term/pose_network_long_term.py [0:0]

• 55 lines of code
• 7 McCabe index (conditional complexity)

    def generate_motion(self, spline, action_prefix):
        """
        Generate a locomotion sequence using this model.
        Arguments:
         -- spline: the path to follow (an object of type Spline), annotated with all extra inputs (e.g. local speed).
         -- action_prefix: the action which will be used to initialize the model for the first "prefix_length" frames.
        """
        with torch.no_grad():
            start_time = time()
            rot0 = action_prefix['rotations'][:, :self.prefix_length].reshape(-1, self.skeleton.num_joints()*4)
            extra0 = action_prefix['extra_features'][:, :self.prefix_length]
            phase0 = np.arctan2(extra0[-1, 3], extra0[-1, 2])
            
            inputs = torch.from_numpy(np.concatenate((rot0, extra0), axis=1).astype('float32')).unsqueeze(0)
            if self.use_cuda:
                inputs = inputs.cuda()

            predicted, hidden = self.model(inputs)
            
            def next_frame(phase, amplitude, direction, facing_direction):
                nonlocal predicted, hidden
                features = torch.Tensor([np.cos(phase)*amplitude, np.sin(phase)*amplitude,
                                         direction[0]*amplitude, direction[1]*amplitude,
                                         facing_direction[0], facing_direction[1]]).view(1, 1, -1)
                if self.use_cuda:
                    features = features.cuda()
                predicted, hidden = self.model(torch.cat((predicted, features), dim=2), hidden)
                return predicted[:, :, :self.skeleton.num_joints()*4].cpu().numpy().reshape(-1, 4), \
                       predicted[0, 0, -2].item(), predicted[0, 0, -1].item() # Spline offset and height

            current_phase = phase0

            travel_distance = 0 # Longitudinal distance traveled along the spline
            rotations = [] # Joint rotations at each time step
            positions = [] # Root joint positions at each time step
            speed = 0
            stop = False

            while not stop:
                direction = spline.interpolate(travel_distance, 'tangent')
                facing_direction = spline.interpolate(travel_distance, 'direction')
                avg_speed = spline.interpolate(travel_distance, 'amplitude').squeeze()
                freq = spline.interpolate(travel_distance, 'frequency').squeeze()

                rot, displacement, height = next_frame(current_phase, avg_speed, direction, facing_direction)

                current_phase = (current_phase + freq) % (2*np.pi)
                travel_distance += avg_speed
                next_distance = travel_distance + displacement

                if next_distance < spline.length():
                    rotations.append(rot)
                    positions.append((next_distance, height))
                else:
                    # End of spline reached
                    stop = True

            out_trajectory = torch.FloatTensor(len(positions), 3)
            distances, heights = zip(*positions)
            out_trajectory[:, [0, 2]] = torch.from_numpy(spline.interpolate(distances).astype('float32'))
            out_trajectory[:, 1] = torch.FloatTensor(list(heights))
            out_trajectory = out_trajectory.unsqueeze(0)
            rotations = torch.FloatTensor(rotations).unsqueeze(0)
            
            if self.use_cuda:
                out_trajectory = out_trajectory.cuda()
                rotations = rotations.cuda()

            result = self.skeleton.forward_kinematics(rotations, out_trajectory).squeeze(0).cpu().numpy()
            end_time = time()
            duration = end_time - start_time
            fps = result.shape[0] / duration
            print('%d frames generated in %.3f seconds (%.2f FPS)' % (result.shape[0], duration, fps))
            return result