fish bite detection

capture.py

@@ -1,39 +1,10 @@
 import pyautogui
-import cv2
-import numpy as np
-import time
-
-def match_img(img, target):
-    h, w = target.shape[:2]
-    res = cv2.matchTemplate(img, target, cv2.TM_CCOEFF)
-    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
-    print(max_val)
-    return (*max_loc, max_loc[0] + w, max_loc[1] + h, max_loc[0] + w//2, max_loc[1] + h//2)
-
-#img = cv2.imread('imgs/a.png')[94:94 + 103, 712:712 + 496, :]
-img = cv2.imread('fish_bar_dataset/78.jpg')[2:34,:,:]
-t_l = cv2.imread('imgs/target_left.png')
-t_l_n = cv2.imread('imgs/target_l+m_2.png')
-t_r = cv2.imread('imgs/target_right.png')
-t_n = cv2.imread('imgs/target_now.png')
-
-start=time.time()
-#img2 = pyautogui.screenshot(region=[712, 94, 496, 103])
-bbox_l=match_img(img, t_l)
-bbox_l_n=match_img(img, t_l_n)
-bbox_r=match_img(img, t_r)
-bbox_n=match_img(img, t_n)
-cv2.rectangle(img, bbox_r[0:2], bbox_r[2:4], (0,255,0), 2)  # 画出矩形位置
-cv2.rectangle(img, bbox_l[0:2], bbox_l[2:4], (255,0,0), 2)  # 画出矩形位置
-cv2.rectangle(img, bbox_n[0:2], bbox_n[2:4], (0,0,255), 2)  # 画出矩形位置
-#cv2.rectangle(img, bbox_l_n[0:2], bbox_n[2:4], (0,255,255), 2)  # 画出矩形位置
-end=time.time()
-print(end-start)
-cv2.imshow('a',img)
-cv2.waitKey()
-
-'''
-#959,166 r=21 ring center
-img = pyautogui.screenshot(region=[712, 94, 496, 103])  # x,y,w,h
-img.save('bar.png')'''
+import keyboard
+
+i=0
+while True:
+    keyboard.wait('t')
+    img = pyautogui.screenshot()
+    img.save(f'img_tmp/{i}.png')
+    i+=1
 

environment.py

@@ -12,6 +12,8 @@ def __init__(self, delay=0.1, max_step=100, show_det=True):
         self.t_l = cv2.imread('imgs/target_left.png')
         self.t_r = cv2.imread('imgs/target_right.png')
         self.t_n = cv2.imread('imgs/target_now.png')
+        self.im_bar = cv2.imread('imgs/bar2.png')
+        self.bite = cv2.imread('imgs/bite.png', cv2.IMREAD_GRAYSCALE)
         self.std_color=np.array([192,255,255])
         self.r_ring=21
         self.delay=delay
@@ -22,7 +24,8 @@ def __init__(self, delay=0.1, max_step=100, show_det=True):
         self.add_vec=[0,2,0,2,0,2]
 
     def reset(self):
-        self.img=cap([712 - 10, 94, 496 + 20, 103])
+        self.y_start = self.find_bar()[0]
+        self.img=cap([712 - 10, self.y_start, 496 + 20, 103])
 
         self.fish_start=False
         self.zero_count=0
@@ -42,6 +45,21 @@ def do_action(self, action):
     def scale(self, x):
         return (x-5-10)/484
 
+    def find_bar(self):
+        img = cap(region=[700, 0, 520, 300])
+        bbox_bar = match_img(img, self.im_bar)
+        if self.show_det:
+            img=deepcopy(img)
+            cv2.rectangle(img, bbox_bar[:2], bbox_bar[2:4], (0, 0, 255), 1)  # 画出矩形位置
+            cv2.imwrite(f'./img_tmp/bar.jpg',img)
+        return bbox_bar[1]-9, bbox_bar
+
+    def is_bite(self):
+        img = cap(region=[1595, 955, 74, 74])
+        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        edge_output = cv2.Canny(gray, 50, 150)
+        return psnr(self.bite, edge_output)>10
+
     def get_state(self):
         bar_img=self.img[2:34,:,:]
         bbox_l = match_img(bar_img, self.t_l)
@@ -85,7 +103,7 @@ def step(self, action):
         self.do_action(action)
 
         time.sleep(self.delay-0.05)
-        self.img=cap([712 - 10, 94, 496 + 20, 103])
+        self.img=cap([712 - 10, self.y_start, 496 + 20, 103])
         self.step_count+=1
 
         score=self.get_score()
@@ -97,7 +115,7 @@ def step(self, action):
         self.reward=score-self.last_score
         self.last_score=score
 
-        return self.get_state(), self.reward, (self.step_count>self.max_step or (self.zero_count>=15 and self.fish_start) or score>177)
+        return self.get_state(), self.reward, (self.step_count>self.max_step or (self.zero_count>=15 and self.fish_start) or score>176)
 
     def render(self):
         pass

test.py

@@ -1,3 +1,5 @@
+import time
+
 import keyboard
 import winsound
 from models import FishNet
@@ -16,14 +18,22 @@
 if __name__ == '__main__':
 
     net = FishNet(in_ch=args.n_states, out_ch=args.n_actions)
-    env = Fishing(delay=0.1, max_step=10000, show_det=False)
+    env = Fishing(delay=0.1, max_step=10000, show_det=True)
 
     net.load_state_dict(torch.load(args.model_dir))
     net.eval()
 
     while True:
         winsound.Beep(500, 500)
         keyboard.wait('r')
+        while True:
+            if env.is_bite():
+                break
+            time.sleep(0.5)
+        winsound.Beep(700, 500)
+        env.drag()
+        time.sleep(1)
+
         state = env.reset()
         for i in range(10000):
             state = torch.FloatTensor(state).unsqueeze(0)

utils.py

@@ -16,4 +16,11 @@ def list_add(li, num):
     if isinstance(num, int) or isinstance(num, float):
         return [x+num for x in li]
     elif isinstance(num, list) or isinstance(num, tuple):
-        return [x+y for x,y in zip(li,num)]
+        return [x+y for x,y in zip(li,num)]
+
+def psnr(img1, img2):
+   mse = np.mean( (img1/255. - img2/255.) ** 2 )
+   if mse < 1.0e-10:
+      return 100
+   PIXEL_MAX = 1
+   return 20 * np.log10(PIXEL_MAX / np.sqrt(mse))