Making Subnautica Scarier with Real-Life Effects

import cv2
import pytesseract
import pyautogui
import socket
import time
import re
import numpy as np
import os
import pygame

###############################################################################
# Configuration
###############################################################################

# Turn logs on/off
DEBUG_LOG = False  # set True to see debug-level logs
INFO_LOG = True    # set True to see info-level logs

# UDP broadcast for the ESP32
UDP_BROADCAST_IP = "255.255.255.255"
UDP_PORT = 12345

# Health bounding box and color thresholds
HEALTH_REGION = (96, 808, 60, 60)
HEALTH_LOWER_COLOR = np.array([0, 100, 100])
HEALTH_UPPER_COLOR = np.array([10, 255, 255])
HEALTH_RAW_MIN = 4.7
HEALTH_RAW_MAX = 28.88

# O2 bounding boxes
O2_REGION = (176, 779, 125, 125)
O2_LABEL_REGION = (220, 810, 37, 26)
O2_MASK_FILE = "oxygen_icon_mask.png"
O2_RAW_MIN = 5.05
O2_RAW_MAX = 97.3

LOW_O2_THRESHOLD = 15.0

# MP3 path for the Sonic alarm (provide your own file name here).
LOW_O2_AUDIO = "sonic_underwater.mp3"

# Depth bounding box
DEPTH_REGION = (896, 37, 125, 59)
DEPTH_CONFIG = r'--psm 7 -c tessedit_char_whitelist=0123456789m'
VALID_DEPTH_RANGE = (0, 3000)

###############################################################################
# Logging Helpers
###############################################################################

def log_debug(message):
    if DEBUG_LOG:
        print("[DEBUG]", message)

def log_info(message):
    if INFO_LOG:
        print("[INFO]", message)

###############################################################################
# Pygame Audio Helpers
###############################################################################

def init_audio_system():
    """
    Initialize pygame.mixer so we can load and play MP3 files.
    """
    pygame.mixer.init()
    log_debug("Pygame mixer initialized.")

def load_sound(file_path):
    """
    Load a sound (MP3 or WAV) into pygame.
    Returns a pygame.mixer.Sound object or None if load fails.
    """
    try:
        sound = pygame.mixer.Sound(file_path)
        log_debug(f"Audio file loaded: {file_path}")
        return sound
    except pygame.error as e:
        print(f"Failed to load audio file '{file_path}': {e}")
        return None

def play_alarm(sound):
    """
    Start playing the alarm in a loop if it's not already playing.
    """
    if not pygame.mixer.get_busy():
        log_info("O2 < 15% => Starting alarm.")
        sound.play(loops=-1)
    else:
        log_debug("Alarm already playing, not restarting.")

def stop_alarm(sound):
    """
    Stop the alarm if it's currently playing.
    """
    if pygame.mixer.get_busy():
        log_info("O2 recovered or icon gone => Stopping alarm.")
        sound.stop()
    else:
        log_debug("No alarm playing to stop.")

###############################################################################
# Image Capture and OCR
###############################################################################

def capture_screenshot():
    screenshot = pyautogui.screenshot()
    screenshot.save("temp_screenshot.png")
    return cv2.imread("temp_screenshot.png")

def safe_crop(img, region):
    if img is None:
        return None
    x, y, w, h = map(int, region)
    H, W = img.shape[:2]
    if x < 0: x = 0
    if y < 0: y = 0
    if x + w > W: w = max(0, W - x)
    if y + h > H: h = max(0, H - y)
    if w <= 0 or h <= 0:
        return None
    return img[y:y+h, x:x+w]

def read_text(roi, config=None):
    """
    Convert the region to grayscale, enlarge, threshold,
    then run Tesseract for text. Returns the text string.
    """
    if roi is None:
        return ""
    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    scale_factor = 3
    bigger = cv2.resize(gray, (gray.shape[1]*scale_factor, 
                               gray.shape[0]*scale_factor),
                        interpolation=cv2.INTER_CUBIC)
    _, thresh = cv2.threshold(bigger, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
    thresh = cv2.bitwise_not(thresh)
    kernel = np.ones((2,2), np.uint8)
    cleaned = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
    text = pytesseract.image_to_string(cleaned, config=config).strip()
    return text

###############################################################################
# Health / O2 / Depth Measurement
###############################################################################

def measure_health_fill(roi):
    """
    Returns scaled health % from 0..100, based on color threshold in HSV.
    """
    hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
    mask = cv2.inRange(hsv, HEALTH_LOWER_COLOR, HEALTH_UPPER_COLOR)
    total_px = roi.size // 3
    fill_px = np.count_nonzero(mask)
    raw_pct = (fill_px / total_px) * 100.0

    if raw_pct <= HEALTH_RAW_MIN:
        return 0.0
    elif raw_pct >= HEALTH_RAW_MAX:
        return 100.0
    else:
        return ((raw_pct - HEALTH_RAW_MIN) /
                (HEALTH_RAW_MAX - HEALTH_RAW_MIN)) * 100.0

def otsu_threshold(gray):
    _, th = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
    return th

def measure_o2_otsu_with_mask(roi, ring_mask):
    """
    Convert ROI to grayscale, Otsu threshold,
    then mask with ring_mask to isolate the ring's fill.
    """
    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    th = otsu_threshold(gray)
    H_roi, W_roi = roi.shape[:2]

    if ring_mask is None:
        return 0.0, th, None

    H_mask, W_mask = ring_mask.shape[:2]
    if (W_mask != W_roi) or (H_mask != H_roi):
        ring_mask_resized = cv2.resize(ring_mask, (W_roi, H_roi), 
                                       interpolation=cv2.INTER_NEAREST)
    else:
        ring_mask_resized = ring_mask

    masked = cv2.bitwise_and(th, th, mask=ring_mask_resized)
    ring_pixels = np.count_nonzero(ring_mask_resized)
    if ring_pixels == 0:
        return 0.0, th, masked

    white_in_ring = np.count_nonzero(masked)
    raw_pct = (white_in_ring / ring_pixels) * 100.0

    if raw_pct <= O2_RAW_MIN:
        scaled = 0.0
    elif raw_pct >= O2_RAW_MAX:
        scaled = 100.0
    else:
        scaled = ((raw_pct - O2_RAW_MIN) /
                  (O2_RAW_MAX - O2_RAW_MIN)) * 100.0
    return scaled, th, masked

def read_depth(roi):
    text = read_text(roi, config=DEPTH_CONFIG)
    match = re.search(r'(\d+)m', text)
    if match:
        return int(match.group(1))
    match2 = re.search(r'\d+', text)
    if match2:
        return int(match2.group(0))
    return None

###############################################################################
# Main
###############################################################################

def main():
    log_debug("Starting main function.")

    # Setup network
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)

    # Init pygame audio
    pygame.mixer.init()
    log_debug("pygame.mixer initialized.")

    # Load ring mask
    ring_mask = cv2.imread(O2_MASK_FILE, cv2.IMREAD_GRAYSCALE)
    if ring_mask is None:
        print(f"Failed to load ring mask: {O2_MASK_FILE}")
        return

    # Load alarm sound
    try:
        alarm_sound = pygame.mixer.Sound(LOW_O2_AUDIO)
        log_debug(f"Loaded alarm sound: {LOW_O2_AUDIO}")
    except pygame.error as e:
        print(f"Could not load alarm sound: {e}")
        return

    prev_health = None
    o2_alarm_playing = False

    while True:
        img = capture_screenshot()
        if img is None:
            log_debug("No screenshot captured, sleeping.")
            time.sleep(1)
            continue

        # --- Health ---
        health_roi = safe_crop(img, HEALTH_REGION)
        if health_roi is not None:
            health_val = measure_health_fill(health_roi)
            log_info(f"Health: {health_val:.1f}%")
            if prev_health is not None and health_val > 0 and health_val < prev_health:
                log_info("Health decreased => broadcast BUZZ")
                sock.sendto(b"BUZZ\n", (UDP_BROADCAST_IP, UDP_PORT))
            prev_health = health_val
        else:
            log_debug("Health ROI is None.")

        # --- O2 ---
        label_roi = safe_crop(img, O2_LABEL_REGION)
        o2_icon_text = ""
        if label_roi is not None:
            o2_icon_text = read_text(label_roi).upper()
        log_info(f"O2 icon text => '{o2_icon_text}'")

        # We'll consider multiple potential OCR outputs that represent O2
        is_o2 = any(sub in o2_icon_text for sub in ["O2","O.","OS","OS."])
        if is_o2:
            o2_roi = safe_crop(img, O2_REGION)
            if o2_roi is not None:
                o2_fill, _, _ = measure_o2_otsu_with_mask(o2_roi, ring_mask)
                log_info(f"O2 fill: {o2_fill:.1f}%")

                if o2_fill < LOW_O2_THRESHOLD:
                    if not o2_alarm_playing:
                        log_info("O2 < 15 => starting alarm with pygame")
                        alarm_sound.play(loops=-1)
                        o2_alarm_playing = True
                else:
                    if o2_alarm_playing:
                        log_info("O2 recovered => stopping alarm")
                        alarm_sound.stop()
                        o2_alarm_playing = False
            else:
                log_debug("O2 ROI is None.")
        else:
            log_debug("Not O2 => skip O2 fill.")
            # If not recognized as O2, also stop if playing
            if o2_alarm_playing:
                log_info("No O2 circle => stopping alarm")
                alarm_sound.stop()
                o2_alarm_playing = False

        # --- Depth ---
        depth_roi = safe_crop(img, DEPTH_REGION)
        if depth_roi is not None:
            depth_val = read_depth(depth_roi)
            if depth_val is not None and VALID_DEPTH_RANGE[0] <= depth_val <= VALID_DEPTH_RANGE[1]:
                log_debug(f"Depth: {depth_val}m")
                msg = f"DEPTH:{depth_val}\n".encode('utf-8')
                sock.sendto(msg, (UDP_BROADCAST_IP, UDP_PORT))
            else:
                log_debug(f"Depth invalid: {depth_val}")
        else:
            log_debug("Depth ROI is None.")

        # Check if user pressed 'q'
        key = cv2.waitKey(1000) & 0xFF
        if key == ord('q'):
            log_debug("'q' pressed => break loop")
            break

    # On exit
    if o2_alarm_playing:
        log_info("Stopping alarm on exit.")
        alarm_sound.stop()

    cv2.destroyAllWindows()
    sock.close()
    pygame.quit()  # optional to close the mixer
    log_debug("main() ended, script done.")

if __name__ == "__main__":
    main()