Juartaurus/poi_Engineering

1from __future__ import division, print_function, absolute_import2 3import argparse4import os5 6import cv27import numpy as np8 9from application_util import preprocessing10from application_util import visualization11from deep_sort import nn_matching12from deep_sort.detection import Detection13from deep_sort.tracker import Tracker14 15 16def gather_sequence_info(sequence_dir, detection_file):17    """Gather sequence information, such as image filenames, detections,18    groundtruth (if available).19    Parameters20    ----------21    sequence_dir : str22        Path to the MOTChallenge sequence directory.23    detection_file : str24        Path to the detection file.25    Returns26    -------27    Dict28        A dictionary of the following sequence information:29        * sequence_name: Name of the sequence30        * image_filenames: A dictionary that maps frame indices to image31          filenames.32        * detections: A numpy array of detections in MOTChallenge format.33        * groundtruth: A numpy array of ground truth in MOTChallenge format.34        * image_size: Image size (height, width).35        * min_frame_idx: Index of the first frame.36        * max_frame_idx: Index of the last frame.37    """38    image_dir = os.path.join(sequence_dir, "img1")39    image_filenames = {40        int(os.path.splitext(f)[0]): os.path.join(image_dir, f)41        for f in os.listdir(image_dir)}42    groundtruth_file = os.path.join(sequence_dir, "gt/gt.txt")43 44    detections = None45    if detection_file is not None:46        detections = np.load(detection_file)47    groundtruth = None48    if os.path.exists(groundtruth_file):49        groundtruth = np.loadtxt(groundtruth_file, delimiter=',')50 51    if len(image_filenames) > 0:52        image = cv2.imread(next(iter(image_filenames.values())),53                           cv2.IMREAD_GRAYSCALE)54        image_size = image.shape55    else:56        image_size = None57 58    if len(image_filenames) > 0:59        min_frame_idx = min(image_filenames.keys())60        max_frame_idx = max(image_filenames.keys())61    else:62        min_frame_idx = int(detections[:, 0].min())63        max_frame_idx = int(detections[:, 0].max())64 65    info_filename = os.path.join(sequence_dir, "seqinfo.ini")66    if os.path.exists(info_filename):67        with open(info_filename, "r") as f:68            line_splits = [l.split('=') for l in f.read().splitlines()[1:]]69            info_dict = dict(70                s for s in line_splits if isinstance(s, list) and len(s) == 2)71 72        update_ms = 1000 / int(info_dict["frameRate"])73    else:74        update_ms = None75 76    feature_dim = detections.shape[1] - 10 if detections is not None else 077    seq_info = {78        "sequence_name": os.path.basename(sequence_dir),79        "image_filenames": image_filenames,80        "detections": detections,81        "groundtruth": groundtruth,82        "image_size": image_size,83        "min_frame_idx": min_frame_idx,84        "max_frame_idx": max_frame_idx,85        "feature_dim": feature_dim,86        "update_ms": update_ms87    }88    return seq_info89 90 91def create_detections(detection_mat, frame_idx, min_height=0):92    """Create detections for given frame index from the raw detection matrix.93    Parameters94    ----------95    detection_mat : ndarray96        Matrix of detections. The first 10 columns of the detection matrix are97        in the standard MOTChallenge detection format. In the remaining columns98        store the feature vector associated with each detection.99    frame_idx : int100        The frame index.101    min_height : Optional[int]102        A minimum detection bounding box height. Detections that are smaller103        than this value are disregarded.104    Returns105    -------106    List[tracker.Detection]107        Returns detection responses at given frame index.108    """109    frame_indices = detection_mat[:, 0].astype(np.int)110    mask = frame_indices == frame_idx111 112    detection_list = []113    for row in detection_mat[mask]:114        bbox, confidence, feature = row[2:6], row[6], row[10:]115        if bbox[3] < min_height:116            continue117        detection_list.append(Detection(bbox, confidence, feature))118    return detection_list119 120 121def run(sequence_dir, detection_file, output_file, min_confidence,122        nms_max_overlap, min_detection_height, max_cosine_distance,123        nn_budget, display):124    """Run multi-target tracker on a particular sequence.125    Parameters126    ----------127    sequence_dir : str128        Path to the MOTChallenge sequence directory.129    detection_file : str130        Path to the detections file.131    output_file : str132        Path to the tracking output file. This file will contain the tracking133        results on completion.134    min_confidence : float135        Detection confidence threshold. Disregard all detections that have136        a confidence lower than this value.137    nms_max_overlap: float138        Maximum detection overlap (non-maxima suppression threshold).139    min_detection_height : int140        Detection height threshold. Disregard all detections that have141        a height lower than this value.142    max_cosine_distance : float143        Gating threshold for cosine distance metric (object appearance).144    nn_budget : Optional[int]145        Maximum size of the appearance descriptor gallery. If None, no budget146        is enforced.147    display : bool148        If True, show visualization of intermediate tracking results.149    """150    seq_info = gather_sequence_info(sequence_dir, detection_file)151    metric = nn_matching.NearestNeighborDistanceMetric(152        "cosine", max_cosine_distance, nn_budget)153    tracker = Tracker(metric)154    results = []155 156    def frame_callback(vis, frame_idx):157        print("Processing frame %05d" % frame_idx)158 159        # Load image and generate detections.160        detections = create_detections(161            seq_info["detections"], frame_idx, min_detection_height)162        detections = [d for d in detections if d.confidence >= min_confidence]163 164        # Run non-maxima suppression.165        boxes = np.array([d.tlwh for d in detections])166        scores = np.array([d.confidence for d in detections])167        indices = preprocessing.non_max_suppression(168            boxes, nms_max_overlap, scores)169        detections = [detections[i] for i in indices]170 171        # Update tracker.172        tracker.predict()173        tracker.update(detections)174 175        # Update visualization.176        if display:177            image = cv2.imread(178                seq_info["image_filenames"][frame_idx], cv2.IMREAD_COLOR)179            vis.set_image(image.copy())180            vis.draw_detections(detections)181            vis.draw_trackers(tracker.tracks)182 183        # Store results.184        for track in tracker.tracks:185            if not track.is_confirmed() or track.time_since_update > 1:186                continue187            bbox = track.to_tlwh()188            results.append([189                frame_idx, track.track_id, bbox[0], bbox[1], bbox[2], bbox[3]])190 191    # Run tracker.192    if display:193        visualizer = visualization.Visualization(seq_info, update_ms=5)194    else:195        visualizer = visualization.NoVisualization(seq_info)196    visualizer.run(frame_callback)197 198    # Store results.199    f = open(output_file, 'w')200    for row in results:201        print('%d,%d,%.2f,%.2f,%.2f,%.2f,1,-1,-1,-1' % (202            row[0], row[1], row[2], row[3], row[4], row[5]),file=f)203 204 205def bool_string(input_string):206    if input_string not in {"True","False"}:207        raise ValueError("Please Enter a valid Ture/False choice")208    else:209        return (input_string == "True")210 211def parse_args():212    """ Parse command line arguments.213    """214    parser = argparse.ArgumentParser(description="Deep SORT")215    parser.add_argument(216        "--sequence_dir", help="Path to MOTChallenge sequence directory",217        default=None, required=True)218    parser.add_argument(219        "--detection_file", help="Path to custom detections.", default=None,220        required=True)221    parser.add_argument(222        "--output_file", help="Path to the tracking output file. This file will"223        " contain the tracking results on completion.",224        default="/tmp/hypotheses.txt")225    parser.add_argument(226        "--min_confidence", help="Detection confidence threshold. Disregard "227        "all detections that have a confidence lower than this value.",228        default=0.8, type=float)229    parser.add_argument(230        "--min_detection_height", help="Threshold on the detection bounding "231        "box height. Detections with height smaller than this value are "232        "disregarded", default=0, type=int)233    parser.add_argument(234        "--nms_max_overlap",  help="Non-maxima suppression threshold: Maximum "235        "detection overlap.", default=1.0, type=float)236    parser.add_argument(237        "--max_cosine_distance", help="Gating threshold for cosine distance "238        "metric (object appearance).", type=float, default=0.2)239    parser.add_argument(240        "--nn_budget", help="Maximum size of the appearance descriptors "241        "gallery. If None, no budget is enforced.", type=int, default=None)242    parser.add_argument(243        "--display", help="Show intermediate tracking results",244        default=True, type=bool_string)245    return parser.parse_args()246 247 248if __name__ == "__main__":249    args = parse_args()250    run(251        args.sequence_dir, args.detection_file, args.output_file,252        args.min_confidence, args.nms_max_overlap, args.min_detection_height,253        args.max_cosine_distance, args.nn_budget, args.display)