def inference(model, img, bs=5):
    """Inference DWPose model implemented in TorchScript.

    Args:
        model : TorchScript Model.
        img : Input image in shape.

    Returns:
        outputs : Output of DWPose model.
    """
    all_out = []
    # build input
    orig_img_count = len(img)
    #Pad zeros to fit batch size
    for _ in range(bs - (orig_img_count % bs)):
        img.append(np.zeros_like(img[0]))
    input = np.stack(img, axis=0).transpose(0, 3, 1, 2)
    device, dtype = next(model.parameters()).device, next(model.parameters()).dtype
    input = torch.from_numpy(input).to(device, dtype)

    out1, out2 = [], []
    for i in range(input.shape[0] // bs):
        curr_batch_output = model(input[i*bs:(i+1)*bs])
        out1.append(curr_batch_output[0].float())
        out2.append(curr_batch_output[1].float())
    out1, out2 = torch.cat(out1, dim=0)[:orig_img_count], torch.cat(out2, dim=0)[:orig_img_count]
    out1, out2 = out1.float().cpu().detach().numpy(), out2.float().cpu().detach().numpy()
    all_outputs = out1, out2

    for batch_idx in range(len(all_outputs[0])):
        outputs = [all_outputs[i][batch_idx:batch_idx+1,...] for i in range(len(all_outputs))]
        all_out.append(outputs)
    return all_out
Inference DWPose model implemented in TorchScript.

    Args:
        model : TorchScript Model.
        img : Input image in shape.

    Returns:
        outputs : Output of DWPose model.
    
mimicmotion/dwpose/jit_pose.py