vector_trigger_pipeline.py

# TODO:(nhendy) script level docstring
import tempfile
import galleries
from detectors import FasterRCNN
from attribute_extractors import MgnWrapper
import loaders
from PIL import Image
from utils import crop_image
import config_parser
import argparse
import functools
import os
import sys
import numpy as np
import cv2
from third_party.sort import Sort
from bbox_trigger import VectorTrigger
from scipy.stats import mode
from scipy.optimize import linear_sum_assignment
from tqdm import tqdm
from constants import *
import collections
from out_writer import RESULT, write_result_file, OutWriter
from third_party.re_ranking_ranklist import re_ranking
import displayer


def init_args():
    parser = argparse.ArgumentParser(description="multi-camera re-id system")
    parser.add_argument("-d",
                        "--detector",
                        help="Object detection model",
                        default='FasterRCNN',
                        choices=['FasterRCNN'])
    parser.add_argument("-r",
                        "--distance",
                        default='dot_product',
                        help="Distance metric used for retrieval",
                        choices=['dot_product'])
    parser.add_argument("-l",
                        "--loader",
                        default='videos',
                        help="Type of data loading",
                        choices=['videos', 'frames'])
    parser.add_argument("-g",
                        "--gallery",
                        default='trigger',
                        help="Type of Gallery",
                        choices=['trigger'])
    parser.add_argument("-v",
                        "--vect_gen",
                        default='MGN',
                        help="Attribute extraction model",
                        choices=['MGN'])
    parser.add_argument("-i",
                        "--interval",
                        default=2,
                        help="Sampling interval",
                        type=int)
    parser.add_argument("-nr",
                        "--no_rerank",
                        action="store_true",
                        help="Turns off reranking during assignment")
    parser.add_argument("--video_path",
                        required=True,
                        help="Path to the video to run the pipeline on")
    parser.add_argument(
        "--weights_path",
        required=True,
        help="Path to MGN weigths",
    )
    parser.add_argument(
        "--gallery_path",
        default="tmpgal/",
        help="Path to gallery",
    )
    parser.add_argument("--config",
                        required=True,
                        help="json file with configuration info")
    parser.add_argument("--debug", action="store_true",
                        help="Intermediate information saved for manual analyzing")
    parser.add_argument("--phase2", default=None, help="debug file")
    parser.add_argument("--make_videos", action="store_true",
                        help="will create output videos at end of analysis")
    parser.add_argument("--email", action="store_true",
                        help="email user on completion")
    return parser.parse_args()

def read_img_and_compute_feat_vector(path, attribute_extractor):
    if os.path.isfile(path):
        img = Image.open(path)
        return attribute_extractor(img)
    return None


def get_max_index(array, k=1):
    maxnums = array[0:k]
    maxinds = np.arange(0, k)
    minins = min(maxnums)
    mininin = maxnums.index(minins)
    for i, val in enumerate(array):
        if val > minins:
            maxnums[mininin] = val
            maxinds[mininin] = i
            minins = min(maxnums)
            mininin = maxnums.index(minins)
    if len(maxnums) == 1:
        return maxinds[0]
    else:
        return [x for _, x in sorted(zip(maxnums, maxinds), reverse=True)]


def load_gallery_feat_vectors(imgs_dir, attribute_extractor):
    if not os.path.exists(imgs_dir):
        raise ValueError("path doesn't exist")
    feat_vectors = list()
    for name in os.listdir(imgs_dir):
        feat_vector = read_img_and_compute_feat_vector(
            os.path.join(imgs_dir, name), attribute_extractor)
        feat_vectors.append(feat_vector)
    return feat_vectors


def cosine_similarity(x, y):
    return np.dot(x, y) / (np.sqrt(np.dot(x, x)) * np.sqrt(np.dot(y, y)))


def write_gallery_imgs(imgs, path):
    if not os.path.exists(path):
        os.makedirs(path)
    for i, img in enumerate(imgs):
        cv2.imwrite(os.path.join(path, "{:5d}.jpg".format(i)), img)


def run_mot_and_fill_gallery(video_loader, gallery, detector, sort_trackers,
                             outwriter, ref_img_dict):

    lineid = 0
    # Iterate through frames of all cameras
    iterator = tqdm(video_loader)
    for findex, frames in iterator:
        """ early termination for gallery testing"
        if findex >= 400:
            iterator.close()
            break
        """

        # Iterate through each camera
        for vidname, frame in frames.items():
            # Get real frame index for this video
            vid_findex = video_loader.get_vid_frame_num(vidname, findex)

            # Get bounding boxes of all people
            boxes, scores = detector.get_bboxes(frame)

            # Send people bounding boxes to tracker
            # Get three things: normal Sort output (tracking bounding boxes it wants to send), corresponding track objects, and objects of new tracks
            tracker = sort_trackers[vidname]
            dets = np.column_stack((np.reshape(boxes, [-1, 4]), scores))
            matched_tracks, matched_kb_trackers, new_kb_trackers = tracker.update(
                dets)

            gallery.update(vidname, frame, matched_tracks)

            # Find indexes of returned bounding boxes that meet ideal ratio
            trkbboxes = np.array(matched_tracks)
            widths = trkbboxes[:, 2] - trkbboxes[:, 0]
            heights = trkbboxes[:, 3] - trkbboxes[:, 1]
            aspectratio = heights / widths
            readybools = np.isclose(aspectratio, 2, rtol=0.25)
            indexes = np.arange(len(matched_tracks))[readybools]

            # Iterate through returned bounding boxes
            for ind, trk in enumerate(matched_tracks):
                lineid += 1

                box = ((int(trk[0]), int(trk[1])), (int(trk[2]), int(trk[3])))
                use4id = 0

                # If bounding box meets ideal ratio, save image of person as reference
                if ind in indexes:
                    cropimg = crop_image(frame, box)
                    if cropimg.size > 5:
                        use4id = 1
                        temp_file_name = outwriter.save_reid_ref_img(cropimg)
                        matched_kb_trackers[ind].save_img(temp_file_name)
                        ref_img_dict[lineid] = temp_file_name

                # Write frame number, id, trackid, ideal ratio flag, bounding box, and lineid to file
                outwriter.write_result_file(vidname, vid_findex, findex, -1, trk[4], use4id, box, lineid)

            # Iterate through new tracks and add their current bounding box to list of track references
            for trk in new_kb_trackers:
                lineid += 1
                d = trk.get_state()[0]
                box = ((int(d[0]), int(d[1])), (int(d[2]), int(d[3])))
                cropimg = crop_image(frame, box)
                if cropimg.size > 5:
                    temp_file_name = outwriter.save_reid_ref_img(cropimg)
                    trk.save_img(temp_file_name)
                    ref_img_dict[lineid] = temp_file_name
                    outwriter.write_result_file(vidname, vid_findex, findex, -1, trk.id, 1, box, lineid)

def create_cost_matrix(feats1, feats2):
    cost_matrix = np.zeros((len(feats1), len(feats2)))
    for i, vect1 in enumerate(feats1):
        for j, vect2 in enumerate(feats2):
            cost_matrix[i,j] = 1 - np.average(np.dot(vect1, np.transpose(vect2)))
    return cost_matrix

def create_square_cost_matrix(feats):
    cost_matrix = np.zeros((len(feats), len(feats)))
    for i, vect in enumerate(feats):
        for j in range(i, len(feats)):
            cost_matrix[i,j] = 1 - np.average(np.dot(feats[i], np.transpose(feats[j])))
    cost_matrix = np.maximum(cost_matrix, cost_matrix.transpose())
    return cost_matrix

def run_reid_model_and_assign_ids(attribute_extractor,
                                  output_files, gallery_feature_vectors, ref_img_dict, args):

    # create gallery cost matrix
    g_g_dist = create_square_cost_matrix(gallery_feature_vectors)

    # get a list of [sync] frames
    sframes = [viddata[:, RESULT.sindex] for viddata in output_files.values()]
    sframes = np.unique(np.concatenate(sframes))

    # array of vid names, index is number used in arrays below
    vidname_array = list(output_files.keys())
    vidname_num = {vidname: num for num, vidname in enumerate(vidname_array)}

    prev_track_ref = collections.defaultdict(dict)
    convertdict = collections.defaultdict(lambda: collections.defaultdict(list))

    # iterate through sync frames
    for sframe in tqdm(sframes):
        # count number of bounding boxes across videos for this sync frame
        num_entries = sum(np.count_nonzero(viddata[:, RESULT.sindex] == sframe) for viddata in output_files.values())

        assignment_index = 0
        assignments = np.zeros((num_entries, 4))
        cost_matrices = dict()
        for vidname, viddata in output_files.items():
            frameinfo = viddata[viddata[:, RESULT.sindex] == sframe]
            if len(frameinfo) == 0: continue
            cost_matrix = list()
            for info in frameinfo:
                # update prev_track_ref with new representations of track
                if info[RESULT.ref_flag] == 1:
                    prev_track_ref[vidname][info[RESULT.trackid]] = read_img_and_compute_feat_vector(ref_img_dict[info[RESULT.lineid]], attribute_extractor)
                 
                # construct cost matrix for Hungarian Algorithm
                uniqvect = prev_track_ref[vidname][info[RESULT.trackid]] 

                # TODO: (msrasheed) move distance calc within Gallery class
                cost_matrix.append([
                    1 - np.average(np.dot(uniqvect, np.transpose(feat_vector)))
                    for feat_vector in gallery_feature_vectors
                ])

            cost_matrix = np.array(cost_matrix)

            # create query cost matrix
            q_q_dist = create_square_cost_matrix([prev_track_ref[vidname][info[RESULT.trackid]] for info in frameinfo])

            # get new cost_matrix from re ranking
            if not args.no_rerank:
                cost_matrix = re_ranking(cost_matrix, q_q_dist, g_g_dist)

            # get single id to track assignments using Hungarian Algorithm
            track_ind, id_val = linear_sum_assignment(cost_matrix)

            # put assignments for one camera in assignments array to compare later
            for tind, reid in zip(track_ind, id_val):
                vidnum = vidname_num[vidname]
                tid = frameinfo[tind, RESULT.trackid]
                conf = cost_matrix[tind, reid]
                assignments[assignment_index, :] = np.array([vidnum, tid, reid, conf])
                assignment_index += 1

        # get list of ids
        frames_reids = np.unique(assignments[:, 2])

        # iterate through ids and compare across cameras
        for reid in frames_reids:
            idinfo = assignments[assignments[:, 2] == reid]
            sel = np.argmin(idinfo[:, 3])
            convertdict[vidname_array[int(idinfo[sel, 0])]][int(idinfo[sel, 1])].append(reid)

    print(convertdict)

    # find mode of tracks and write files
    for vidname, viddata in output_files.items():
        vid_convertdict = {key: mode(val)[0][0] for key, val in convertdict[vidname].items()
                            if not len(val) == 0}

        # Vectorize mapping from track id to reid id.
        track_id_to_reid_id = np.vectorize(
            lambda value: vid_convertdict.get(int(value), -1))

        # Use vectorized mapping function to change trackIDs to Re-IDs in output file
        viddata[:, RESULT.reid] = track_id_to_reid_id(viddata[:, RESULT.trackid])

        # Save
        np.savetxt(vidname + ".txt", viddata)


def convert_files_to_numpy(temp_dir, output_files):
    for video_name, file_handle in output_files.items():
        file_handle.close()
        output_files[video_name] = np.loadtxt(os.path.join(
            temp_dir, "{}.txt".format(video_name)),
                                              delimiter=',')


def main():
    args = init_args()

    # TODO: (msrasheed) more control over type of triggering mechanism
    # needed in config file
    json_dict = config_parser.parse_json_filename(args.config)

    # Input Checks
    #  check if gallery directory is empty
    if os.path.exists(args.gallery_path) and len(os.listdir(args.gallery_path)) != 0:
        response = input("Gallery path {} is nonempty. This may affect results. Would you like to proceed? [y/n]: ".format(args.gallery_path)).lower()
        if response != 'y':
            print("aborting...")
            exit()

    dataloader = loaders.get_loader(args.video_path, args.loader,
                                    args.interval, json_dict)
    outwriter = OutWriter(dataloader.get_vid_names(), args.debug, "debug", args.phase2 != None)
    attribute_extractor = MgnWrapper(args.weights_path)
    if not args.phase2:

        dir_files = os.listdir()
        outfile_names = ["{}.txt".format(vidname) for vidname in dataloader.get_vid_names()]
        for val in outfile_names:
            if val in dir_files:
                response = input("Output file {} already exists. Would you like to overwrite this file? [y/n]: ".format(val)).lower()
                if response != 'y':
                    print("aborting...")
                    exit()

        detector = FasterRCNN()

        # TODO: (msrasheed) more control over type of triggering mechanism
        # needed in config file
        json_dict = config_parser.parse_json_filename(args.config)
        trigger_causes = config_parser.extract_line_trigger_list(json_dict)

        gallery = galleries.TriggerLineGallery(attribute_extractor, trigger_causes)

        sort_trackers = {
            vidnames: Sort()
            for vidnames in dataloader.get_vid_names()
        }

        # Dict for finding ref image given lineid
        ref_img_dict = dict()

        # Run detector, Sort and fill up gallery
        run_mot_and_fill_gallery(dataloader, gallery, detector, sort_trackers,
                                 outwriter, ref_img_dict)

        # Save images from gallery captured throughout video
        write_gallery_imgs(gallery.people(), args.gallery_path)
        if args.debug:
            outwriter.save_ref_img_dict(ref_img_dict)
            outwriter.write_debug_file(args.gallery_path)
    else: 
        debug_json = config_parser.parse_json_filename(args.phase2)
        args.gallery_path = debug_json["gallery_dir"]
        ref_img_dict = outwriter.load_debug(debug_json)

    # Load up the gallery feature vectors
    gallery_feature_vectors = load_gallery_feat_vectors(
        args.gallery_path, attribute_extractor)

    # TODO: (nhendy) this is needed because downstream functions
    # assume the dict contain numpy arrays not files. Remove later
    output_files = outwriter.return_numpy()

    # Run reid model and map track ids to reid ids
    run_reid_model_and_assign_ids(attribute_extractor,
                                  output_files, gallery_feature_vectors, ref_img_dict, args)

    if args.make_videos:
        for vidname in dataloader.get_vid_names():
            displayer.create_vid(vidname,
                                 os.path.join(args.video_path, vidname),
                                 vidname+'.txt')

    if args.email:
        ostream = os.popen("whoami")
        ecnuname = ostream.read()[:-1]
        os.system("echo done | mail -s \"ee220clnx1\" \"{}@ecn.purdue.edu\"".format(ecnuname))

if __name__ == "__main__":
    sys.exit(main())