benchmark.py

from __future__ import print_function, division
from builtins import str
from builtins import object
import pandas as pd
import glob
import PIL
import numpy as np
import hwang
import cv2
from collections import namedtuple, deque, OrderedDict
import torch
import tqdm as tq
import argparse
import maskrcnn_benchmark
from maskrcnn_benchmark.config import cfg
import sys
import os
import glob
import multiprocessing

import time
import tqdm

## need this in the pythonpath:
#     sys.path.append('/nvme_drive/orm/maskrcnn-benchmark/demo/')
### comes from maskrcnn-benchmark/demo folder. not part of package.
## make sure to have maskrcnn-benchmark/demo in the pythonpath
from predictor import COCODemo
#from vsms.bench_utils import *

import vsms.dataset
import json
import time
import pandas as pd
import numpy as np
from collections import namedtuple
import hwang
#from vsms.evaluation_tools import *
import signal,sys,time


dump_current_results = False
def signal_handling(signum,frame):
    global dump_current_results
    dump_current_results = True


def flatten_batch(tensor):  # batch of batches -> batch
    return tensor.view([-1] + list(tensor.shape[2:]))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Benchmark script to measure end to end performance of sampling etc.')

    parser.add_argument('--maskrcnn_config', dest='maskrcnn_config',
                        default= maskrcnn_benchmark.__path__[0] + '/../configs/caffe2/e2e_mask_rcnn_R_101_FPN_1x_caffe2.yaml',
                        help='config for detector')

    parser.add_argument('--object_class', dest='object_class',
                        default='person',
                        help='string for coco object class')

    parser.add_argument('--dataset_name', dest='dataset_name',
                        default='dashcam_unified_ng',
                        help='dataset name')

    parser.add_argument('--method', dest='method',
                        required=True,
                        help='method for search')

    parser.add_argument('--data_root', dest='DATA_ROOT',
                        required=True,
                        help='must have subdirs videos/ and csv/. If in noscope mode, '
                             'also needs cnn-avg/ cnn-models/, experiments/, as output by noscope')

    parser.add_argument('--num_workers', dest='num_workers',
                        default=5,
                        type=int,
                        help='num worker for data loaders',
                        )

    parser.add_argument('--test', dest='test',
                        default=False,
                        action='store_true',
                        help='run only a few iterations and see if everything works.')

    parser.add_argument('--noscope_accuracy_level', dest='noscope_accuracy_level',
                        default='0.1',
                        help='(string) used to pick a distilled cnn from among those trained')

    parser.add_argument('--batch_size', dest='batch_size',
                        type=int, default=8,
                        help='batch size')

    parser.add_argument('--skip_time', dest='skip_time',
                        type=str, default='1 s',
                        help='for sequential fixed rate benchmark, time interval in format readable by pandas timedelta')


    parser.add_argument('--time_limit', dest='time_limit',
                        type=str, default='2.5 hours',
                        help='time string for limit')

    parser.add_argument('--use_noscope_start', dest='use_noscope_start',
                        action='store_true', default=False,
                        help='if the chosen non-noscope sampler should use the noscope start point')

    parser.add_argument('--instance_limit', dest='instance_limit', default=1000,
                        type=int, help='quit this loop if this number is reached')

    parser.add_argument('--noscope_batch_size', dest='noscope_batch_size',
                        type=int, default=200,
                        help='batch size for diff loop')

    parser.add_argument('--decode_device', dest='decode_device', type=str, required=True)


    ## there is also a CSV generated by the
    args = parser.parse_args()
    print(vars(args))
    assert args.method in ['video_bandit', 'random', 'noscope', 'seq']\
        + ['random_read', 'noscope_read30fps', 'noscope_read', 'noscope_diff', 'noscope_cnn', 'noscope_all']

    pretty_ts = str(pd.to_datetime(int(time.time()), unit='s')).replace(' ', 'T').replace(':','-')
    odir = '{rt}/sampling_experiments/{ds}/{method}/{pretty_ts}/'.format(rt=args.DATA_ROOT, ds=args.dataset_name,
                                                                        pretty_ts=pretty_ts, method=args.method)

    dev_type = hwang.DeviceType.GPU if args.decode_device == 'GPU' else hwang.DeviceType.CPU
    g_time_limit = pd.to_timedelta(args.time_limit).total_seconds() # check this argument works

    assert not os.path.exists(odir) # no overwrite
    os.makedirs(odir)
    assert os.path.isdir(odir)

    dumpnum = 0
    def dump_results(dflist):
        start = time.time()
        global dumpnum
        global dump_current_results

        for (prefix, df) in dflist:
            opath = '{o}/{prefix}_{dumpnum}.parquet'.format(o=odir,
                                                    prefix=prefix,
                                                    dumpnum=dumpnum)

            df.to_parquet(opath)
            print('Dumped ', opath)

        dumpnum +=1
        dump_current_results = False
        end = time.time()
        print('dump took {:.2f}s'.format(end-start))


    print('testing output path...')
    argpth = '{o}/arg_vals.json'.format(o=odir)
    print('saving args to ', argpth)
    json.dump(vars(args), open(argpth, 'w'))

    ds = Dataset(data_root=args.DATA_ROOT,
                 dataset_name=args.dataset_name)


    maxscore = ds.max_scores
    assert maxscore.index.duplicated().sum() < 0.001 * maxscore.shape[0]
    match_ids = ds.gt

    start_frame = 0
    if args.use_noscope_start or args.method.startswith('noscope'):
        ns = Noscope(dataset=ds)
        start_frame = ns.params.start

    b1 = Bench(ds, start_frame=start_frame)
    skip_order = b1.gen_skip_ilocs(args.skip_time)

    print ('method: ', args.method)
    if args.method == 'seq':
        print(' ', skip_order.name)

    def init_maskrcnn():
        # update the config options with the config file
        h, w = (ds.video.h, ds.video.w)
        cfg.merge_from_file(args.maskrcnn_config)
        torch.cuda.set_device(0)
        cfg.merge_from_list(["MODEL.DEVICE", "cuda"])
        # torch.cuda.set_device(torch.device("cuda:0")) # this caused problems due to bug in fasterrcnn nms.

        coco_demo = COCODemo(
            cfg,
            min_image_size=min(h, w),  # keep original size?
            show_mask_heatmaps=False,
            confidence_threshold=0.0)
        return coco_demo

    maskrcnn = init_maskrcnn()
    text2label = OrderedDict([(cname, i) for (i, cname) in enumerate(maskrcnn.CATEGORIES)])
    kClass = text2label[args.object_class]
    print('Using scores for category=', maskrcnn.CATEGORIES[kClass])
    vds = dataset.VideoDataset(ds.video_path, transform=dataset.tocv2,
                               device_type=dev_type)


    class LoopStats(object):
        def __init__(self, chkpoint_fun, limit_str):
            self.start_time = time.time()
            self.last_chk = self.start_time
            self.frame_total = 0
            self.batch_total = 0
            self.body_total = 0
            self.body_start = 0
            self.chkpoint_fun = chkpoint_fun
            self.chkpoint_time = 3600 // 2
            self.time_limit = pd.to_timedelta(limit_str).total_seconds()

        def mark_body_start(self):
            global dump_current_results
            self.body_start = time.time()
            if dump_current_results \
                    or self.body_start - self.last_chk > self.chkpoint_time\
                    or self.batch_total == 1:
                self.chkpoint_fun()
                self.last_chk = self.body_start

            return (self.body_start - self.start_time > self.time_limit)

        def mark_body_end(self, num_frames):
            self.body_total += time.time() - self.body_start
            self.frame_total += num_frames
            self.batch_total += 1
            time_total = time.time() - self.start_time

            return [('loop body(%)', 100. *self.body_total / time_total),
                    ('avg_fps', self.frame_total / time_total),
                    ('fps(loop body)', self.frame_total/self.body_total),
                    ('batches', self.batch_total)]


    def run_maskrcnn(frmlist):
        return maskrcnn.max_scores(frmlist)[:, kClass].cpu().numpy()

    def run_sampling_loop(b1, method):
        sampler = b1.samplers[0]
        grader = b1.graders[0]

        def chkpoint():
            mres = get_results(b1)
            sres = sampler.get_sampler_history()
            dump_results([('grader', mres), ('hist', sres)])

        bds = dataset.BenchDataset(vds)
        loader = torch.utils.data.DataLoader(dataset=bds,
                                            sampler=sampler,
                                            batch_size = 1, # each sample is a batch
                                            num_workers = args.num_workers,
                                            pin_memory = False,
                                            drop_last = False,
                                            timeout = 0)

        assert sampler.batch_size == args.batch_size

        chkpoint() # testing
        stats = LoopStats(chkpoint, args.time_limit)
        #total = len(sampler) # if len(sampler) is not None else int(grader.NF)
        with tqdm.tqdm(unit='frames') as pbar:
            for (i,sbatch) in enumerate(loader):
                if stats.mark_body_start():
                    break
                frmlist = list(sbatch['frames'].squeeze().numpy())

                if method != 'random_read':
                    ids = sbatch['ids'].squeeze().numpy() if 'ids' in list(sbatch.keys()) else None
                    ilocs = sbatch['ilocs'].squeeze().numpy()
                    ref_scores = run_maskrcnn(frmlist)

                    # ref gt for some datasets is not maskrcnn
                    max_scores = grader.factory.max_scores.values[ilocs]
                    match_ids = grader.grade_batch(ilocs, max_scores)
                    #print(ilocs, ids, match_ids, max_scores, ref_scores)
                    sampler.feedback(SampleBatch(ilocs=ilocs, ids=ids), match_ids, ref_scores)

                st = stats.mark_body_end(len(frmlist))
                pbar.set_postfix(st + [('ni', grader.ni)], refresh=False)
                pbar.update(len(frmlist))


        print('Loop ended')
        chkpoint()


    def run_noscope_loop(b1, method):
        import tensorflow as tf

        if method == 'noscope_read30fps':
            # test throughput reading everything
            rng = np.arange(ns.params.start, ds.max_scores.shape[0], 1, dtype=np.int)
        else:
            rng = np.arange(ns.params.start, ds.max_scores.shape[0], ns.params.skip, dtype=np.int)

        base_ilocs = pd.Index(rng, name='noscope')
        cfg = named_prod(order=[base_ilocs],
                       score_method=['gt'],
                       batch_size=[args.noscope_batch_size],
                       split=[pd.CategoricalIndex(b1.nosplit)],
                       object_class=[args.object_class])

        add_combos(b1, cfg)
        grader = b1.graders[0]
        base_sampler = b1.samplers[0]
        ## tf model init
        with tf.device('/gpu:0'):
            tfconfig = tf.ConfigProto(log_device_placement=False)
            tfconfig.gpu_options.per_process_gpu_memory_fraction = 0.3
            sess = tf.Session(config=tfconfig)
            print("loading graph from:\n", ns.params.model)
            with tf.gfile.GFile(ns.params.model, 'rb') as f:
                graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
            sess.graph.as_default()
            tf.import_graph_def(graph_def, name='')
            input_img = sess.graph.get_operation_by_name('input_img').outputs[0]
            output_prob = sess.graph.get_operation_by_name('output_prob').outputs[0]

        Hist = namedtuple('Hist', 'frame,status,diff_confidence,cnn_confidence,yolo_confidence,maskrcnn_confidence')
        results = Hist(frame=deque([[]]),
                       status=deque([[]]),
                       diff_confidence=deque([[]]),
                       cnn_confidence=deque([[]]),
                       yolo_confidence=deque([[]]),
                       maskrcnn_confidence=deque([[]]))



        if method == 'noscope_diff':  # up to diff every time
            diff_upper = np.inf
            cnn_lower = np.inf
        elif method == 'noscope_cnn': # up to cnn every time
            diff_upper = -1.
            cnn_lower = np.inf
        elif method == 'noscope_all': # everything every time
            diff_upper = -1.
            cnn_lower  = -1.
        elif method == 'noscope':
            diff_upper = ns.params.diff_thresh
            cnn_lower = ns.params.distill_thresh_lower


        bds = dataset.BenchDataset(vds, transform_dict={})

        def cnn_normalize(cnn_ims):
            return cnn_ims.astype('float') / 255. - ns.avg_frm
        import torchvision.transforms as T
        diff_preproc = T.Compose([dataset.Resize(Noscope.kDiffRes),dataset.mse_diff])
        cnn_frms_preproc = T.Compose([dataset.Resize(Noscope.kDistRes),cnn_normalize])

        loader = torch.utils.data.DataLoader(dataset=bds,
                                            sampler=base_sampler,
                                            batch_size = 1,
                                            num_workers = args.num_workers,
                                            pin_memory = False,
                                            drop_last = False,
                                            timeout = 0)

        def chkpoint():
            cols = OrderedDict([(k, np.concatenate(v)) for (k, v) in list(results._asdict().items())])
            mres = pd.DataFrame(cols)
            mres2 = get_results(b1)
            dump_results([('hist', mres), ('grader', mres2)])


        chkpoint() # testing
        stats = LoopStats(chkpoint, args.time_limit)
        total = np.zeros(7)

        with tqdm.tqdm(unit='frames') as pbar:
            for (i, sbatch) in enumerate(loader):
                if stats.mark_body_start():
                    break

                frms = flatten_batch(sbatch['frames']).numpy()
                ilocs = flatten_batch(sbatch['ilocs']).numpy()
                assert ilocs.shape[0] == frms.shape[0]

                diff_scores = np.zeros(frms.shape[0])
                cnn_scores = np.zeros_like(diff_scores)
                yolo_scores = np.zeros_like(diff_scores)
                mask_scores = np.zeros_like(diff_scores)
                status = np.zeros_like(ilocs)

                if method in ['noscope_diff', 'noscope_cnn', 'noscope_all', 'noscope']:
                    status += 2
                    diff_scores = diff_preproc(frms)
                    cnn_mask = (diff_scores > diff_upper)

                    if cnn_mask.any():
                        status[cnn_mask] += 2
                        cnn_ims = cnn_frms_preproc(frms[cnn_mask])
                        cnn_res= sess.run(output_prob, feed_dict={input_img:cnn_ims})[:,1]
                        cnn_scores[cnn_mask] = cnn_res

                        yolo_mask = cnn_scores > cnn_lower

                        if yolo_mask.any():
                            status[yolo_mask] += 2
                            expensive_frames = list(frms[yolo_mask])
                            expensive_ilocs = ilocs[yolo_mask]

                            ## hack: for quality metric, use the scores used to train the model (bc
                            ## coral dataset was trained using yolo csv)
                            ys = grader.factory.max_scores.values[expensive_ilocs]

                            # we want the outer batch to be large, and this one needs to be smaller
                            # b/c tf model + torch model + large torch batch sometimes don't fit in gpu.
                            # TODO: allow specifying non-1 inner batch size
                            ms = []
                            for frm in expensive_frames:
                                ms.append(run_maskrcnn([frm]))

                            grader.grade_batch(expensive_ilocs, ys)
                            yolo_scores[yolo_mask] = ys
                            mask_scores[yolo_mask] = np.concatenate(ms)


                results.frame.append(ilocs)
                results.status.append(status)
                results.diff_confidence.append(diff_scores)
                results.cnn_confidence.append(cnn_scores)
                results.yolo_confidence.append(yolo_scores)
                results.maskrcnn_confidence.append(mask_scores)

                vals, counts = np.unique(status, return_counts=True)
                assert counts.sum() == frms.shape[0]

                total[vals] += counts
                assert total[6] == grader.nf, '{} != {}'.format(total[6], grader.nf)

                if i == 1:
                    chkpoint() # test

                st = stats.mark_body_end(frms.shape[0])
                assert total[[0,2,4,6]].sum() == stats.frame_total
                pct = total*100./total.sum()

                pbar.set_postfix(st +
                                 [
                                    ('ni', grader.ni),
                                    ('diffd_out(%)', pct[2]),
                                    ('cnnd_out(%)', pct[4]),
                                    ('maskrcnn(%)', pct[6]),
                                 ], refresh=False)

                pbar.update(frms.shape[0])

        print('Loop ended')
        chkpoint()


    configs = {
        'video_bandit':named_prod(order=[b1.lsb_ilocs],
                        score_method=['gt'],
                        batch_size=[args.batch_size],
                        split=[pd.CategoricalIndex(b1.split_30min)],
                        object_class=[args.object_class]),

        'random': named_prod(order=[b1.random_ilocs],
                        score_method=['gt'],
                        batch_size=[args.batch_size],
                        split=[pd.CategoricalIndex(b1.nosplit)],
                        object_class=[args.object_class]),

        'seq': named_prod(order=[skip_order],
                        score_method=['gt'],
                        batch_size=[args.batch_size],
                        split=[pd.CategoricalIndex(b1.nosplit)],
                        object_class=[args.object_class])
    }

    signal.signal(signal.SIGINT, signal_handling) # set statdump handler
    if args.method in ['noscope', 'noscope_read30fps', 'noscope_read', 'noscope_diff', 'noscope_cnn', 'noscope_all']:
        ## compare based on the part of the dataset that was not used for training
        run_noscope_loop(b1, args.method)
    else:
        if args.method == 'random_read':
            k = 'random'
        else:
            k = args.method
        cfg = configs[k]
        assert len(list(cfg)) == 1
        assert len(b1.samplers) == 0
        add_combos(b1, cfg)
        assert len(b1.samplers) == 1
        print('sampler params', b1.param_names[0])
        run_sampling_loop(b1, args.method)


## can check gpu activity over time with
# nvidia-smi -i 0 --query-gpu=timestamp,gpu_name,index,utilization.gpu,utilization.memory,memory.free,memory.used --format=csv,nounits -m 0.5 -lms 101