Vision_2018/thresh.py at master · IgniteRobotics/Vision_2018 · 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
118
119
120
121
122
import cv2
import pickle
import numpy as np

cap = cv2.VideoCapture(0)

lower_hsv = np.array([20, 71, 0])
upper_hsv = np.array([45, 255, 255])


def validate_contour(cnt, cntArea, rect):
    return True
    target = 1
    tolerance = 1

    (x, y, w, h) = rect

    rectArea = (w * h)

    ratio = rectArea / cntArea

    print(ratio)

    if abs(ratio - target) < tolerance:
        return True
    return False


def get_box(contours):
    result = []
    if len(contours) == 0:
        return None

    areaArray = []
    for i, c in enumerate(contours):
        area = cv2.contourArea(c)
        areaArray.append(area)

    sortedContours = sorted(zip(areaArray, contours), key=lambda x: x[0], reverse= True)

    hull = None
    for i in sortedContours:
        area = i[0]
        cnt = i[1]
        # epsilon = 0.02*cv2.arcLength(cnt, True)
        # hull = cv2.approxPolyDP(cnt, epsilon, True)

        hul = cv2.convexHull(cnt)
        x,y,w,h = cv2.boundingRect(hul)

        rect = (x,y,w,h)

        if validate_contour(hul, area, rect):
            hull = hul
            break

    M = cv2.moments(hull)
    try:
        x = int(M['m10']/M['m00'])
        y = int(M['m01']/M['m00'])
    except ZeroDivisionError:
        return None
    result = [hull, area, (x,y), rect]
    return result

features = []
labels = []

while True:
    _, OG_img = cap.read()
    OG_img = cv2.resize(OG_img, (320,240))

    img_width = 320
    img_height = 240

# img = cv2.GaussianBlur(OG_img, (5,5), 0)
    # img = cv2.medianBlur(img, 5)
    img = OG_img
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    thresh = cv2.inRange(hsv, lower_hsv, upper_hsv)
    kernal = np.ones((9,9), np.uint8)
    thresh = cv2.erode(thresh, kernal, iterations=1)
    thresh = cv2.medianBlur(thresh, 15)
    thresh = cv2.dilate(thresh, kernal, iterations=1)

    _, contours, heirarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)

    result = get_box(contours)
    if result is not None:
        [hull, area, center, rect] = result
        (x, y ,w, h ) = rect

        cv2.drawContours(OG_img, [hull], -1, (255, 0, 0), 3)
        cv2.rectangle(OG_img, (x, y), (x+w, y+h), (0, 0, 255), 2)

        radius = 2
        cv2.circle(OG_img, center, radius, (0, 0, 255), 2)
        center_x = -int(img_width/2) + center[0]

        crop_img = thresh[y:y+h, x:x+w]
        feature = [h, w, np.count_nonzero(thresh) / (h*w)]
        print(feature)
        features.append(feature)
        labels.append(1)
        cv2.imshow('crop', crop_img)

    cv2.imshow('result', OG_img)
    cv2.imshow('thresh', thresh)
    if cv2.waitKey(0) & 0xFF == ord('q'):
        break

f = open('training.pkl','wb')

pickle.dump(features,f)
print()
print('here is the data')

f = open('training_y_1.pkl','wb')

pickle.dump(labels,f)