RCMrobot-CV/loop_thread.py at main · TimSong412/RCMrobot-CV · 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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
# from PIL import Image
# from propose import BBProposer
import threading
from circlepropose import Circle
import cv2
import numpy as np
from matting import BcRd
import time
import network

DEBUG = True
lastimg = None
end = False
CAMIDX = 4


class cap(threading.Thread):
    def __init__(self):
        global lastimg
        threading.Thread.__init__(self)
        self.camera = cv2.VideoCapture(CAMIDX)
        ret, lastimg = self.camera.read()
        print("INIT")
        if not ret:
            print("READ IMG ERROR")

    def run(self):
        global lastimg
        while True:
            ret, lastimg = self.camera.read()
            if ret:
                time.sleep(0.003)
            else:
                print("ERROR")
            if end:
                break


class Loop():
    def __init__(self) -> None:
        global lastimg
        self.Predictor = Circle()
        self.preimg = lastimg.copy()
        if DEBUG:
            cv2.imshow("ini_img", self.preimg)
            cv2.waitKey(0)
        if True:
            self.pre_gray, self.pre_mask, self.pre_center = self.imgprocess(
                self.preimg)
            self.boarder = np.zeros(
                (self.preimg.shape[0], self.preimg.shape[1]), dtype=np.uint8)
            self.boarder[10:self.preimg.shape[0] -
                         10, 10:self.preimg.shape[1]-10] += 1
        else:
            print("Error! can not open camera!")
            exit()

    def imgprocess(self, img):
        bb = self.Predictor.predict(img)
        bbox = np.array(bb[0][0], dtype=int)
        bg = BcRd(img)
        mask = bg.GenMask(bbox)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        totalmask = np.zeros(gray.shape[0:2], dtype=np.uint8)
        # totalmask[bbox[0]:bbox[2], bbox[1]:bbox[3]] = mask
        totalmask[bbox[0]:bbox[2], bbox[1]:bbox[3]] = 1
        totalmask[int((bbox[0]+bbox[2])/2), int((bbox[1]+bbox[3])/2)] = 0
        if DEBUG:
            cv2.imshow("mask", totalmask*250)
        # (y, x)
        return gray, totalmask, np.array([(bbox[1]+bbox[3])/2, (bbox[0]+bbox[2])/2])

    def refresh(self, newimg, newgray, newmask):
        self.preimg = newimg.copy()
        self.pre_gray = newgray.copy()
        self.pre_mask = newmask.copy()

    def norm(self, vec):
        if abs(vec[0]) > 0.01 or abs(vec[1]) > 0.01:
            deno = max(abs(vec[0]), abs(vec[1]))*100
        else:
            deno = 1
        return vec/deno


if __name__ == '__main__':
    sender = network.Sender('192.168.137.187')
    cam = cap()
    cam.start()
    mainloop = Loop()
    while True:
        st = time.time()
        newimg = lastimg.copy()
        ed = time.time()
        print("cam= ", ed-st)

        newgray, newmask, newcenter = mainloop.imgprocess(newimg)
        if DEBUG:
            overlap = np.array(mainloop.pre_gray*0.3 +
                               newgray*0.7, dtype=np.uint8)
            cv2.imshow("overlap", overlap)

        flow = cv2.calcOpticalFlowFarneback(
            mainloop.pre_gray, newgray, None, 0.5, 3, 15, 3, 5, 1.2, 0)
        flow *= np.expand_dims(mainloop.pre_mask, 2).repeat(2, axis=2)
        flow *= np.expand_dims(mainloop.boarder, 2).repeat(2, axis=2)
        flowdata = flow.reshape(1, -1, 2)[0]
        shift = np.mean(flowdata, axis=0)  # (y_width, x_height)
        move = (newcenter - mainloop.pre_center)/10000
        shift /= 100
        print("shift= ", shift)
        print("move= ", move)

        # nomalization
        shift = mainloop.norm(shift)
        move = mainloop.norm(move)

        sender.sendvec(newcenter[0], newcenter[1])
        print("center= ", newcenter)
        '''
        if abs(shift[0]) < 0.004 and abs(shift[1]) < 0.004:
            if abs(move[0]) > 0.004 or abs(move[1]) > 0.004:
                print("MOVE!")
                # sender.sendvec(move[0], move[1])
                sender.sendvec(newcenter[0], newcenter[1])
                print("center= ", newcenter)
            # mainloop.refresh(newimg, newgray, newmask)
                pass
            else:
                print("CONVERGE!")
        else:
            print("SHIFT!")
            # sender.sendvec(shift[0], shift[1])
            sender.sendvec(newcenter[0], newcenter[1])
            print("center= ", newcenter)
            pass
        '''

        if cv2.waitKey(500) == 115:
            end = True
            break
        ed = time.time()
        print("t=", ed-st, "--------------------------")
    cam.join()