Writing your modules

[bropines]

API Documentation for Module Creation

Introduction

This guide describes how to create and integrate custom modules into the application. The system features a modular architecture that allows developers to extend its functionality by adding new algorithms for the following tasks:

1. Text Detection (TextDetector): Finding text regions within an image.
2. Text Recognition (OCR): Extracting text from the found regions.
3. Inpainting (Inpainter): Removing the original text from the image.
4. Translation (Translator): Translating the extracted text.

Key Concepts

1. Module Registry

The system uses a "Registry" pattern to discover and manage modules. For your extension to become available in the program, it must be registered using a special decorator.

• Place your .py file in the corresponding directory (modules/textdetector, modules/ocr, etc.).
• Add the decorator above your class definition. The name you provide in the decorator will be displayed in the user interface.

Module Type	Directory	Decorator
Text Detector	modules/textdetector	@register_textdetectors('my-detector')
OCR	modules/ocr	@register_OCR('my-ocr')
Inpainter	modules/inpaint	@register_inpainter('my-inpainter')
Translator	modules/translators	@register_translator('my-translator')

2. The BaseModule Base Class

All modules must inherit from a base class corresponding to their type (TextDetectorBase, OCRBase, etc.), which in turn inherit from modules.base.BaseModule. This class provides a common interface and functionality.

Key Class Attributes:

• params (dict): A dictionary defining the module's parameters that will be user-configurable. Structure:

  params = {
      'my_parameter': {
          'type': 'selector',  # can be 'checkbox', 'editor', or just a default value
          'options': ['option1', 'option2'],
          'value': 'option1',
          'description': 'A description for the parameter.'
      },
      'simple_param': 5.0 # A simple parameter with a default value
  }

• download_file_list (list): A list of dictionaries describing files (e.g., models) that need to be downloaded for the module to work. The system will automatically handle the download and verification process.

  download_file_list = [{
      'url': 'https://example.com/models/my_model.ckpt',
      'sha256_pre_calculated': '...', # Optional hash for verification
      'files': 'data/models/my_model_folder/my_model.ckpt', # Save path
  }]

• _load_model_keys (set): A set of attribute names that store large objects (models). This helps the system manage memory by loading and unloading models as needed.

  _load_model_keys = {'model', 'processor'} # Refers to self.model and self.processor

Core Methods to Implement:

• _load_model(): This is where you should load your heavy models and assign them to the attributes listed in _load_model_keys. It is called automatically the first time the module is used.
• unload_model(): The system manages unloading automatically via _load_model_keys, but you can override this method for more complex cleanup logic.

3. Inheriting from Specialized Base Classes

Each module type has its own base class that inherits from BaseModule and defines the primary method you must implement.

---

How to Create Your Module: Step-by-Step Guides

1. Creating a Custom Text Detector

A text detector scans an image for text blocks and returns their mask and coordinates.

1. File Location: modules/textdetector/my_detector.py
2. Inheritance: Your class must inherit from TextDetectorBase.
3. Registration: Use the @register_textdetectors('my-detector') decorator.
4. Core Method: Implement the _detect(self, img, proj) method.
   • Input:
     • img (np.ndarray): The image in BGR format (3 channels).
     • proj (ProjImgTrans): A context object (for advanced use).
   • Output: Tuple[np.ndarray, List[TextBlock]]
     • mask (np.ndarray): A single-channel mask (0-255), where white areas denote detected text.
     • blk_list (List[TextBlock]): A list of TextBlock objects, one for each detected text region.

Example (A simple detector based on thresholding):

import cv2
import numpy as np
from typing import Tuple, List
from modules.textdetector.base import register_textdetectors, TextDetectorBase, TextBlock, ProjImgTrans

@register_textdetectors('simple-threshold-detector')
class SimpleDetector(TextDetectorBase):

    def _load_model(self):
        # Model loading would happen here if needed
        pass

    def _detect(self, img: np.ndarray, proj: ProjImgTrans) -> Tuple[np.ndarray, List[TextBlock]]:
        # Convert to grayscale and apply a threshold
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        _, mask = cv2.threshold(gray, 128, 255, cv2.THRESH_BINARY_INV)

        # Find contours
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        blk_list = []
        for cnt in contours:
            x, y, w, h = cv2.boundingRect(cnt)
            # Create a TextBlock for each contour
            block = TextBlock(xyxy=[x, y, x + w, y + h])
            blk_list.append(block)

        return mask, blk_list

2. Creating a Custom OCR Module

An OCR module takes an image and a list of text blocks and fills them with recognized text.

1. File Location: modules/ocr/my_ocr.py
2. Inheritance: Your class must inherit from OCRBase.
3. Registration: Use the @register_OCR('my-ocr') decorator.
4. Core Method: Implement the _ocr_blk_list(self, img, blk_list) method.
   • Input:
     • img (np.ndarray): The original image in RGB format.
     • blk_list (List[TextBlock]): A list of text blocks to recognize.
   • Action: You must process each TextBlock, crop the corresponding region from img, recognize the text, and write it to blk.text.

Example (A dummy OCR that returns coordinates):

import numpy as np
from typing import List
from modules.ocr.base import register_OCR, OCRBase, TextBlock

@register_OCR('dummy-ocr')
class DummyOCR(OCRBase):

    def _load_model(self):
        # No model to load
        pass

    def _ocr_blk_list(self, img: np.ndarray, blk_list: List[TextBlock]):
        im_h, im_w = img.shape[:2]
        for blk in blk_list:
            x1, y1, x2, y2 = blk.xyxy
            # Ensure coordinates are valid
            if 0 <= y1 < y2 <= im_h and 0 <= x1 < x2 <= im_w:
                # Crop the image using the block's coordinates
                cropped_img = img[y1:y2, x1:x2]
                # Instead of real OCR, just write info about the block
                blk.text = f"OCR for block at ({x1},{y1}) with size {cropped_img.shape}x{cropped_img.shape}"
            else:
                blk.text = ""

3. Creating a Custom Inpainter Module

An Inpainter takes an image and a text mask, then fills in the text regions so they can be replaced with translated text.

1. File Location: modules/inpaint/my_inpainter.py
2. Inheritance: Your class must inherit from InpainterBase.
3. Registration: Use the @register_inpainter('my-inpainter') decorator.
4. Core Method: Implement the _inpaint(self, img, mask, textblock_list) method.
   • Input:
     • img (np.ndarray): The RGB image to be processed.
     • mask (np.ndarray): A single-channel mask where white areas correspond to text that needs to be inpainted.
     • textblock_list (List[TextBlock]): A list of text blocks (can be None).
   • Output: np.ndarray - the image with text inpainted.

Example (Using cv2.inpaint):

import cv2
import numpy as np
from typing import List
from modules.inpaint.base import register_inpainter, InpainterBase
from modules.textdetector.base import TextBlock

@register_inpainter('simple-cv2-inpainter')
class SimpleCVInpainter(InpainterBase):
    # This module will not process blocks individually
    inpaint_by_block = False

    def _load_model(self):
        pass

    def _inpaint(self, img: np.ndarray, mask: np.ndarray, textblock_list: List[TextBlock] = None) -> np.ndarray:
        # Ensure the mask is single-channel
        if mask.ndim == 3:
            mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)

        # Use the Telea inpainting method from OpenCV
        inpainted_img = cv2.inpaint(img, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
        return inpainted_img

    def moveToDevice(self, device: str, precision: str = None):
        # This module runs on the CPU, so this method is empty
        pass

4. Creating a Custom Translator Module

A Translator takes a list of strings and returns a list of translated strings.

1. File Location: modules/translators/my_translator.py
2. Inheritance: Your class must inherit from BaseTranslator.
3. Registration: Use the @register_translator('my-translator') decorator.
4. Setup: Implement _setup_translator(). Here, you must populate self.lang_map with the languages your API supports and their corresponding codes.
5. Core Method: Implement _translate(self, src_list).
   • Input: src_list (List[str]) - a list of strings to translate.
   • Output: List[str] - a list of translated strings. The length must match src_list.

Example (A "translator" that reverses strings):

from typing import List, Dict
from modules.translators.base import BaseTranslator, register_translator

@register_translator('reverser')
class ReverserTranslator(BaseTranslator):
    # This translator does not concatenate strings
    concate_text = False

    def _setup_translator(self):
        # Our "translator" supports any language
        self.lang_map['简体中文'] = 'zh'
        self.lang_map['日本語'] = 'ja'
        self.lang_map['English'] = 'en'
        # ... and so on for all languages in LANGMAP_GLOBAL

    def _translate(self, src_list: List[str]) -> List[str]:
        # Get the current source and target languages
        source_lang_code = self.lang_map[self.lang_source]
        target_lang_code = self.lang_map[self.lang_target]

        self.logger.info(f"Reversing text from {source_lang_code} to {target_lang_code}")

        translated_list = []
        for text in src_list:
            reversed_text = text[::-1]
            translated_list.append(reversed_text)

        return translated_list

Best Practices: On-Demand Client Initialization

When creating modules that interact with external APIs or load heavy models (e.g., OpenAI, httpx.Client, GPU models), it is strongly recommended not to create the client instance within the __init__ or _setup_translator methods.

The Problem: Premature Initialization

Initializing a client in the constructor (__init__) can lead to several issues:

1. Wasted Resources: The client is created when the application starts, even if the user never selects your module for a task.
2. Stale Parameters: If a user changes module parameters (like an API key, endpoint, or proxy) through the UI after initialization, your client will continue to use the old, incorrect settings.
3. Slower Startup: Initializing many network clients at once can slow down the application's launch time.

The Solution: Lazy Initialization

The correct approach is to create the client only at the moment it is actually needed. This ensures that the most current parameters are used and that resources are not wasted.

How to Implement It:

1. In your class's __init__ method, set the client attribute to None.

   def __init__(self, **params) -> None:
       super().__init__(**params)
       self.client = None

2. Create a separate private method (e.g., _initialize_client) that handles the creation and configuration of the client instance, using the current values from self.params.
3. In your main operational method (_translate, ocr, etc.), add a check: if the client has not been created yet (self.client is None), call your initialization method.
4. In the updateParam method, reset the client to None whenever a key parameter (API key, endpoint, proxy) is changed. This forces the system to re-create the client with the new settings on the next call.

Example: Refactoring llm_api_translator

Here is how this pattern is applied in the LLM_OCR module (file ocr\ocr_llm_api.py):

# ... imports

@register_OCR("llm_ocr")
class LLM_OCR(OCRBase):
    # ... params

    def __init__(self, **params) -> None:
        super().__init__(**params)
        # 1. Initialize the client as None
        self.client = None
        # ... other attributes

    def _initialize_client(self):
        # 2. A method to create the client with current parameters
        # Here, we fetch the up-to-date settings from self.params
        api_key_to_use = self.get_param_value('api_key')
        endpoint = self.get_param_value('endpoint')
        proxy = self.get_param_value('proxy')
        
        # Client creation logic...
        self.client = OpenAI(
            api_key=api_key_to_use,
            base_url=endpoint,
            http_client=httpx.Client(transport=...) if proxy else None
        )
        self.logger.info("LLM client initialized.")

    def ocr(self, img_base64: str, prompt_override: str = None) -> str:
        # 3. Check and create the client "on-demand"
        if self.client is None:
            self.logger.debug("Client is not initialized. Initializing now.")
            self._initialize_client()

        # ... the rest of the ocr logic now safely uses self.client
        self._respect_delay()
        # ...

    def updateParam(self, param_key: str, param_content):
        super().updateParam(param_key, param_content)
        
        # 4. Reset the client when important parameters change
        if param_key in ["api_key", "endpoint", "proxy", "provider"]:
            self.client = None
            self.logger.info(f"Parameter '{param_key}' changed. Client will be re-initialized on next request.")

By following this pattern, you will make your modules more efficient, robust, and responsive to user configuration changes.

Gemini did almost all the work on the preparation of the documentation. Thanks to him)

Version for wibe-coding. Warning, the most stupid modules, which will not even be checked, but simply made so that we can check them later, will be immediately rejected. Spend time debugging and checking.

[bropines]

Defining User-Configurable Parameters: How the UI Interprets params

The params dictionary in your class is not just a data store; it's a direct blueprint for automatically building the settings interface for your module. The UI parses this dictionary and creates corresponding controls based on its structure. Understanding what keys and values the UI expects is crucial for creating a functional and user-friendly interface.

General Structure

A parameter can be defined in one of two ways:

1. Simple (Key-Value): If the value is a basic type (str, int, float, bool), the UI will create a simple input field or checkbox.
2. Complex (Dictionary): If the value is a dictionary, the UI will look for a 'type' key to determine which complex control to create (e.g., a dropdown list or a multi-line editor).

Control Types

Below are all the supported control types and their configurations.

1. Simple Input Field (Text Field / Spin Box)

If 'type' is not specified, the UI infers the control type from the data type of the value.

• For str: A single-line text field is created.
• For int or float: A numeric input field (spin box) with up/down arrows is created.

Example:

params = {
    'api_key': '',  # Will create a text field for a string
    'delay': 1.5,   # Will create a numeric field for a float
    'retries': 3    # Will create a numeric field for an integer
}

2. Selector / Dropdown

• 'type': 'selector'

Creates a dropdown list for selection.

Dictionary Keys:

• options (List[str]): A list of strings that will be the choices in the dropdown.
• value (str): The default value, which must exist in the options list.

Example:

params = {
    'model': {
        'type': 'selector',
        'options': ['gpt-4o', 'gpt-4-turbo', 'gpt-3.5-turbo'],
        'value': 'gpt-4o',
        'description': 'Select the model to use.'
    }
}

3. Checkbox

• 'type': 'checkbox'

Creates a checkbox.

Dictionary Keys:

• value (bool): The default state (True for checked, False for unchecked).

Example:

params = {
    'low vram mode': {
        'type': 'checkbox',
        'value': False,
        'description': 'Enable this if you have low video memory.'
    }
}

4. Multi-line Text Editor

• 'type': 'editor'

Creates a large text area suitable for editing multi-line text, such as prompts or templates.

Dictionary Keys:

• value (str): The default text content.

Example:

params = {
    'system_prompt': {
        'type': 'editor',
        'value': 'You are a professional manga translator...',
        'description': 'The system prompt for the model.'
    }
}

5. Push Button

• 'type': 'pushbtn'

Creates a clickable button. It is typically used to trigger an action, like refreshing a token or updating a model list.

Dictionary Keys:

• display_name (str): The text that will appear on the button.
• description (str): The text for the tooltip.
• value: Usually an empty string, as the button's value is not stored.

Example:

params = {
    'update_token_btn': {
        'type': 'pushbtn',
        'value': '',
        'display_name': 'Update Token',
        'description': 'Click to fetch a new access token.'
    }
}

The action for the button is implemented in the updateParam method, where you check for the param_key.

6. Checkbox Group

• 'type': 'check_group'

Creates a group of related checkboxes.

Dictionary Keys:

• value (Dict[str, bool]): A dictionary where keys are the checkbox labels and values are their default states.

Example:

params = {
    'label': {
        'type': 'check_group',
        'value': {  
            'vertical_textline': True, 
            'horizontal_textline': True, 
            'textblock': False # This type will be unchecked by default
        },
        'description': 'Select the types of text to detect.'
    }
}

7. File/Path Selector

This is an enhanced version of the 'selector'. It creates a text field with a "Browse..." button to select a file or folder.

Dictionary Keys:

• 'type': 'selector'
• 'path_selector': True (required to activate this mode)
• path_filter (str, optional): A file filter for the dialog window (e.g., '*.pt *.ckpt').
• value (str): The default path.

Example:

params = {
    'model path': {
        'type': 'selector',
        'options': ['path/to/default/model.pt'], # Can be one or more default paths
        'value': 'data/models/ysgyolo_S150best.pt',
        'path_selector': True,
        'path_filter': '*.pt *.ckpt *.pth',
        'description': 'Path to the model file.'
    }
}

Important Notes

• The description Key: Almost all control types support the 'description' key. Its value will be displayed as a tooltip when the user hovers over the control. Always add a description to make your settings clear and user-friendly.
• Accessing Values: Inside your module's code, access a parameter's value using self.get_param_value('my_parameter_key'). This ensures you are always working with the current value set by the user.
• Unknown Keys: If you add a key to a parameter's dictionary that the UI does not recognize (e.g., 'my_custom_key': 'hello'), it will simply be ignored and will not cause an error.