TrackAtSea/tool.py at dev · NWalker4483/TrackAtSea · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
import argparse
import pickle
import cv2
import numpy as np
import csv

parser = argparse.ArgumentParser(description='Process some integers.')

parser.add_argument('--video_num', type=int, default=6)

parser.add_argument('--video_file', type=str,
                    help='Video file name (MP4 format)')
parser.add_argument('--num_detections', type=int, default=25,
                    help='')
parser.add_argument('--camera_file_path', type=str, default="generated_data/best_camera_params.pkl",
                    help='')
parser.add_argument('--cam_lat', type=float, default=32.70297,
                    help='Latitude of source in decimal degrees (i.e. where the camera is mounted')
parser.add_argument('--cam_long', type=float, default=-117.23463100000001,
                    help='Longitude of source in decimal degrees (i.e. where the camera is mounted')
parser.add_argument('--tracker_type', type=str, default="ORB",
                    help='')
parser.add_argument('--detection_file', type=str, default=None,
                    help='')


args = parser.parse_args()
if args.video_file == None:
    args.video_file = f"raw_data/video/{args.video_num}.mp4"

video = cv2.VideoCapture(args.video_file)
frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))

if args.tracker_type == "ORB":
    from trackers.ORB import ORBTracker
    tracker = ORBTracker(Land_at=420)
elif args.tracker_type == "MANUAL":
    from trackers.MANUAL import ManualTracker
    def drag(event, x, y, flags, param):
        global dragging, tracker # Defined in fit_transfomr.py drag function
        if dragging:
            tracker.pos = [x, y]
        if event == cv2.EVENT_LBUTTONDOWN:
            dragging = True
        # check to see if the left mouse button was released
        elif event == cv2.EVENT_LBUTTONUP:
            dragging = False
    cv2.namedWindow("Tracker Frame")
    cv2.setMouseCallback("Tracker Frame", drag)
    tracker = ManualTracker(frame_count, args.num_detections)
    dragging = False
# elif args.tracker_type == "SORT":
#     from trackers.SORT import SORTTracker
#     tracker = MVDATracker(init_frames=50, detecting_rate=1,
#                           detections_per_denoising=5, framerate=20, max_recovery_distance=50)
else:
    print("No tracker of the name exists")
    raise(KeyError)

detections = dict()
print("Detecting...")
while video.isOpened():
    ret, frame = video.read()
    if not ret:
        break
    det = tracker.update(frame)
    if len(det) > 0:
        for frame_num, ID, box in det:
            if frame_num not in detections: detections[frame_num] = dict()
            detections[frame_num][ID] = box

# Determine Landmark Track

print("Matching...")
from match_tracks import match_tracks
tracks = match_tracks(detections, args.video_file, thresh=10, scale=5)

# Load Homography and Camera Coefficients
params = pickle.load(open("generated_data/best_camera_params.pkl","rb"))

# Generate and save results
out_file = open(f"outputs/{args.video_num}.{args.tracker_type.lower()}.csv","w+")
fields = ['Frame No.', 'Vessel ID', 'Latitude', 'Longitude', 'X', 'Y']
csvwriter = csv.writer(out_file)
csvwriter.writerow(fields)

out_file2 = open(f"outputs/{args.video_num}.{args.tracker_type.lower()}.main.csv","w+")
mainwriter = csv.writer(out_file2)
mainwriter.writerow(fields)
try:
    print("Projecting...")
    for ID in tracks:
        distorted_points = []
        frames = []
        for frame_num in tracks[ID]:
            X1, Y1, X2, Y2 = tracks[ID][frame_num]
            X, Y = [X1+((X2 - X1)//2),Y2]
            frames.append(frame_num)
            distorted_points.append([[X],[Y]])
        if len(distorted_points) == 0: continue
        undistorted_points = cv2.undistortPoints(
                np.array(distorted_points, dtype=np.float64), params["Intrinsic Matrix"], params["Distortion Coefficients"], P=params["Intrinsic Matrix"])

        gps_projections = cv2.perspectiveTransform(undistorted_points, params["Homography"])
        gps_projections = gps_projections.reshape(-1, 2)

        det_count = len(frames)
        for i in range(det_count):
            if i % (det_count // args.num_detections) != 0: continue
            lat, lon = gps_projections[i]
            x, y = distorted_points[i][0][0], distorted_points[i][1][0]
            csvwriter.writerow([frames[i], ID, lat, lon, x, y])
            if ID == -1:
                mainwriter.writerow([frames[i], ID, lat, lon, x, y])
finally:
    out_file.close()
    out_file2.close()