pretrained model_hboe.pth does not work well

[LilKoke]

I tried to run inference on my own image with the pretrained hboe model, but it does not output good result at all. It always outputs angle between 175 ~ 190 deg even if I input a image of a rear-facing person. Here's the code I wrote. Does anyone know what's wrong with my inference code?

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import os
import pprint
import shutil

import torch
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
from tensorboardX import SummaryWriter

import _init_paths
from config import cfg
from config import update_config

from core.loss import JointsMSELoss
from core.loss import DepthLoss
from core.loss import hoe_diff_loss
from core.loss import Bone_loss

from core.function import train
from core.function import validate

from utils.utils import get_optimizer
from utils.utils import save_checkpoint
from utils.utils import create_logger
from utils.utils import get_model_summary

import dataset
import models
from PIL import Image

def parse_args():
    parser = argparse.ArgumentParser(description='Train keypoints network')
    # general
    parser.add_argument('--cfg',
                        help='experiment configure file name',
                        required=True,
                        type=str)

    parser.add_argument('opts',
                        help="Modify config options using the command-line",
                        default=None,
                        nargs=argparse.REMAINDER)

    # philly
    parser.add_argument('--modelDir',
                        help='model directory',
                        type=str,
                        default='')
    parser.add_argument('--logDir',
                        help='log directory',
                        type=str,
                        default='')
    parser.add_argument('--dataDir',
                        help='data directory',
                        type=str,
                        default='')
    parser.add_argument('--prevModelDir',
                        help='prev Model directory',
                        type=str,
                        default='')
    parser.add_argument('--device', default='cpu')
    parser.add_argument('img_path')

    args = parser.parse_args()

    return args


def main():
    args = parse_args()
    update_config(cfg, args)

    logger, _, _ = create_logger(
        cfg, args.cfg, 'valid')

    logger.info(pprint.pformat(args))
    logger.info(cfg)

    # cudnn related setting
    cudnn.benchmark = cfg.CUDNN.BENCHMARK
    torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
    torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED

    model = eval('models.'+cfg.MODEL.NAME+'.get_pose_net')(
        cfg, is_train=False
    ).to(args.device)

    logger.info('=> loading model from {}'.format(cfg.TEST.MODEL_FILE))
    model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE, map_location=torch.device(args.device)), strict=False)

    # Data loading code
    normalize = transforms.Normalize(
        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
    )

    transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Resize((640, 480)),
            normalize])

    img = Image.open(args.img_path).convert("RGB")
    input = transform(img).unsqueeze(0)
    _, hoe_output = model(input)
    sum_hoe = 0
    for i in range(72):
        deg = i * 5
        sum_hoe += deg * hoe_output[0, i]
    print(sum_hoe)

if __name__ == '__main__':
    main()

[ChenyanWu]

Maybe due to the following problems: 1. You should find the peak value of hoe_output to get the hoe prediction, other than sum all the elements. 2. Image size should be 256*192. 3. To be consistent with the original code, use cv2 to read the image. (Using PIL seems not to be a big issue)

[LilKoke]

Thank you for your reply. I fixed three points you mentioned, but it seems like it's not still working...

    transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Resize((256, 192)),
            normalize])

    img = cv2.imread(args.img_path)
    input = transform(img).unsqueeze(0)
    print(input.shape)
    _, hoe_output = model(input)
    ori = torch.argmax(hoe_output[0]) * 5
    print(ori)
    import matplotlib.pyplot as plt
    for i in range(hoe_output.shape[1]):
        plt.scatter(i * 5, hoe_output[0, i].detach().numpy())
    plt.savefig("plot.png")

I might be using inappropriate image for this model, so I share the input image I used. The output of the code above is as follows:

torch.Size([1, 3, 256, 192])
tensor(175)

[ChenyanWu]

I have figured out the issue. Please add model.eval() in front of _, hoe_output = model(input). Based on your code, I write a demo file. Would you mind I upload it to the repo?

[LilKoke]

It worked! The model seems to output the horizontally flipped orientation (comparing with the figure in the paper), so I changed the code like below, and now it's working as expected.

ori = 360 - torch.argmax(hoe_output[0]) * 5

Would you mind I upload it to the repo?

No, I don't mind. Please feel free to use my code.

Thank you very much for your help!