From 80e99220ef90434dfc1678785cab940c34b391d8 Mon Sep 17 00:00:00 2001
From: Geniukx <46901991+Geniukx@users.noreply.github.com>
Date: Sun, 10 Sep 2023 22:44:29 +0800
Subject: [PATCH] detector for enemy and item (#63)

* detector for enemy and item

* Add: detector for enemy and items
---
 tasks/combat/detector.py | 120 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 120 insertions(+)
 create mode 100644 tasks/combat/detector.py

diff --git a/tasks/combat/detector.py b/tasks/combat/detector.py
new file mode 100644
index 000000000..a1ad3bf8f
--- /dev/null
+++ b/tasks/combat/detector.py
@@ -0,0 +1,120 @@
+import os
+import cv2
+import numpy as np
+
+
+class Detector:
+    def __init__(self):
+        """
+        A detector to detect the circle mark of enemy and item.
+        """
+        self.ui_mask = cv2.imread(os.path.join(os.path.dirname(__file__), "mask.png"), 0)
+
+    def update(self, frame):
+        # apply a mask to every frame to block the UI from interfering detection
+        frame_masked = cv2.bitwise_and(frame, frame, mask=self.ui_mask)
+        # update the detector with current frame before detection
+        self.hsv = cv2.cvtColor(frame_masked, cv2.COLOR_BGR2HSV)
+
+    def detect_item(self):
+        filter_params = {
+            'lowerb': np.array([0, 0, 215]),
+            'upperb': np.array([108, 53, 255]),
+        }
+        hough_params = {
+            'minDist': 8,
+            'param1': 200,
+            'param2': 8,
+            'minRadius': 2,
+            'maxRadius': 4,
+        }
+
+        # augment the target marker through mask based on HSV color space
+        mask = cv2.inRange(self.hsv, **filter_params)
+        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2)), iterations=1)
+
+        # resize the mask to reduce computational complexity of hough operation
+        mask = cv2.resize(mask, None, fx=0.25, fy=0.25)
+        mask = cv2.GaussianBlur(mask, (3, 3), 0, 0)
+
+        circles = cv2.HoughCircles(mask, cv2.HOUGH_GRADIENT, 1, **hough_params)
+
+        # if any target exists, return the center coordinates
+        return self.get_coordinates(circles) if circles is not None else None
+
+    def detect_enemy(self):
+        filter_params = {
+            'lowerb': np.array([0, 62, 213]),
+            'upperb': np.array([180, 198, 255]),
+        }
+        hough_params = {
+            'minDist': 18,
+            'param1': 200,
+            'param2': 10,
+            'minRadius': 5,
+            'maxRadius': 9,
+        }
+
+        mask = cv2.inRange(self.hsv, **filter_params)
+        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2)), iterations=1)
+
+        mask = cv2.resize(mask, None, fx=0.25, fy=0.25)
+        mask = cv2.GaussianBlur(mask, (3, 3), 0, 0)
+        circles = cv2.HoughCircles(mask, cv2.HOUGH_GRADIENT, 1, **hough_params)
+
+        return self.get_coordinates(circles) if circles is not None else None
+
+    @staticmethod
+    def get_coordinates(circles):
+        # convert coordinates to original scale
+        circles = np.asarray(np.around(circles * 4), dtype=np.uint16).squeeze(0)
+        return [circle[:2] for circle in circles]
+
+
+def generate_ui_mask():
+    """
+    code to generate ui mask
+    """
+    mask = np.ones([720, 1280]) * 255
+    mask[34:81, 21:61] = 0
+    mask[179:220, 21:51] = 0
+    mask[35:84, 183:218] = 0
+    mask[0:61, 780:1280] = 0
+    mask[145:435, 1153:1240] = 0
+    cv2.circle(mask, (907, 614), 55, 0, -1)
+    cv2.circle(mask, (1033, 542), 67, 0, -1)
+    cv2.imwrite("mask.png", mask)
+
+
+if __name__ == '__main__':
+    # how to use
+    class YourClass:
+        def __init__(self, stream):
+            self.stream = stream
+            self.detector = Detector()  # initiate a detector, recommend to set scale=0.25
+
+            self.run()
+
+        def run(self):
+            while True:
+                # update the detector with current frame before detection
+                self.frame = cv2.cvtColor(self.stream.capture(), cv2.COLOR_RGB2BGR)
+                self.detector.update(self.frame)
+
+                # return a list of coordinates of corresponding target
+                enemies = self.detector.detect_enemy()
+                if enemies is not None:
+                    self.plot_points(enemies)
+
+                items = self.detector.detect_item()
+                if items is not None:
+                    self.plot_points(items)
+
+                cv2.imshow("frame", self.frame)
+                if cv2.waitKey(1) == ord('q'):
+                    cv2.destroyAllWindows()
+                    break
+
+        def plot_points(self, points):
+            for point in points:
+                cv2.circle(self.frame, (point[0], point[1]), 3, (255, 255, 255), -1)