changedetection.io/changedetectionio/processors/image_ssim_diff/util.py

"""
Subprocess-isolated utility functions for image operations.

These functions use multiprocessing to run image operations in separate
processes, ensuring complete memory cleanup. Raw bytes are passed via
Pipe to avoid pickle overhead for large objects.
"""

import multiprocessing
from loguru import logger


def _worker_find_region_with_template_matching(conn, original_bbox, search_tolerance):
    """
    Worker function for template matching (runs in subprocess).

    Receives via pipe: (current_img_bytes, template_bytes)
    Sends via pipe: new_bbox [left, top, right, bottom] or None
    """
    import cv2
    import numpy as np
    from PIL import Image
    import io

    try:
        # Receive image data from parent process
        current_img_bytes, template_bytes = conn.recv()

        # Load images from bytes
        template_img = Image.open(io.BytesIO(template_bytes))
        template_img.load()

        current_img = Image.open(io.BytesIO(current_img_bytes))
        current_img.load()

        # Calculate search region
        left, top, right, bottom = original_bbox
        width = right - left
        height = bottom - top

        margin_x = int(width * search_tolerance)
        margin_y = int(height * search_tolerance)

        search_left = max(0, left - margin_x)
        search_top = max(0, top - margin_y)
        search_right = min(current_img.width, right + margin_x)
        search_bottom = min(current_img.height, bottom + margin_y)

        # Crop search region
        current_img_cropped = current_img.crop((search_left, search_top, search_right, search_bottom)).copy()
        current_img_cropped.load()

        # Convert to numpy arrays
        current_array = np.array(current_img_cropped)
        template_array = np.array(template_img)

        # Convert to grayscale
        if len(current_array.shape) == 3:
            current_gray = cv2.cvtColor(current_array, cv2.COLOR_RGB2GRAY)
        else:
            current_gray = current_array

        if len(template_array.shape) == 3:
            template_gray = cv2.cvtColor(template_array, cv2.COLOR_RGB2GRAY)
        else:
            template_gray = template_array

        logger.debug(f"[subprocess] Searching for template in region: ({search_left}, {search_top}) to ({search_right}, {search_bottom})")

        # Perform template matching
        result = cv2.matchTemplate(current_gray, template_gray, cv2.TM_CCOEFF_NORMED)
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)

        logger.debug(f"[subprocess] Template matching confidence: {max_val:.2%}")

        # Check if match is good enough (80% confidence threshold)
        if max_val >= 0.8:
            # Calculate new bounding box in original image coordinates
            match_x = search_left + max_loc[0]
            match_y = search_top + max_loc[1]

            new_bbox = [match_x, match_y, match_x + width, match_y + height]

            # Calculate movement distance
            move_x = abs(match_x - left)
            move_y = abs(match_y - top)

            logger.info(f"[subprocess] Template found at ({match_x}, {match_y}), "
                       f"moved {move_x}px horizontally, {move_y}px vertically, "
                       f"confidence: {max_val:.2%}")

            conn.send(new_bbox)
        else:
            logger.warning(f"[subprocess] Template match confidence too low: {max_val:.2%} (need 80%)")
            conn.send(None)

    except Exception as e:
        logger.error(f"[subprocess] Template matching error: {e}")
        conn.send(None)
    finally:
        conn.close()


def _worker_regenerate_template(conn, bbox):
    """
    Worker function for template regeneration (runs in subprocess).

    Receives via pipe: (snapshot_img_bytes, output_path)
    Sends via pipe: True/False for success
    """
    from PIL import Image
    import io
    import os

    try:
        # Receive data from parent process
        snapshot_img_bytes, output_path = conn.recv()

        left, top, right, bottom = bbox
        width = right - left
        height = bottom - top

        # Ensure output directory exists
        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        # Load snapshot from bytes
        snapshot_img = Image.open(io.BytesIO(snapshot_img_bytes))
        snapshot_img.load()

        # Crop template region
        template = snapshot_img.crop(tuple(bbox)).copy()
        template.load()

        # Save as PNG
        template.save(output_path, format='PNG', optimize=True)

        logger.info(f"[subprocess] Regenerated template: {output_path} ({width}x{height}px)")
        conn.send(True)

    except Exception as e:
        logger.error(f"[subprocess] Failed to regenerate template: {e}")
        conn.send(False)
    finally:
        conn.close()


def _worker_crop_image(conn, bbox):
    """
    Worker function for image cropping (runs in subprocess).

    Receives via pipe: input_img_bytes
    Sends via pipe: cropped_img_bytes or None
    """
    from PIL import Image
    import io

    try:
        # Receive image data from parent process
        input_img_bytes = conn.recv()

        # Load image from bytes
        img = Image.open(io.BytesIO(input_img_bytes))
        img.load()

        # Crop and force independent buffer
        cropped = img.crop(tuple(bbox)).copy()
        cropped.load()

        # Convert back to bytes
        output = io.BytesIO()
        cropped.save(output, format='PNG', optimize=True)
        cropped_bytes = output.getvalue()

        logger.debug(f"[subprocess] Cropped image: {cropped.size}")
        conn.send(cropped_bytes)

    except Exception as e:
        logger.error(f"[subprocess] Failed to crop image: {e}")
        conn.send(None)
    finally:
        conn.close()


# Public API functions

def find_region_with_template_matching_isolated(
    current_img_bytes,
    template_bytes,
    original_bbox,
    search_tolerance=0.2
):
    """
    Find region using template matching in isolated subprocess.

    Args:
        current_img_bytes: Current screenshot as PNG bytes
        template_bytes: Template image as PNG bytes
        original_bbox: (left, top, right, bottom) tuple
        search_tolerance: How far to search (0.2 = ±20% of region size)

    Returns:
        tuple: New (left, top, right, bottom) region, or None if not found
    """
    parent_conn, child_conn = multiprocessing.Pipe()

    # Start subprocess
    p = multiprocessing.Process(
        target=_worker_find_region_with_template_matching,
        args=(child_conn, original_bbox, search_tolerance)
    )
    p.start()

    try:
        # Send image data to subprocess
        parent_conn.send((current_img_bytes, template_bytes))

        # Wait for result
        result = parent_conn.recv()

        # Convert list back to tuple if we got a result
        if result is not None:
            result = tuple(result)

        return result

    finally:
        parent_conn.close()
        p.join(timeout=30)  # 30 second timeout
        if p.is_alive():
            logger.warning("Template matching subprocess timed out, terminating")
            p.terminate()
            p.join()


def regenerate_template_isolated(snapshot_img_bytes, bbox, output_path):
    """
    Regenerate template file from snapshot in isolated subprocess.

    Args:
        snapshot_img_bytes: Snapshot image as PNG bytes
        bbox: (left, top, right, bottom) tuple
        output_path: Where to save the template PNG

    Returns:
        bool: True if successful, False otherwise
    """
    parent_conn, child_conn = multiprocessing.Pipe()

    # Start subprocess
    p = multiprocessing.Process(
        target=_worker_regenerate_template,
        args=(child_conn, bbox)
    )
    p.start()

    try:
        # Send data to subprocess
        parent_conn.send((snapshot_img_bytes, output_path))

        # Wait for result
        success = parent_conn.recv()
        return success

    finally:
        parent_conn.close()
        p.join(timeout=30)  # 30 second timeout
        if p.is_alive():
            logger.warning("Template regeneration subprocess timed out, terminating")
            p.terminate()
            p.join()


def crop_image_isolated(input_img_bytes, bbox):
    """
    Crop image in isolated subprocess.

    Args:
        input_img_bytes: Input image as PNG bytes
        bbox: (left, top, right, bottom) tuple

    Returns:
        bytes: Cropped image as PNG bytes, or None if failed
    """
    parent_conn, child_conn = multiprocessing.Pipe()

    # Start subprocess
    p = multiprocessing.Process(
        target=_worker_crop_image,
        args=(child_conn, bbox)
    )
    p.start()

    try:
        # Send image data to subprocess
        parent_conn.send(input_img_bytes)

        # Wait for result
        cropped_bytes = parent_conn.recv()
        return cropped_bytes

    finally:
        parent_conn.close()
        p.join(timeout=30)  # 30 second timeout
        if p.is_alive():
            logger.warning("Image crop subprocess timed out, terminating")
            p.terminate()
            p.join()