monoloco/monstereo/predict.py

# pylint: disable=too-many-statements, too-many-branches, undefined-loop-variable

"""
Adapted from https://github.com/vita-epfl/openpifpaf/blob/master/openpifpaf/predict.py
"""

import os
import glob
import json
import logging
from collections import defaultdict


import torch
import PIL
import openpifpaf
import openpifpaf.datasets as datasets
from openpifpaf.predict import processor_factory, preprocess_factory
from openpifpaf import decoder, network, visualizer, show

from .visuals.printer import Printer
from .network import Loco
from .network.process import factory_for_gt, preprocess_pifpaf
from .activity import show_social

LOG = logging.getLogger(__name__)

OPENPIFPAF_PATH = 'data/models/shufflenetv2k30-201104-224654-cocokp-d75ed641.pkl'  # Default model


def factory_from_args(args):

    # Data
    if args.glob:
        args.images += glob.glob(args.glob)
    if not args.images:
        raise Exception("no image files given")

    # Model
    if not args.checkpoint:
        if os.path.exists(OPENPIFPAF_PATH):
            args.checkpoint = OPENPIFPAF_PATH
        else:
            LOG.info("Checkpoint for OpenPifPaf not specified and default model not found in 'data/models'. "
                  "Using a ShuffleNet backbone")
            args.checkpoint = 'shufflenetv2k30'

    # Devices
    args.device = torch.device('cpu')
    args.disable_cuda = False
    args.pin_memory = False
    if torch.cuda.is_available():
        args.device = torch.device('cuda')
        args.pin_memory = True
    args.loader_workers = 8

    # Add visualization defaults
    args.figure_width = 10
    args.dpi_factor = 1.0

    if args.net == 'monstereo':
        args.batch_size = 2
    else:
        args.batch_size = 1

    # Make default pifpaf argument
    args.force_complete_pose = True
    LOG.info("Force complete pose is active")

    # Configure
    decoder.configure(args)
    network.configure(args)
    show.configure(args)
    visualizer.configure(args)

    return args


def predict(args):

    cnt = 0
    args = factory_from_args(args)

    # Load Models
    assert args.net in ('monoloco_pp', 'monstereo', 'pifpaf')
    if args.net in ('monoloco_pp', 'monstereo'):
        net = Loco(model=args.model, net=args.net, device=args.device, n_dropout=args.n_dropout, p_dropout=args.dropout)

    # data
    processor, model = processor_factory(args)
    preprocess = preprocess_factory(args)

    # data
    data = datasets.ImageList(args.images, preprocess=preprocess)
    if args.net == 'monstereo':
        assert len(data.image_paths) % 2 == 0, "Odd number of images in a stereo setting"

    data_loader = torch.utils.data.DataLoader(
        data, batch_size=args.batch_size, shuffle=False,
        pin_memory=False, collate_fn=datasets.collate_images_anns_meta)

    for batch_i, (image_tensors_batch, _, meta_batch) in enumerate(data_loader):
        pred_batch = processor.batch(model, image_tensors_batch, device=args.device)

        # unbatch (only for MonStereo)
        for idx, (pred, meta) in enumerate(zip(pred_batch, meta_batch)):
            print('batch %d: %s', batch_i, meta['file_name'])
            pred = preprocess.annotations_inverse(pred, meta)

            if args.output_directory is None:
                splits = os.path.split(meta['file_name'])
                output_path = os.path.join(splits[0], 'out_' + splits[1])
            else:
                file_name = os.path.basename(meta['file_name'])
                output_path = os.path.join(args.output_directory, 'out_' + file_name)
            print('image', batch_i, meta['file_name'], output_path)

            if idx == 0:
                with open(meta_batch[0]['file_name'], 'rb') as f:
                    cpu_image = PIL.Image.open(f).convert('RGB')
                pifpaf_outs = {
                    'pred': pred,
                    'left': [ann.json_data() for ann in pred],
                    'image': cpu_image}
            else:
                pifpaf_outs['right'] = [ann.json_data() for ann in pred]

        # 3D Predictions
        if args.net in ('monoloco_pp', 'monstereo'):

            im_name = os.path.basename(meta['file_name'])
            im_size = (cpu_image.size[0], cpu_image.size[1])  # Original
            kk, dic_gt = factory_for_gt(im_size, name=im_name, path_gt=args.path_gt)

            # Preprocess pifpaf outputs and run monoloco
            boxes, keypoints = preprocess_pifpaf(pifpaf_outs['left'], im_size, enlarge_boxes=False)

            if args.net == 'monoloco_pp':
                LOG.info("Prediction with MonoLoco++")
                dic_out = net.forward(keypoints, kk)
                dic_out = net.post_process(dic_out, boxes, keypoints, kk, dic_gt)
                if args.social_distance:
                    dic_out = net.social_distance(dic_out, args)

            else:
                LOG.info("Prediction with MonStereo")
                boxes_r, keypoints_r = preprocess_pifpaf(pifpaf_outs['right'], im_size)
                dic_out = net.forward(keypoints, kk, keypoints_r=keypoints_r)
                dic_out = net.post_process(dic_out, boxes, keypoints, kk, dic_gt)

        else:
            dic_out = defaultdict(list)
            kk = None

        # Outputs
        factory_outputs(args, pifpaf_outs, dic_out, output_path, kk=kk)
        LOG.info('Image {}\n'.format(cnt) + '-' * 120)
        cnt += 1


def factory_outputs(args, pifpaf_outs, dic_out, output_path, kk=None):
    """Output json files or images according to the choice"""

    # Verify conflicting options
    if any((xx in args.output_types for xx in ['front', 'bird', 'multi'])):
        assert args.net != 'pifpaf', "please use pifpaf original arguments"
        if args.social_distance:
            assert args.net == 'monoloco_pp', "Social distancing only works with MonoLoco++ network"

    if args.net == 'pifpaf':
        annotation_painter = openpifpaf.show.AnnotationPainter()
        with openpifpaf.show.image_canvas(pifpaf_outs['image'], output_path) as ax:
            annotation_painter.annotations(ax, pifpaf_outs['pred'])

    elif any((xx in args.output_types for xx in ['front', 'bird', 'multi'])):
        LOG.info(output_path)
        if args.social_distance:
            show_social(args, pifpaf_outs['image'], output_path, pifpaf_outs['left'], dic_out)
        else:
            printer = Printer(pifpaf_outs['image'], output_path, kk, args)
            figures, axes = printer.factory_axes(dic_out)
            printer.draw(figures, axes, pifpaf_outs['image'])

    elif 'json' in args.output_types:
        with open(os.path.join(output_path + '.monoloco.json'), 'w') as ff:
            json.dump(dic_out, ff)

    else:
        LOG.info("No output saved, please select one among front, bird, multi, or pifpaf options")