Improved clustering & cropping of Items (#247)

game.ini

@@ -2,7 +2,7 @@
 ; min and max hsv range (opencv format: h: [0-180], s: [0-255], v: [0, 255])
 ; h_min, s_min, v_min, h_max, s_max, v_max
 black=0,0,0,180,255,15
-item_highlight=90,235,135,115,255,160
+item_highlight=100,235,130,105,255,150
 white=0,0,150,180,20,255
 gray=0,0,90,180,20,130
 blue=114,100,190,125,132,255

src/item/item_cropper.py

@@ -0,0 +1,115 @@
+import cv2
+import numpy as np
+from config import Config
+from utils.misc import color_filter
+from dataclasses import dataclass
+import time
+
+
+# TODO: With OCR we can then add a "text" field to this class
+@dataclass
+class ItemText:
+    color_key: str = None
+    roi: list[int] = None
+    data: np.ndarray = None
+
+class ItemCropper:
+    def __init__(self):
+        self._config = Config()
+
+        self._gaus_filter = (19, 1)
+        self._expected_height_range = [int(round(num, 0)) for num in [x / 1.5 for x in [14, 40]]]
+        self._expected_width_range = [int(round(num, 0)) for num in [x / 1.5 for x in [60, 1280]]]
+
+        self._hud_mask = cv2.imread(f"assets/hud_mask.png", cv2.IMREAD_GRAYSCALE)
+        self._hud_mask = cv2.threshold(self._hud_mask, 1, 255, cv2.THRESH_BINARY)[1]
+
+        self._item_colors = ['white', 'gray', 'blue', 'green', 'yellow', 'gold', 'orange']
+
+    def clean_img(self, inp_img: np.ndarray) -> np.ndarray:
+        img = inp_img[:, :, :]
+        if img.shape[0] == self._hud_mask.shape[0] and img.shape[1] == self._hud_mask.shape[1]:
+            img = cv2.bitwise_and(img, img, mask=self._hud_mask)
+        # In order to not filter out highlighted items, change their color to black
+        highlight_mask = color_filter(img, self._config.colors["item_highlight"])[0]
+        img[highlight_mask > 0] = (0, 0, 0)
+        # Cleanup image with erosion image as marker with morphological reconstruction
+        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        thresh = cv2.threshold(gray, 15, 255, cv2.THRESH_BINARY)[1]
+        kernel = np.ones((3, 3), np.uint8)
+        marker = thresh.copy()
+        marker[1:-1, 1:-1] = 0
+        while True:
+            tmp = marker.copy()
+            marker = cv2.dilate(marker, kernel)
+            marker = cv2.min(thresh, marker)
+            difference = cv2.subtract(marker, tmp)
+            if cv2.countNonZero(difference) <= 0:
+                break
+        mask_r = cv2.bitwise_not(marker)
+        mask_color_r = cv2.cvtColor(mask_r, cv2.COLOR_GRAY2BGR)
+        img = cv2.bitwise_and(img, mask_color_r)
+        return img
+
+    def crop(self, inp_img: np.ndarray, padding_y: int = 5) -> list[ItemText]:
+        start = time.time()
+        cleaned_img = self.clean_img(inp_img)
+        debug_str = f" | clean: {time.time() - start}"
+
+        # Cluster item names
+        start = time.time()
+        item_clusters = []
+        for key in self._item_colors:
+            _, filtered_img = color_filter(cleaned_img, self._config.colors[key])
+            filtered_img_gray = cv2.cvtColor(filtered_img, cv2.COLOR_BGR2GRAY)
+            blured_img = np.clip(cv2.GaussianBlur(filtered_img_gray, self._gaus_filter, cv2.BORDER_DEFAULT), 0, 255)
+            contours = cv2.findContours(blured_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+            contours = contours[0] if len(contours) == 2 else contours[1]
+            for count, cntr in enumerate(contours):
+                x, y, w, h = cv2.boundingRect(cntr)
+                expected_height = 1 if (self._expected_height_range[0] < h < self._expected_height_range[1]) else 0
+                # increase height a bit to make sure we have the full item name in the cluster
+                y = y - padding_y if y > padding_y else 0
+                h += padding_y * 2
+                cropped_item = filtered_img[y:y+h, x:x+w]
+                # save most likely item drop contours
+                avg = int(np.average(filtered_img_gray[y:y+h, x:x+w]))
+                contains_black = True if np.min(cropped_item) < 14 else False
+                expected_width = True if (self._expected_width_range[0] < w < self._expected_width_range[1]) else False
+                mostly_dark = True if 4 < avg < 25 else False
+                if contains_black and mostly_dark and expected_height and expected_width:
+                    # double-check item color
+                    color_averages=[]
+                    for key2 in self._item_colors:
+                        _, extracted_img = color_filter(cropped_item, self._config.colors[key2])
+                        extr_avg = np.average(cv2.cvtColor(extracted_img, cv2.COLOR_BGR2GRAY))
+                        color_averages.append(extr_avg)
+                    max_idx = color_averages.index(max(color_averages))
+                    if key == self._item_colors[max_idx]:
+                        item_clusters.append(ItemText(
+                            color_key=self._item_colors[max_idx],
+                            roi=[x, y, w, h],
+                            data=cropped_item
+                        ))
+        debug_str += f" | cluster: {time.time() - start}"
+        # print(debug_str)
+        return item_clusters
+
+
+if __name__ == "__main__":
+    import keyboard
+    import os
+    from screen import Screen
+
+    keyboard.add_hotkey('f12', lambda: os._exit(1))
+    cropper = ItemCropper()
+    screen = Screen(cropper._config.general["monitor"])
+
+    while 1:
+        img = screen.grab().copy()
+        res = cropper.crop(img)
+        for cluster in res:
+            x, y, w, h = cluster.roi
+            cv2.rectangle(img, (x, y), (x+w, y+h), (0, 255, 0), 1)
+        cv2.imshow("res", img)
+        cv2.waitKey(1)

src/item/item_finder.py

@@ -7,6 +7,7 @@ from dataclasses import dataclass
 import math
 from config import Config
 from utils.misc import color_filter, cut_roi
+from item.item_cropper import ItemCropper
 
 
 @dataclass
@@ -26,6 +27,7 @@ class Item:
 class ItemFinder:
     def __init__(self):
         config = Config()
+        self._item_cropper = ItemCropper()
         # color range for each type of item
         # hsv ranges in opencv h: [0-180], s: [0-255], v: [0, 255]
         self._template_color_ranges = {
@@ -37,16 +39,7 @@ class ItemFinder:
             "unique": [np.array([23, 80, 140]), np.array([23, 89, 216])],
             "runes": [np.array([21, 251, 190]), np.array([22, 255, 255])]
         }
-        self._game_color_ranges = {
-            "white": config.colors["white"],
-            "gray": config.colors["gray"],
-            "magic": config.colors["blue"],
-            "set": config.colors["green"],
-            "rare": config.colors["yellow"],
-            "unique": config.colors["gold"],
-            "runes": config.colors["orange"]
-        }
-        self._gaus_filter = (17, 5)
+
         self._folder_name = "items"
         self._min_score = 0.86
         # load all templates
@@ -79,37 +72,12 @@ class ItemFinder:
     def search(self, inp_img: np.ndarray) -> List[Item]:
         img = inp_img[:,:,:]
         start = time.time()
-        # Pre filter black and highlight
-        mask1, _ = color_filter(img, self._config.colors["black"])
-        mask2, _ = color_filter(img, self._config.colors["item_highlight"])
-        filtered_img = cv2.bitwise_or(mask1, mask2)
-        contours = cv2.findContours(filtered_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        contours = contours[0] if len(contours) == 2 else contours[1]
-        new_img = np.zeros(img.shape, np.uint8)
-        for cntr in contours:
-            x, y, w, h = cv2.boundingRect(cntr)
-            new_img[y:y+h, x:x+w] = img[y:y+h, x:x+w]
-        img = new_img
-        # Filter by item colors
-        filtered_img = np.zeros(img.shape, np.uint8)
-        for key in self._game_color_ranges:
-            _, extracted_img = color_filter(img, self._game_color_ranges[key])
-            filtered_img = cv2.bitwise_or(filtered_img, extracted_img)
-        filtered_img_gray = cv2.cvtColor(filtered_img, cv2.COLOR_BGR2GRAY)
-        # Cluster item names
-        cluster_img = np.clip(cv2.GaussianBlur(filtered_img_gray, self._gaus_filter, cv2.BORDER_DEFAULT), 0, 255)
-        contours = cv2.findContours(cluster_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        contours = contours[0] if len(contours) == 2 else contours[1]
+        item_text_clusters = self._item_cropper.crop(img, 7)
         item_list = []
-        for cntr in contours:
-            x, y, w, h = cv2.boundingRect(cntr)
-            x -= 5
-            y -= 5
-            w += 10
-            h += 10
+        for cluster in item_text_clusters:
+            x, y, w, h = cluster.roi
             # cv2.rectangle(inp_img, (x, y), (x+w, y+h), (0, 255, 0), 1)
-
-            cropped_input = filtered_img[y:y+h, x:x+w]
+            cropped_input  = cluster.data
             best_score = None
             item = None
             for key in self._templates:
@@ -132,10 +100,11 @@ class ItemFinder:
                                 if template.blacklist:
                                     item = None
                                 else:
-                                    max_loc = [max_loc[0] + x, max_loc[1] + y]
                                     # Do another color hist check with the actuall found item template
-                                    cropped_roi = [*max_loc, template.data.shape[1], template.data.shape[0]]
-                                    cropped_item = cut_roi(filtered_img, cropped_roi)
+                                    # TODO: After cropping the "cropped_input" with "cropped_item", check if "cropped_input" might need to be
+                                    #       checked for other items. This would solve the issue of many items in one line being in one cluster
+                                    roi = [max_loc[0], max_loc[1], template.data.shape[1], template.data.shape[0]]
+                                    cropped_item = cut_roi(cropped_input, roi)
                                     grayscale = cv2.cvtColor(cropped_item, cv2.COLOR_BGR2GRAY)
                                     _, mask = cv2.threshold(grayscale, 0, 255, cv2.THRESH_BINARY)
                                     hist = cv2.calcHist([cropped_item], [0, 1, 2], mask, [8, 8, 8], [0, 256, 0, 256, 0, 256])
@@ -143,10 +112,10 @@ class ItemFinder:
                                     same_type = hist_result > 0.65 and hist_result is not np.inf
                                     if same_type:
                                         item = Item()
-                                        item.center = (int(max_loc[0] + int(template.data.shape[1] * 0.5)), int(max_loc[1] + int(template.data.shape[0] * 0.5)))
+                                        item.center = (int(max_loc[0] + x + int(template.data.shape[1] * 0.5)), int(max_loc[1] + y + int(template.data.shape[0] * 0.5)))
                                         item.name = key
                                         item.score = max_val
-                                        item.roi = [*max_loc, template.data.shape[1], template.data.shape[0]]
+                                        item.roi = [max_loc[0] + x, max_loc[1] + y, template.data.shape[1], template.data.shape[0]]
                                         center_abs = (item.center[0] - (inp_img.shape[1] // 2), item.center[1] - (inp_img.shape[0] // 2))
                                         item.dist = math.dist(center_abs, (0, 0))
             if item is not None and self._config.items[item.name]:
@@ -168,10 +137,10 @@ if __name__ == "__main__":
         img = screen.grab().copy()
         item_list = item_finder.search(img)
         for item in item_list:
-            print(item.name + " " + str(item.score))
+            # print(item.name + " " + str(item.score))
             cv2.circle(img, item.center, 5, (255, 0, 255), thickness=3)
             cv2.rectangle(img, item.roi[:2], (item.roi[0] + item.roi[2], item.roi[1] + item.roi[3]), (0, 0, 255), 1)
-            cv2.putText(img, item.name, item.center, cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 1, cv2.LINE_AA)
+            # cv2.putText(img, item.name, item.center, cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 1, cv2.LINE_AA)
         # img = cv2.resize(img, None, fx=0.5, fy=0.5)
         cv2.imshow('test', img)
         cv2.waitKey(1)

src/utils/item_extractor.py

@@ -0,0 +1,62 @@
+"""
+Script to autocrop items. Input image with items in the correct resolution and the script will auto crop it for you and ask for names for each of them.
+"""
+import argparse
+import os
+import cv2
+import numpy as np
+from config import Config
+from utils.misc import color_filter
+from item.item_cropper import ItemCropper
+import time
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Script to autocrop items.")
+    parser.add_argument("--file_path", type=str, help="Path to screenshots e.g. C:/data")
+    args = parser.parse_args()
+
+    args.file_path = "C:\\Users\\aliig\\Desktop\\bot\\botty-gleed-ocr\\input_images"
+    gen_truth = 1
+
+
+    item_cropper = ItemCropper()
+
+    for filename in os.listdir(args.file_path):
+        if filename.endswith(".png"):
+            start = time.time()
+            inp_img = cv2.imread(f"{args.file_path}\\{filename}")
+            filename = filename[:-4]
+            img = inp_img[:,:,:]
+            img_clean = item_cropper.clean_img(img)
+            item_clusters = item_cropper.crop(img)
+            for count, cluster in enumerate(item_clusters):
+                x, y, w, h = cluster.roi
+                key = cluster.color_key
+                if gen_truth:
+                    cv2.namedWindow("item")
+                    cv2.moveWindow("item", 100, 100)
+                    cv2.imshow("item", img_clean[y:y+h, x:x+w])
+                    cv2.waitKey(1)
+                    print(f"{count} Input item name and press enter (converts to all caps)...")
+                    item_name = input()
+                    if item_name != "":
+                        out_filename = f"{key}_{item_name.replace(' ','_')}"
+                        if not os.path.exists(f"./ground_truth/{out_filename}.png"):
+                            cv2.imwrite(f"./ground_truth/{out_filename}.png", img_clean[y:y+h, x:x+w])
+                            file1 = open(f"./ground_truth/{out_filename}.gt.txt","w")
+                            file1.write(item_name.upper())
+                            file1.close()
+                    else:
+                        print("Skipping")
+                    time.sleep(0.1)
+                    cv2.destroyAllWindows()
+                else:
+                    avg = int(np.average(cv2.cvtColor(cluster.data, cv2.COLOR_BGR2GRAY)))
+                    cv2.imwrite(f"./generated/z_{filename}_{key}_{count}_{avg}.png", cluster.data)
+                    cv2.rectangle(inp_img, (x, y), (x+w, y+h), (0, 255, 0), 1)
+                    cv2.putText(inp_img, key, (x+5, y+5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
+
+            finish=time.time()
+            print(f"{filename} total: {finish-start}s")
+            cv2.imwrite(f"./generated/{filename}.png", inp_img)