Hackathon/visionTest.py at master · jabel3141/Hackathon · 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
from cueballfinder2 import CueBallFinder
from yellowballfinder import YellowFinder
from redballfinder import RedFinder
from darkredballfinder import DarkRedFinder
from blueballfinder import BlueFinder
from greenballfinder import GreenFinder
from purpleballfinder import PurpleFinder
from distance import CornerDistance
import cv2
import numpy as np
import math

def main():
    # cue =  CueBallFinder();
    #
    #
    # output = cue.process(img)
    # # for contour in output:
    # #     (x,y,w,h) = cv2.boundingRect(contour)
    # #     cv2.rectangle(img, (x,y), (x+w,y+h), (255, 0, 0), 2)
    #
    # cv2.imshow("image", img)
    # cv2.waitKey(0)

    global coordinates

    pipeline = CueBallFinder()
    yellowBall = YellowFinder()
    redBall = RedFinder()
    darkredBall = DarkRedFinder()
    blueBall = BlueFinder()
    greenBall = GreenFinder()
    purpleBall = PurpleFinder()
    cornerDistance = CornerDistance()
    cap = cv2.VideoCapture(0)
    leftCorner =[]

    # while cap.isOpened():
    #     ret, frame = cap.read()
    #     if ret:
    #         corners = cornerDistance.process(frame)
    #         #print(corners)
    #         # cv2.imshow('frame', frame)
    #         cornerDist = math.sqrt(math.pow(corners[0][0]+ corners[1][0], 2)+math.pow(corners[0][1]+corners[1][1],2))
    #         if corners[0][0]<corners[1][0]:
    #             leftCorner = corners[0]
    #         else:
    #             leftCorner = corners[1]
    #
    #         break

    while(cap.isOpened()):
        ret, frame = cap.read()
        if(ret):
            # import pdb;pdb.set_trace()
            output= pipeline.process(frame)
            oneBall = yellowBall.process(frame)
            threeBall = redBall.process(frame)
            darkred = darkredBall.process(frame)
            blueball = blueBall.process(frame)
            greenball = greenBall.process(frame)
            purpleball = purpleBall.process(frame)
            cv2.imshow('frame',frame)
            #print (output)
            #print(oneBall)
            #print(threeBall)
            # if(len(output) != 0):
            #     coordinates = [output[0]-leftCorner[0],output[1]-leftCorner[1]]
            # #print (coordinates)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()