Amazing-Python-Scripts

1
import os
2
import numpy as np
3
import cv2
4
import pickle
5
import argparse
6
import matplotlib.pyplot as plt
7
from matplotlib.patches import Polygon
8
from matplotlib.widgets import PolygonSelector
9
from matplotlib.collections import PatchCollection
10
from shapely.geometry import box
11
from shapely.geometry import Polygon as shapely_poly
12

13
points = []
14
prev_points = []
15
patches = []
16
total_points = []
17
breaker = False
18

19

20
class SelectFromCollection(object):
21
    def __init__(self, ax):
22
        self.canvas = ax.figure.canvas
23

24
        self.poly = PolygonSelector(ax, self.onselect)
25
        self.ind = []
26

27
    def onselect(self, verts):
28
        global points
29
        points = verts
30
        self.canvas.draw_idle()
31

32
    def disconnect(self):
33
        self.poly.disconnect_events()
34
        self.canvas.draw_idle()
35

36

37
def break_loop(event):
38
    global breaker
39
    global globSelect
40
    global savePath
41
    if event.key == 'b':
42
        globSelect.disconnect()
43
        if os.path.exists(savePath):
44
            os.remove(savePath)
45

46
        print("Data saved in " + savePath + " file")
47
        with open(savePath, 'wb') as f:
48
            pickle.dump(total_points, f, protocol=pickle.HIGHEST_PROTOCOL)
49
        exit()
50

51

52
def onkeypress(event):
53
    global points, prev_points, total_points
54
    if event.key == 'n':
55
        pts = np.array(points, dtype=np.int32)
56
        if points != prev_points and len(set(points)) == 4:
57
            print("Points: " + str(pts))
58
            patches.append(Polygon(pts))
59
            total_points.append(pts)
60
            prev_points = points
61

62

63
def process_video(video_path, out_file):
64
    global globSelect
65
    global savePath
66
    savePath = out_file if out_file.endswith(".p") else out_file + ".p"
67

68
    print("\n> Select a region in the figure by enclosing it within a quadrilateral.")
69
    print("> Press the 'f' key to go full screen.")
70
    print("> Press the 'esc' key to discard the current quadrilateral.")
71
    print("> Try holding the 'shift' key to move all of the vertices.")
72
    print("> Try holding the 'ctrl' key to move a single vertex.")
73
    print("> After marking a quadrilateral, press 'n' to save the current quadrilateral, and then press 'q' to start marking a new quadrilateral.")
74
    print("> When you are done, press 'b' to exit the program.\n")
75

76
    video_capture = cv2.VideoCapture(video_path)
77
    cnt = 0
78
    rgb_image = None
79
    while video_capture.isOpened():
80
        success, frame = video_capture.read()
81
        if not success:
82
            break
83
        if cnt == 5:
84
            rgb_image = frame[:, :, ::-1]
85
        cnt += 1
86
    video_capture.release()
87

88
    fig, ax = plt.subplots()
89
    image = rgb_image
90
    ax.imshow(image)
91

92
    p = PatchCollection(patches, alpha=0.7)
93
    p.set_array(10 * np.ones(len(patches)))
94
    ax.add_collection(p)
95

96
    globSelect = SelectFromCollection(ax)
97
    bbox = plt.connect('key_press_event', onkeypress)
98
    break_event = plt.connect('key_press_event', break_loop)
99
    plt.show()
100
    globSelect.disconnect()
101

102

103
if __name__ == '__main__':
104
    parser = argparse.ArgumentParser()
105
    parser.add_argument('video_path', help="Path of the video file")
106
    parser.add_argument(
107
        '--out_file', help="Name of the output file", default="regions.p")
108
    args = parser.parse_args()
109

110
    video_path = args.video_path
111
    out_file = args.out_file
112

113
    process_video(video_path, out_file)
114