d4-altar-overlay.py

import cv2
import json
import math
import time
import sys
import os
import numpy as np
from mss import mss
import pygame as pg
import win32gui
import win32con
import win32api
import argparse
import pyWinhook
import threading
import pythoncom

global clear_overlay
global update_overlay
global should_exit
global markers

def on_key_down(event):
	global clear_overlay
	global update_overlay
	global update_overlay
	global markers
	
	# Handle key events as needed
	if event.Key in ['m', 't', 'Escape']:
		clear_overlay = True
		if opt.debug == True:
			print("Key '" + event.Key + "' pressed")
		
	if event.Key in ['F1']:
		update_overlay = True
		markers = markers_altar
		if opt.debug == True:
			print("Key '" + event.Key + "' pressed")
	
	if event.Key in ['F2']:
		update_overlay = True
		markers = markers_mysterious
		if opt.debug == True:
			print("Key '" + event.Key + "' pressed")
	
	if event.Key in ['F3']:
		update_overlay = True
		markers = markers_waygate
		if opt.debug == True:
			print("Key '" + event.Key + "' pressed")
	
	if event.Key in ['F4']:
		global should_exit
		should_exit = True
		if opt.debug == True:
			print("Key '" + event.Key + "' pressed")
		
	return True  # Continue the hook

def on_mouse_event(event):
	global clear_overlay
	clear_overlay = True
	if opt.debug == True:
		print("Mouse button pressed")

	return True # Continue the hook

# Thread function for pyWinhook event handling
def pywinhook_thread_func(exit_flag):
	hm = pyWinhook.HookManager()
	
	hm.KeyDown = on_key_down
	hm.HookKeyboard()
	
	hm.MouseLeftDown = on_mouse_event
	hm.HookMouse()

	while not exit_flag.is_set():
		pythoncom.PumpWaitingMessages()

	hm.UnhookKeyboard()
	hm.UnhookMouse()

exit_flag = threading.Event()

def resource_path(relative):
	if hasattr(sys, "_MEIPASS") == True:
		return os.path.join(
			sys._MEIPASS,
			relative
		)
	else:
		return relative

class MapCoords():
	def __init__(self):
		self.MAX_LATITUDE = 85.0511287798
		self.R = 6378137
		self._a = 2.495320233665337e-8
		self._b = 0.5
		self._c = -2.495320233665337e-8
		self._d = 0.5

	def latLngToPoint(self, pt, zoom = 1):
		x, y = self.project(pt)
		n = self.scale(zoom)
		return self.transform(x, y, n)
	
	def scale(self, n):
		return 256 * math.pow(2, n)
	
	def project(self, pt):
		lat = pt[0]
		lng = pt[1]
		i = math.pi / 180
		e = self.MAX_LATITUDE
		n = max(min(e, lat), -e)
		o = math.sin(n * i)
		return self.R * lng * i, self.R * math.log((1 + o) / (1 - o)) / 2
	
	def transform(self, x, y, n):
		return n * (self._a * x + self._b), n * (self._c * y + self._d)

class MapSIFTMatcher():
	def __init__(self, map_image):
		self.sift = cv2.SIFT_create()
		self.map_image = map_image
		self.kp_map, self.des_map = self.sift.detectAndCompute(map_image, None)

		FLANN_INDEX_KDTREE = 1
		index_params = dict(algorithm = FLANN_INDEX_KDTREE, trees = 5)
		search_params = dict(checks = 50)
		self.flann = cv2.FlannBasedMatcher(index_params, search_params)

	def query(self, query_image, markers, MIN_MATCH_COUNT = 10, debug = False):
		img = None
		height, width = None, None
		out_markers = []

		kp_query, des_query = self.sift.detectAndCompute(query_image, None)
		matches = self.flann.knnMatch(des_query, self.des_map, k=2)

		good = []
		for m, n in matches:
			if m.distance < 0.7 * n.distance:
				good.append(m) 

		if len(good) > MIN_MATCH_COUNT:
			h, w, d = query_image.shape

			src_pts = np.float32([ kp_query[m.queryIdx].pt for m in good ]).reshape(-1, 1, 2)
			dst_pts = np.float32([ self.kp_map[m.trainIdx].pt for m in good ]).reshape(-1, 1, 2)

			M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
			self.last_homography_matrix = M

			matchesMask = mask.ravel().tolist()

			try:
				transform_points = cv2.perspectiveTransform(np.float32([[0, 0], [0, h-1], [w - 1, h - 1], [w - 1, 0]]).reshape(-1, 1, 2), M)
			except:
				matchesMask = None
				if debug == True:
					temp = self.map_image.copy()
			else:
				topleft,bottomleft,bottomright,topright = [np.int32(x[0]) for x in transform_points]
				width = bottomright[0] - topleft[0]
				height = bottomright[1] - topleft[1]

				if debug == True:
					temp = cv2.polylines(self.map_image.copy(), [np.int32(transform_points)], True, (255, 0, 0), 3, cv2.LINE_AA)

		else:
			matchesMask = None
			if debug == True:
				print( "Not enough matches are found - {}/{}".format(len(good), MIN_MATCH_COUNT) )
				temp = self.map_image.copy()

		if debug == True:
			draw_params = dict(matchColor = (0, 255, 0), singlePointColor = None, matchesMask = matchesMask, flags = 2)
			query_image_temp = query_image.copy()
		
		for m in markers:
			x = int(round(m[0], 0))
			y = int(round(m[1], 0))
			if debug == True:
				cv2.circle(temp, (x, y), 10, (255,255,255), -1)

			if width != None:
				if x >= topleft[0] and x <= topright[0] and y >= topleft[1] and y <= bottomleft[1]:
					x -= topleft[0]
					y -= topleft[1]
					x = int(round(x * (w / width), 0))
					y = int(round(y * (h / height), 0))
					if debug == True:
						cv2.circle(query_image_temp, (x, y), 10, (255, 255, 255), -1)
					
					out_markers.append([x, y])

		if debug == True:
			img = cv2.drawMatches(query_image_temp, kp_query, temp, self.kp_map, good, None, **draw_params)
			img = cv2.resize(img, (0,0), fx = 0.5, fy = 0.5)

		return img, out_markers, len(good)

if __name__ == '__main__':
	parser = argparse.ArgumentParser()
	parser.add_argument('--use-large-map', default=False, type=lambda x: x.lower() not in ['false', 'no', '0', 'None'], help='use large map')
	parser.add_argument('--debug', default=False, type=lambda x: x.lower() not in ['false', 'no', '0', 'None'], help='toggle debug mode')
	opt = parser.parse_args()

	print("")
	print("======================================")
	print("Diablo IV: Altar of Lilith Map Overlay")
	print("======================================")
	print("")
	print("Press F1 for Altars of Lilith")
	print("Press F2 for Mysterious Chests")
	print("Press F3 for Waygates")
	print("Press F4 to exit...")

	if opt.use_large_map == True:
		zoom = 6
		map_image_path = resource_path('map_images/map_5.jpg')
	else:
		zoom = 5
		map_image_path = resource_path('map_images/map_5_small.jpg')

	map_scale = 0.5
	map_image = cv2.imread(map_image_path)
	map_image = cv2.resize(map_image, (0,0), fx = map_scale, fy = map_scale)

	mc = MapCoords()
	
	markers_altar = []
	markers_mysterious = []
	markers_waygate = []
	
	with open(resource_path('data.json'), 'r') as f:
		data = json.load(f)
		jmarkers = data["markers"]
		for m in jmarkers:
			if m["type"] == "Altar of Lilith":
				x, y = mc.latLngToPoint(m["coords"], zoom) # 5 = small map, 6 = large map ::: tiled map (saved from web) zoom level + 1
				x /= (1 / map_scale)
				y /= (1 / map_scale)
				markers_altar.append([x, y])
			if m["type"] == "Mysterious Chest":
				x, y = mc.latLngToPoint(m["coords"], zoom) # 5 = small map, 6 = large map ::: tiled map (saved from web) zoom level + 1
				x /= (1 / map_scale)
				y /= (1 / map_scale)
				markers_mysterious.append([x, y])
			if m["type"] == "Waygate":
				x, y = mc.latLngToPoint(m["coords"], zoom) # 5 = small map, 6 = large map ::: tiled map (saved from web) zoom level + 1
				x /= (1 / map_scale)
				y /= (1 / map_scale)
				markers_waygate.append([x, y])

	stm = MapSIFTMatcher(map_image)
	pg.init()
	info = pg.display.Info()
	screen = pg.display.set_mode((info.current_w, info.current_h), pg.FULLSCREEN)
	screen_rect = screen.get_rect()
	FPS = 1
	clock = pg.time.Clock()

	# Create and start the pyWinhook thread
	pywinhook_thread = threading.Thread(target=pywinhook_thread_func, args=(exit_flag,))
	pywinhook_thread.start()

	transparency_color = (255, 0, 128)
	hwnd = pg.display.get_wm_info()["window"]
	styles = win32gui.GetWindowLong(hwnd, win32con.GWL_EXSTYLE)
	styles = win32con.WS_EX_LAYERED | win32con.WS_EX_TRANSPARENT | win32con.WS_EX_TOOLWINDOW
	win32gui.SetWindowLong(hwnd, win32con.GWL_EXSTYLE, styles)
	win32gui.SetLayeredWindowAttributes(hwnd, win32api.RGB(*transparency_color), 0, win32con.LWA_COLORKEY)
	win32gui.SetWindowPos(hwnd, win32con.HWND_TOPMOST, 0,0,0,0, win32con.SWP_NOMOVE | win32con.SWP_NOSIZE)

	screen.fill((255, 0, 128))
	pg.display.flip()

	sct = mss()
	w, h = info.current_w, info.current_h
	monitor = {'top': 0, 'left': 0, 'width': w, 'height': h}
	
	
	running = True
	if opt.debug == True:
		print("Run loop initiated...")
		
	update_overlay = False
	clear_overlay = False
	should_exit = False
	
	while running:		
		if update_overlay:
			if opt.debug == True:
				print("update_overlay triggered")
			start_ss = time.time()
			screenshot = np.array(sct.grab(monitor))
			screenshot = cv2.cvtColor(screenshot, cv2.COLOR_RGB2BGR)
			end_ss = time.time()

			start_query = time.time()
			debug_img, out_markers, num_matches = stm.query(screenshot, markers, MIN_MATCH_COUNT =  40, debug = opt.debug)
			end_query = time.time()
			if opt.debug == True:
				print(f"Query: {end_query-start_query}s, Grab: {end_ss-start_ss}s, Num matches: {num_matches}")

			screen.fill((255, 0, 128))
			pg.draw.circle(screen, '#00ff00', (10, 10), 5)
			for m in out_markers:
				pg.draw.circle(screen, '#ff00ff', (m[0] + monitor["left"], m[1] + monitor["top"]), 3)

			pg.display.flip()
			
			update_overlay = False
		
		if clear_overlay:
			if opt.debug == True:
				print("clear_overlay triggered")
			screen.fill((255, 0, 128))
			pg.display.flip()
			
			clear_overlay = False
		
		if should_exit:
			if opt.debug == True:
				print("should_exit triggered")
			running = False
			exit_flag.set()
			
			should_exit = False
	
		clock.tick(FPS)

	# Wait for the pyWinhook thread to complete
	pywinhook_thread.join()
	
	if opt.debug == True:
		print("exiting pygame")
	
	pg.quit()