diff --git a/assets/hud_mask.png b/assets/hud_mask.png new file mode 100644 index 0000000..08e28a3 Binary files /dev/null and b/assets/hud_mask.png differ diff --git a/assets/items/magic_gg_club.png b/assets/items/magic_gg_club.png index c9ae8dd..a2241e7 100644 Binary files a/assets/items/magic_gg_club.png and b/assets/items/magic_gg_club.png differ diff --git a/game.ini b/game.ini index 236e9b7..ce0cfaf 100644 --- a/game.ini +++ b/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 diff --git a/src/item/item_cropper.py b/src/item/item_cropper.py new file mode 100644 index 0000000..6368c63 --- /dev/null +++ b/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) diff --git a/src/item/item_finder.py b/src/item/item_finder.py index 92887dd..3575bc1 100644 --- a/src/item/item_finder.py +++ b/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) diff --git a/src/utils/item_extractor.py b/src/utils/item_extractor.py new file mode 100644 index 0000000..4a3be89 --- /dev/null +++ b/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)