CellTracking/object_tracking.py at master · nivedhasivakumar/CellTracking · 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
import cv2
import copy
from detectors import Detectors
from tracker import Tracker


def main():
    # Create opencv video capture object
    # cap = cv2.VideoCapture('data/TrackingBugs.mp4')
    cap = cv2.VideoCapture('data/video_3_bin.mp4')

    # Create Object Detector
    detector = Detectors()

    # Create Object Tracker
    tracker = Tracker(160, 30, 5, 100)

    # Variables initialization
    skip_frame_count = 0
    track_colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0),
                    (0, 255, 255), (255, 0, 255), (255, 127, 255),
                    (127, 0, 255), (127, 0, 127)]
    pause = False

    frame_count = 1
    # Infinite loop to process video frames
    while True:
        # Capture frame-by-frame
        ret, frame = cap.read()

        # Make copy of original frame
        orig_frame = copy.copy(frame)

        # Convert binary image to greyscale
        frame[frame != 0] = 255

        # Skip initial frames that display logo
        if skip_frame_count < 1:
            skip_frame_count += 1
            continue

        # Detect and return centroids of the objects in the frame
        if ret:
            print "Processing frame " + format(frame_count)
            frame_count += 1
            centers = detector.Detect(frame)

            # If centroids are detected then track them
            if len(centers) > 0:

                # Track object using Kalman Filter
                tracker.Update(centers)

                # For identified object tracks draw tracking line
                # Use various colors to indicate different track_id
                for i in range(len(tracker.tracks)):
                    if len(tracker.tracks[i].trace) > 1:
                        for j in range(len(tracker.tracks[i].trace)-1):
                            # Draw trace line
                            x1 = tracker.tracks[i].trace[j][0][0]
                            y1 = tracker.tracks[i].trace[j][1][0]
                            x2 = tracker.tracks[i].trace[j+1][0][0]
                            y2 = tracker.tracks[i].trace[j+1][1][0]
                            clr = tracker.tracks[i].track_id % 9
                            cv2.line(frame, (int(x1), int(y1)), (int(x2), int(y2)),
                                     track_colors[clr], 2)

                # Display the resulting tracking frame
                cv2.imshow('Tracking', frame)

            # Display the original frame
            cv2.imshow('Original', orig_frame)

            # Slower the FPS
            cv2.waitKey(50)

            # # # Check for key strokes
            # k = cv2.waitKey(50) & 0xff
            # if k == 27:  # 'esc' key has been pressed, exit program.
            #     break
            # if k == 112:  # 'p' has been pressed. this will pause/resume the code.
            #     pause = not pause
            #     if pause is True:
            #         print "Code is paused. Press 'p' to resume.."
            #         while pause is True:
            #             # stay in this loop until
            #             key = cv2.waitKey(30) & 0xff
            #             if key == 112:
            #                 pause = False
            #                 print "Resume code..!!"
            #                 break
        else:
            print "All frames were processed"
            break

    # When everything done, release the capture
    cap.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    # execute main
    main()