A speedometer for `Zelda - Tears of the Kingdom` to overlays speed information on top of a gameplay recording.

import sys
import datetime

from PyQt6 import QtCore, QtWidgets
from PyQt6.QtGui import QPainter, QPen, QColor
from PyQt6.QtWidgets import QMainWindow

import settings

class SpeedometerOverlay(QMainWindow):
    def __init__(self):
        QMainWindow.__init__(self)
        self.setWindowFlags(
            QtCore.Qt.WindowType.WindowStaysOnTopHint |
            QtCore.Qt.WindowType.FramelessWindowHint |
            QtCore.Qt.WindowType.X11BypassWindowManagerHint
        )
        self.setAttribute(
            QtCore.Qt.WidgetAttribute.WA_TranslucentBackground
        )
        self.setGeometry(QtCore.QRect(0, 0, settings.overlay_width, settings.overlay_height))
        self.widget = QtWidgets.QWidget()
        self.widget.setStyleSheet('color: white;')
        self.layout = QtWidgets.QVBoxLayout()
        self.layout.setAlignment(QtCore.Qt.AlignmentFlag.AlignCenter)

        self.setStyleSheet("QLabel{font-size: 16pt;}")

        current_time = str(datetime.datetime.now().strftime("%H:%M:%S"))
        self.l_time = QtWidgets.QLabel('Time: '+current_time)

        self.l_total = QtWidgets.QLabel('Total')
        self.l_total.setStyleSheet("QLabel{font-size: 20pt; padding-top: 1em}")
        self.l_speed = QtWidgets.QLabel('Speed: 0.00 m/s')
        self.l_avg_speed = QtWidgets.QLabel('Avg: 0.00 m/s')
        self.l_max_speed = QtWidgets.QLabel('Max: 0.00 m/s')

        self.l_horizontal = QtWidgets.QLabel('Horizontal')
        self.l_horizontal.setStyleSheet("QLabel{font-size: 20pt; padding-top: 1em}")
        self.l_speed_h = QtWidgets.QLabel('Speed: 0.00 m/s')
        self.l_avg_speed_h = QtWidgets.QLabel('Avg: 0.00 m/s')
        self.l_max_speed_h = QtWidgets.QLabel('Max: 0.00 m/s')

        self.l_vertical = QtWidgets.QLabel('Vertical')
        self.l_vertical.setStyleSheet("QLabel{font-size: 20pt; padding-top: 1em}")
        self.l_speed_v = QtWidgets.QLabel('Speed: 0.00 m/s')
        self.l_avg_speed_v = QtWidgets.QLabel('Avg: 0.00 m/s')
        self.l_max_speed_v = QtWidgets.QLabel('Max: 0.00 m/s')

        self.layout.addWidget(self.l_time)
        self.layout.addWidget(self.l_total)
        self.layout.addWidget(self.l_speed)
        self.layout.addWidget(self.l_avg_speed)
        self.layout.addWidget(self.l_max_speed)
        self.layout.addWidget(self.l_horizontal)
        self.layout.addWidget(self.l_speed_h)
        self.layout.addWidget(self.l_avg_speed_h)
        self.layout.addWidget(self.l_max_speed_h)
        self.layout.addWidget(self.l_vertical)
        self.layout.addWidget(self.l_speed_v)
        self.layout.addWidget(self.l_avg_speed_v)
        self.layout.addWidget(self.l_max_speed_v)

        self.widget.setLayout(self.layout)
        self.setCentralWidget(self.widget)

    def paintEvent(self, event=None):
        painter = QPainter(self)

        painter.setOpacity(0.2)
        painter.setBrush(QColor(0,0,0))
        painter.setPen(QPen(QColor(0,0,0)))
        painter.drawRect(self.rect())

    def update_labels(self, stats):
        current_time = str(datetime.datetime.now().strftime("%H:%M:%S"))
        self.l_time.setText('Time: '+current_time)
        self.l_speed.setText(f'Speed: {stats["total"]["Speed"]:0.2f} m/s')
        self.l_avg_speed.setText(f'Avg: {stats["total"]["Avg"]:0.2f} m/s')
        self.l_max_speed.setText(f'Max: {stats["total"]["Max"]:0.2f} m/s')
        self.l_speed_h.setText(f'Speed: {stats["horizontal"]["Speed"]:0.2f} m/s')
        self.l_avg_speed_h.setText(f'Avg: {stats["horizontal"]["Avg"]:0.2f} m/s')
        self.l_max_speed_h.setText(f'Max: {stats["horizontal"]["Max"]:0.2f} m/s')
        self.l_speed_v.setText(f'Speed: {stats["vertical"]["Speed"]:0.2f} m/s')
        self.l_avg_speed_v.setText(f'Avg: {stats["vertical"]["Avg"]:0.2f} m/s')
        self.l_max_speed_v.setText(f'Max: {stats["vertical"]["Max"]:0.2f} m/s')
        self.widget.setLayout(self.layout)

    def mousePressEvent(self, event):
        sys.exit(0)

overlay.py
miguelqncosta/TotK-Speedometer

### GENERAL SETTINGS ###

# Read corrdinates and calculate speed every x frames
# Reading every frame is very incacurate. For better results is is best to read every 5 or 10 frames.
calc_every_x_frames = 10

# Length of the lists used to calculate the averages and maximums.
# Use the last 10 speed values to calculate the average and maximum.
avg_length = 10


# Text overlay colors
title_color = (255, 255, 255)
text_color_ok = (255, 255, 255) # white text when the coordinates are OK
text_color_fail = (200, 200, 200) # grayed out text when the coordinates are not valid



### VIDEO MODE SETTINGS ###
# Only used when running 'python3 totk-speedometer.py -f ...'

# Name of the directory where the videos will be saved.
output_directory = 'totk-speedometer-videos'

# Show preview while creating the video
show_preview = False



### REAL-TIME OVERLAY SETTINGS ###
# Only used when running 'python3 totk-speedometer.py -s'

# Position for the real-time mode overlay
# 0,0 is the top left corner of the screen
overlay_vertical_pos = 500
overlay_horizontal_pos = 2150
overlay_width = 300
overlay_height = 400



### DEBUGGING SETTINGS ###

# Save an image for each pre-processing stage to images/preprocessing/
save_preprocessing_images = False
settings.py
miguelqncosta/TotK-Speedometer

import os
import sys
import time

import cv2
import mss.tools
import numpy as np
import pytesseract
from moviepy.editor import *
from PyQt6 import QtWidgets
from PyQt6.QtCore import QRunnable, QThreadPool
from PyQt6.QtGui import QScreen

import settings
from overlay import SpeedometerOverlay

# Lists to keep historical values
speed_list = []
speed_h_list = []
speed_v_list = []


def detect_circle(map_img, width, height):
    img = cv2.cvtColor(map_img,cv2.COLOR_BGR2GRAY)

    circles = cv2.HoughCircles(img, cv2.HOUGH_GRADIENT, 1, (height*0.2),
                            param1=80,
                            param2=60,
                            minRadius=int(height*0.04),
                            maxRadius=int(height*0.12)
                            )

    if circles is not None:
        circles = np.uint16(np.around(circles))
        circles_img = img.copy()
        circles_img = cv2.cvtColor(circles_img, cv2.COLOR_GRAY2BGR)

        for i in circles[0,:]:
            # draw the outer circle
            cv2.circle(circles_img,(i[0],i[1]),i[2],(0,255,0),2)
            # draw the center of the circle
            cv2.circle(circles_img,(i[0],i[1]),2,(0,0,255),3)

        return circles, circles_img
    else:
        return None, None



def get_coord_img(map_img, map_circle):
    scaling = 16

    polar_img = cv2.warpPolar(map_img, dsize=(
                                        map_img.shape[1]*scaling,
                                        map_img.shape[0]*scaling),
                                        center=(map_circle[0],map_circle[1]
                                    ),
                                    maxRadius=map_circle[2],
                                    flags=cv2.WARP_POLAR_LOG | cv2.INTER_CUBIC )

    polar_img = cv2.rotate(polar_img, cv2.ROTATE_90_CLOCKWISE)
    cropped_img = polar_img[int(polar_img.shape[0]*0.95):, int(polar_img.shape[1]*0.58):int(polar_img.shape[1]*0.92)]

    if settings.save_preprocessing_images: 
        cv2.imwrite('images/map_polar.png', polar_img)
        cv2.imwrite('images/map_cropped.png', cropped_img)

    return cropped_img



def preprocess_coord_img(coord_img):
    coord_img = cv2.cvtColor(coord_img,cv2.COLOR_BGR2GRAY)
    blur_img = cv2.GaussianBlur(coord_img,(5,5),0)
    avg_brightness = cv2.mean(blur_img)[0]
    beta = -2*avg_brightness-50
    contrast_img = cv2.convertScaleAbs(blur_img, alpha=3.5, beta=beta)
    ret, threshold_img = cv2.threshold(contrast_img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    erode_img = cv2.erode(threshold_img, np.ones((3,3)), iterations=1)
    dilate_img = cv2.dilate(erode_img, np.ones((3,3)), iterations=2)
    processed_img = cv2.erode(dilate_img, np.ones((3,3)), iterations=2)
    
    if settings.save_preprocessing_images:
        os.makedirs(os.path.join('images', 'preprocessing'), exist_ok=True)
        cv2.imwrite('images/preprocessing/1_blur_img.png', blur_img)
        cv2.imwrite('images/preprocessing/2_contrast_img.png', contrast_img)
        cv2.imwrite('images/preprocessing/3_threshold_img.png', threshold_img)
        cv2.imwrite('images/preprocessing/4_erode_img.png', erode_img)
        cv2.imwrite('images/preprocessing/5_dilate_img.png', dilate_img)
        cv2.imwrite('images/preprocessing/6_processed_img.png', processed_img)

    return processed_img



def extract_coordinates(coord_img):
    text = pytesseract.image_to_string(coord_img, config='--psm 7 --oem 3 -c tessedit_char_whitelist=" -0123456789"')
    coord_list = text.strip().replace('--', '-').split(' ')

    if '-' in coord_list:
        coord_list.remove('-')

    try:
        if isinstance(text, str) and text != '':
            coord = [int(c) for c in coord_list]

            if len(coord) == 3:
                if (5000 > coord[0] > -5000 and 4000 > coord[1] > -4000 and 3500 > coord[2] > -3000):
                    return coord
    except ValueError:
        return None

    return None



def process_coordinates(coord, last_coord, time_delta):
    if len(coord)==3 and len(last_coord)==3 and all(isinstance(val, int) for val in coord+last_coord):
        distance = np.sqrt((coord[0]-last_coord[0])**2 + (coord[1]-last_coord[1])**2 + (coord[2]-last_coord[2])**2)
        distance_h = np.sqrt((coord[0]-last_coord[0])**2 + (coord[1]-last_coord[1])**2)
        distance_v = np.sqrt((coord[2]-last_coord[2])**2)

        speed = distance/time_delta
        if speed >= 100:
            speed = float('nan')
        else:
            speed_list.append(speed)
            if len(speed_list) > settings.avg_length:
                speed_list.pop(0)

        speed_h = distance_h/time_delta
        if speed_h >= 100:
            speed_h = float('nan')
        else:
            speed_h_list.append(speed_h)
            if len(speed_h_list) > settings.avg_length:
                speed_h_list.pop(0)

        speed_v = distance_v/time_delta
        if speed_v >= 100:
            speed_v = float('nan')
        else:
            speed_v_list.append(speed_v)
            if len(speed_v_list) > settings.avg_length:
                speed_v_list.pop(0)

        if len(speed_list) > 0:
            avg_speed = sum(speed_list)/len(speed_list)
            max_speed = max(speed_list)
        else:
            avg_speed = float('nan')
            max_speed = float('nan')

        if len(speed_h_list) > 0:
            avg_speed_h = sum(speed_h_list)/len(speed_h_list)
            max_speed_h = max(speed_h_list)
        else:
            avg_speed_h = float('nan')
            max_speed_h = float('nan')

        if len(speed_v_list) > 0:
            avg_speed_v = sum(speed_v_list)/len(speed_v_list)
            max_speed_v = max(speed_v_list)
        else:
            avg_speed_v = float('nan')
            max_speed_v = float('nan')

        results = {
                    'distance': {
                        'Distance': distance,
                        'Distance H': distance_h,
                        'Distance V': distance_v
                    },
                    'total': {
                        'Speed': speed,
                        'Avg': avg_speed,
                        'Max': max_speed
                    },
                    'horizontal': {
                        'Speed': speed_h,
                        'Avg': avg_speed_h,
                        'Max': max_speed_h
                    },
                    'vertical': {
                        'Speed': speed_v,
                        'Avg': avg_speed_v,
                        'Max': max_speed_v
                    }
                }

        return results



def export_video_with_overlay(video_path):
    print('Processing video: ' + video_path)

    video = cv2.VideoCapture(video_path)
    width  = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = video.get(cv2.CAP_PROP_FPS)
    h = int(video.get(cv2.CAP_PROP_FOURCC))
    codec = chr(h&0xff) + chr((h>>8)&0xff) + chr((h>>16)&0xff) + chr((h>>24)&0xff)

    print(f'Video - Width: {width}, Height: {height}, FPS: {fps}, Enconding: {codec}')

    path = os.path.dirname(video_path)
    filename = os.path.basename(video_path)
    name = os.path.splitext(filename)[0]
    ext = os.path.splitext(filename)[1]
    os.makedirs(os.path.join(path, settings.output_directory), exist_ok=True)
    tmp_filename = os.path.join(path, settings.output_directory, name+'_speedometer_tmp'+ext)
    output_filename = os.path.join(path, settings.output_directory, name+'_speedometer'+ext)
    video_output = cv2.VideoWriter(tmp_filename, cv2.VideoWriter_fourcc(*'avc1'), fps, (width, height))

    count = 0
    map_circle = None
    last_coord = None

    while video.isOpened():
        ret,frame = video.read()
        if not ret:
            break

        map_img = frame[int(height*0.71):int(height*0.96), int(width*0.829):int(width*0.97)]
        map_circle = [int(width*0.06875), int(height*0.1223), int(width*0.0641)]

        if settings.save_preprocessing_images:
            cv2.imwrite('images/map.png', map_img)

        if count > settings.calc_every_x_frames:
            coord_img = get_coord_img(map_img, map_circle)
            processed_img = preprocess_coord_img(coord_img)
            coord = extract_coordinates(processed_img)
            text_color = settings.text_color_fail # Set text color to gray when the coordinates are not valid

            if coord is not None:
                if last_coord is None:
                    last_coord = coord
                    count = 0
                else:
                    tmp_results = process_coordinates(coord, last_coord, (1/fps)*count)
                    if tmp_results is not None:
                        print(f'Coord: {[f"{c:5}" for c in coord]}'
                        f' - Last Coord {[f"{c:5}" for c in last_coord]}'
                        f' - Distance: {tmp_results["distance"]["Distance"]:6.2f} m'
                        f' - Speed: {tmp_results["total"]["Speed"]:5.2f} m/s'
                        f' - AvgSpeed: {tmp_results["total"]["Avg"]:5.2f} m/s'
                        f' - Speed H: {tmp_results["horizontal"]["Speed"]:5.2f} m/s'
                        f' - Speed V: {tmp_results["vertical"]["Speed"]:5.2f} m/s')

                        last_coord = coord
                        count = 0

                        if not any(np.isnan(b) for v in tmp_results.values() for b in v.values()):
                            results = tmp_results
                            text_color = settings.text_color_ok # Set text to white when the coordinates are valid

        count = count + 1

        if 'results' in locals() and results is not None:
            text_w_start = int(width*0.842)
            text_h_start = int(height*0.20)
            overlay_w_start = int(width*0.836)
            overlay_w_end = int(width*0.961)
            overlay_h_start = int(height*0.16)
            overlay_h_end = int(height*0.694)
            spacing = int(height/25)
            text_size = height/1500

            alpha = 0.25
            overlay = frame.copy()
            cv2.rectangle(overlay, (overlay_w_start, overlay_h_start), (overlay_w_end, overlay_h_end), (0,0,0), cv2.FILLED)
            cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0, frame)

            text_h = text_h_start
            frame = cv2.putText(frame, 'Total', (text_w_start, text_h), cv2.FONT_HERSHEY_TRIPLEX, text_size, settings.title_color, 1, cv2.LINE_AA)
            for key,value in results['total'].items():
                text_h = text_h + spacing
                text = f'{key}: {value:.2f} m/s'
                frame = cv2.putText(frame, text, (text_w_start, text_h), cv2.FONT_HERSHEY_SIMPLEX, text_size, text_color, 1, cv2.LINE_AA)

            text_h = int(text_h + (1.5*spacing))
            frame = cv2.putText(frame, 'Horizontal', (text_w_start, text_h), cv2.FONT_HERSHEY_TRIPLEX, text_size, settings.title_color, 1, cv2.LINE_AA)
            for key,value in results['horizontal'].items():
                text_h = text_h + spacing
                text = f'{key}: {value:.2f} m/s'
                frame = cv2.putText(frame, text, (text_w_start, text_h), cv2.FONT_HERSHEY_SIMPLEX, text_size, text_color, 1, cv2.LINE_AA)

            text_h = int(text_h + (1.5*spacing))
            frame = cv2.putText(frame, 'Vertical', (text_w_start, text_h), cv2.FONT_HERSHEY_TRIPLEX, text_size, settings.title_color, 1, cv2.LINE_AA)
            for key,value in results['vertical'].items():
                text_h = text_h + spacing
                text = f'{key}: {value:.2f} m/s'
                frame = cv2.putText(frame, text, (text_w_start, text_h), cv2.FONT_HERSHEY_SIMPLEX, text_size, text_color, 1, cv2.LINE_AA)

        video_output.write(frame)

        if settings.show_preview:
            cv2.imshow('TotK Speedometer', frame)
            # Press 'q' key to quit
            if cv2.waitKey(1) & 0xFF == ord('q'):
                return

    video.release()
    video_output.release()
    cv2.destroyAllWindows()

    # Add audio from original video
    orig_clip = VideoFileClip(video_path)
    orig_audioclip = orig_clip.audio
    clip = VideoFileClip(tmp_filename)
    clip_with_audio = clip.set_audio(orig_audioclip)
    clip_with_audio.write_videofile(output_filename, codec='libx264', audio_codec='aac', fps=fps)
    orig_clip.close()
    clip.close()

    # Remove temporary video file
    os.remove(tmp_filename)



def detect_map(monitor_number):
    with mss.mss() as sct:
        mon = sct.monitors[monitor_number]

        monitor = {
                'top': mon['top'],
                'left': mon['left'],
                'width': mon['width'],
                'height': mon['height'],
                'mon': monitor_number,
            }

        print('Searching map position...')

        while True:
            sct_img = sct.grab(monitor)
            img = np.array(sct_img)
            cv2.imwrite('images/screenshot.png', img)

            h, w, c = img.shape
            if mon['width']*2 == w:
                scaling = 2
            else:
                scaling = 1

            circles, circles_img = detect_circle(img, w, h)
            if circles_img is not None:
                cv2.imwrite('images/detected_circles.png', circles_img)

            if circles is not None:
                for map_circle in circles[0]:
                    # Capture only the screen part the map occupies
                    monitor_region = {
                        'top': mon['top']+(map_circle[1]-map_circle[2])/scaling - 20,
                        'left': mon['left']+(map_circle[0]-map_circle[2])/scaling - 20,
                        'width': ((map_circle[2]*2)/scaling) + 40,
                        'height': ((map_circle[2]*2)/scaling) + 40,
                        'mon': monitor_number,
                    }

                    sct_img = sct.grab(monitor_region)
                    map_img = np.array(sct_img)
                    cv2.imwrite('images/map.png', map_img)

                    circles, circles_img = detect_circle(map_img, w, h)
                    if circles_img is not None:
                        cv2.imwrite('images/detected_map_circles.png', circles_img)

                    if circles is not None:
                        if len(circles[0])==1:
                            map_circle = circles[0][0]
                            coord_img = get_coord_img(map_img, map_circle)
                            processed_img = preprocess_coord_img(coord_img)
                            coord = extract_coordinates(processed_img)
                            if coord is not None:
                                tmp_results = process_coordinates(coord, coord, 1)
                                if tmp_results is not None:
                                    print('Map position detected.')
                                    return map_circle, monitor_region



class SpeedometerRunnable(QRunnable):
    def __init__(self, mainwindow, monitor_sct, map_circle):
        super().__init__()
        self.mainwindow = mainwindow
        self.monitor_sct = monitor_sct
        self.map_circle = map_circle
        self.running = True

    def stop(self):
        self.running = False

    def run(self):
        last_coord = None
        last_coord_time = None

        with mss.mss() as sct:
            while self.running:
                # Grab the map screenshot
                sct_img = sct.grab(self.monitor_sct)

                # Save the map screenshot
                mss.tools.to_png(sct_img.rgb, sct_img.size, output='images/map.png')
                map_img = np.array(sct_img)

                coord_img = get_coord_img(map_img, self.map_circle)
                processed_img = preprocess_coord_img(coord_img)
                coord = extract_coordinates(processed_img)

                if coord is not None:
                    if last_coord is None:
                        last_coord = coord
                        last_coord_time = time.time()
                    else:
                        t = time.time()
                        results = process_coordinates(coord, last_coord, t-last_coord_time)

                        if results is not None:
                            print(f'Coord: {[f"{c:5}" for c in coord]}'
                            f' - Last Coord {[f"{c:5}" for c in last_coord]}'
                            f' - Distance: {results["distance"]["Distance"]:6.2f} m'
                            f' - Speed: {results["total"]["Speed"]:5.2f} m/s'
                            f' - AvgSpeed: {results["total"]["Avg"]:5.2f} m/s'
                            f' - Speed H: {results["horizontal"]["Speed"]:5.2f} m/s'
                            f' - Speed V: {results["vertical"]["Speed"]:5.2f} m/s')

                            last_coord_time = t
                            last_coord = coord

                            if not any(np.isnan(b) for v in results.values() for b in v.values()):
                                self.mainwindow.update_labels(results)
                                time.sleep(0.01)



def main():
    import argparse

    parser = argparse.ArgumentParser(description='Speedometer for Zelda TotK')
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('-f', dest='files', metavar='file', type=str, nargs='+', help='path to video file. Accepts multiple files')
    group.add_argument('-s', dest='screen', action='store_true', help='use screen capture and overlay stats')
    parser.add_argument('-m', dest='monitor', default=1, type=int, help='monitor number to capture and display the overlay.')
    group.add_argument('--test', dest='test', action='store_true', help='test image pre-processing')
    args = parser.parse_args()

    parser.add_argument('file', type=argparse.FileType('r'), nargs='+')

    if args.files is None and args.screen is False and args.test is False:
        parser.print_help()
        exit()

    if args.files is not None:
        for f in args.files:
            if os.path.isfile(f): 
                export_video_with_overlay(f)

    elif args.screen:
        map_circle, monitor_sct = detect_map(args.monitor)
        app = QtWidgets.QApplication(sys.argv)
        mainwindow = SpeedometerOverlay()
        monitors = QScreen.virtualSiblings(mainwindow.screen())
        monitor = monitors[args.monitor-1].availableGeometry()
        mainwindow.move(monitor.left() + settings.overlay_horizontal_pos, monitor.top() + settings.overlay_vertical_pos)
        mainwindow.show()
        runnable = SpeedometerRunnable(mainwindow, monitor_sct, map_circle)
        QThreadPool.globalInstance().start(runnable)
        sys.exit(app.exec())

    elif args.test:
        map_img = cv2.imread('images/map.png')
        width = 1280
        height = 720
        map_circle = [int(width*0.06875), int(height*0.1223), int(width*0.0641)]
        coord_img = get_coord_img(map_img, map_circle)
        processed_img = preprocess_coord_img(coord_img)
        coord = extract_coordinates(processed_img)
        print(coord)


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(0)

totk-speedometer.py
miguelqncosta/TotK-Speedometer