GeoTouch/polygon.py at main · mayankzach/GeoTouch · 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

import pickle
import cv2
import numpy as np

#############################
map_file_path = "D:\python programs\interactive_map\map.p"
countries_file_path = "D:\python programs\interactive_map\countries.p"
cam_id = 1
width, height = 1920,1080
#############################

cap = cv2.VideoCapture(cam_id)
cap.set(3, width)
cap.set(4, height)

file_obj = open(map_file_path, 'rb')
map_points = pickle.load(file_obj)
file_obj.close()
print(f"Loaded map coordinates.", map_points)

current_polygon = []
counter = 0

if countries_file_path:
    file_obj = open(countries_file_path, 'rb')
    polygons = pickle.load(file_obj)
    file_obj.close()
    print(f"Loaded {len(polygons)} countries.")
else:
    polygons = []


def warp_image(img, points, size=[1600, 800]):

    pts1 = np.float32(points)
    pts2 = np.float32([[0, 0], [size[0], 0], [0, size[1]], [size[0], size[1]]])
    matrix = cv2.getPerspectiveTransform(pts1, pts2)
    imgOutput = cv2.warpPerspective(img, matrix, (size[0], size[1]))
    return imgOutput, matrix


def mousePoints(event, x, y, flags, params):

    global counter, current_polygon

    if event == cv2.EVENT_LBUTTONDOWN:
        current_polygon.append((x, y))


while True:
    success, img = cap.read()
    imgWarped, _ = warp_image(img, map_points)

    key = cv2.waitKey(1)
    if key == ord("s") and len(current_polygon) > 2:
        country_name = input("Enter the Country name: ")
        polygons.append([current_polygon, country_name])
        current_polygon = []
        counter += 1
        print("Number of countries saved: ", len(polygons))


    if key == ord("q"):
        fileObj = open('countries.p', 'wb')
        pickle.dump(polygons, fileObj)
        fileObj.close()
        print(f"Saved {len(polygons)} countries")
        break

    if key == ord("d"):
        polygons.pop()

    if current_polygon:
        cv2.polylines(imgWarped, [np.array(current_polygon)], isClosed=True, color=(0, 0, 255), thickness=2)

    overlay = imgWarped.copy()

    for polygon, name in polygons:
        cv2.polylines(imgWarped, [np.array(polygon)], isClosed=True, color=(0, 255, 0), thickness=2)
        cv2.fillPoly(overlay, [np.array(polygon)], (0, 255, 0))

    cv2.addWeighted(overlay, 0.35, imgWarped, 0.65, 0, imgWarped)

    cv2.imshow("Warped Image", imgWarped)
    cv2.imshow("Original Image", img)
    cv2.setMouseCallback("Warped Image", mousePoints)