add object segmentation

trex/main.py

@@ -1,13 +1,15 @@
 import webbrowser
 from mss import mss
 import numpy as np
-from time import sleep
+from time import sleep, time
 import pyautogui as gui
+import sched
 import cv2
 
 
 cv2.namedWindow('View', cv2.WINDOW_NORMAL)
 DEFAULT_MONITOR = 3
+scheduler = sched.scheduler(time, sleep)
 
 
 def take_screenshot(sct, monitor=DEFAULT_MONITOR):
@@ -21,9 +23,16 @@ def take_binaryScreenshot(sct, monitor: dict):
     binaryScreen = np.zeros(image.shape[:-1])
     binaryScreen[np.all(image == (83, 83, 83, 0), axis=2)] = 1
     binaryScreen = cv2.erode(binaryScreen, None, iterations=1)
+    binaryScreen = cv2.dilate(binaryScreen, None, iterations=1)
     return binaryScreen
 
 
+def binary_to_RGB(image):
+    rgb_image = np.zeros((*image.shape, 3))
+    rgb_image[image == 1] = (255, 255, 255)
+    return rgb_image
+
+
 def open_dino(sct):
     prevScreen = take_screenshot(sct)
     webbrowser.open('https://chromedino.com/')
@@ -59,8 +68,16 @@ def get_playing_bbox(sct):
     x, y = np.where(binScreen == 1)
     mw = int((y.max()-y.min()) / 100)
     mh = int((x.max()-x.min()) / 100)
-    playing_bbox = {'left':y.min()-2*mw, 'top':x.min()-5*mh, 'width':(y.max()-y.min())+4*mw, 'height':(x.max()-x.min())+10*mh}
-    return playing_bbox
+    playing_bbox = {'left':y.min()-2*mw, 'top':x.min()+15*mh, 'width':(y.max()-y.min())+4*mw, 'height':(x.max()-x.min())+10*mh}
+    return playing_bbox, mw, mh
+
+
+def to_bbox(bin_image):
+    labels, labeledScreen = cv2.connectedComponents(bin_image.astype('uint8'))
+    for label in range(1, labels):
+        x,y = np.where(labeledScreen==label)
+        bin_image[x.min():x.max()+1, y.min():y.max()+1] = 1
+    return bin_image
 
 
 
@@ -72,11 +89,30 @@ if __name__ == '__main__':
     open_dino(sct)
     gui.hotkey('fn', 'f11')
 
-    playing_bbox = get_playing_bbox(sct)
+    playing_bbox, mw, mh = get_playing_bbox(sct)
+    dino_pos = None
+    awaitable = None
+
+    binScreen = take_binaryScreenshot(sct, playing_bbox)
+    binScreen = to_bbox(binScreen)
+    labels, labeledScreen = cv2.connectedComponents(binScreen.astype('uint8'))
+    rgbScreen = binary_to_RGB(binScreen)
+    dino = min(range(1,labels), key=lambda x: np.where(labeledScreen==x)[1].min())
+    dino_pos = np.where(labeledScreen == dino)[1].max()
+    rgbScreen[labeledScreen==dino] = (0, 255, 0)
+
 
     while cv2.waitKey(1) != ord('q'):
         binScreen = take_binaryScreenshot(sct, playing_bbox)
-        cv2.imshow('View', binScreen)
+        binScreen = to_bbox(binScreen)
+        labels, labeledScreen = cv2.connectedComponents(binScreen.astype('uint8'))
+        rgbScreen = binary_to_RGB(binScreen)
+        for label in range(1, labels):
+            pos = np.where(labeledScreen == label)
+            if pos[1].max() > dino_pos and (pos[1].max()-pos[1].min()) > mw:
+                rgbScreen[labeledScreen == label] = (0, 0, 255)
+            else: rgbScreen[labeledScreen == label] = (0, 255, 0)
+        cv2.imshow('View', rgbScreen)
         cv2.waitKey(1)
     
     cv2.destroyAllWindows()