diff --git a/capture.py b/capture.py index 2f346a9..e42e915 100644 --- a/capture.py +++ b/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 diff --git a/environment.py b/environment.py index 002c07d..d098543 100644 --- a/environment.py +++ b/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 diff --git a/test.py b/test.py index 9c9fa60..fcd075b 100644 --- a/test.py +++ b/test.py @@ -1,3 +1,5 @@ +import time + import keyboard import winsound from models import FishNet @@ -16,7 +18,7 @@ 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() @@ -24,6 +26,14 @@ 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) diff --git a/utils.py b/utils.py index d7c37c6..738e297 100644 --- a/utils.py +++ b/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)] \ No newline at end of file + 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)) \ No newline at end of file