From 769371a1a34584c1a1708d4b8abf03f8179aa3d5 Mon Sep 17 00:00:00 2001 From: Siwat Sirichai Date: Tue, 15 Jul 2025 00:30:09 +0700 Subject: [PATCH] feat: integrate Redis support in pipeline execution; add actions for saving images and publishing messages --- pympta.md | 200 ++++++++++++++++++++++++++++++++++++++++++ requirements.txt | 3 +- siwatsystem/pympta.py | 52 ++++++++++- 3 files changed, 250 insertions(+), 5 deletions(-) create mode 100644 pympta.md diff --git a/pympta.md b/pympta.md new file mode 100644 index 0000000..2136480 --- /dev/null +++ b/pympta.md @@ -0,0 +1,200 @@ +# pympta: Modular Pipeline Task Executor + +`pympta` is a Python module designed to load and execute modular, multi-stage AI pipelines defined in a special package format (`.mpta`). It is primarily used within the detector worker to run complex computer vision tasks where the output of one model can trigger a subsequent model on a specific region of interest. + +## Core Concepts + +### 1. MPTA Package (`.mpta`) + +An `.mpta` file is a standard `.zip` archive with a different extension. It bundles all the necessary components for a pipeline to run. + +A typical `.mpta` file has the following structure: + +``` +my_pipeline.mpta/ +├── pipeline.json +├── model1.pt +├── model2.pt +└── ... +``` + +- **`pipeline.json`**: (Required) The manifest file that defines the structure of the pipeline, the models to use, and the logic connecting them. +- **Model Files (`.pt`, etc.)**: The actual pre-trained model files (e.g., PyTorch, ONNX). The pipeline currently uses `ultralytics.YOLO` models. + +### 2. Pipeline Structure + +A pipeline is a tree-like structure of "nodes," defined in `pipeline.json`. + +- **Root Node**: The entry point of the pipeline. It processes the initial, full-frame image. +- **Branch Nodes**: Child nodes that are triggered by specific detection results from their parent. For example, a root node might detect a "vehicle," which then triggers a branch node to detect a "license plate" within the vehicle's bounding box. + +This modular structure allows for creating complex and efficient inference logic, avoiding the need to run every model on every frame. + +## `pipeline.json` Specification + +This file defines the entire pipeline logic. The root object contains a `pipeline` key for the pipeline definition and an optional `redis` key for Redis configuration. + +### Top-Level Object Structure + +| Key | Type | Required | Description | +| ---------- | ------ | -------- | ------------------------------------------------------- | +| `pipeline` | Object | Yes | The root node object of the pipeline. | +| `redis` | Object | No | Configuration for connecting to a Redis server. | + +### Redis Configuration (`redis`) + +| Key | Type | Required | Description | +| ---------- | ------ | -------- | ------------------------------------------------------- | +| `host` | String | Yes | The hostname or IP address of the Redis server. | +| `port` | Number | Yes | The port number of the Redis server. | +| `password` | String | No | The password for Redis authentication. | +| `db` | Number | No | The Redis database number to use. Defaults to `0`. | + +### Node Object Structure + +| Key | Type | Required | Description | +| ------------------- | ------------- | -------- | -------------------------------------------------------------------------------------------------------------------------------------- | +| `modelId` | String | Yes | A unique identifier for this model node (e.g., "vehicle-detector"). | +| `modelFile` | String | Yes | The path to the model file within the `.mpta` archive (e.g., "yolov8n.pt"). | +| `minConfidence` | Float | Yes | The minimum confidence score (0.0 to 1.0) required for a detection to be considered valid and potentially trigger a branch. | +| `triggerClasses` | Array | Yes | A list of class names that, when detected by the parent, can trigger this node. For the root node, this lists all classes of interest. | +| `crop` | Boolean | No | If `true`, the image is cropped to the parent's detection bounding box before being passed to this node's model. Defaults to `false`. | +| `branches` | Array | No | A list of child node objects that can be triggered by this node's detections. | +| `actions` | Array | No | A list of actions to execute upon a successful detection in this node. | + +### Action Object Structure + +Actions allow the pipeline to interact with Redis. + +#### `redis_save_image` + +Saves the current image frame (or cropped sub-image) to a Redis key. + +| Key | Type | Required | Description | +| ----- | ------ | -------- | ------------------------------------------------------------------------------------------------------- | +| `type`| String | Yes | Must be `"redis_save_image"`. | +| `key` | String | Yes | The Redis key to save the image to. Can contain placeholders like `{class}` or `{id}` to be formatted with detection results. | + +#### `redis_publish` + +Publishes a message to a Redis channel. + +| Key | Type | Required | Description | +| --------- | ------ | -------- | ------------------------------------------------------------------------------------------------------- | +| `type` | String | Yes | Must be `"redis_publish"`. | +| `channel` | String | Yes | The Redis channel to publish the message to. | +| `message` | String | Yes | The message to publish. Can contain placeholders like `{class}` or `{id}` to be formatted with detection results. | + +### Example `pipeline.json` with Redis + +```json +{ + "redis": { + "host": "localhost", + "port": 6379, + "password": "your-password" + }, + "pipeline": { + "modelId": "vehicle-detector", + "modelFile": "vehicle_model.pt", + "minConfidence": 0.5, + "triggerClasses": ["car", "truck"], + "actions": [ + { + "type": "redis_save_image", + "key": "detection:image:{id}" + }, + { + "type": "redis_publish", + "channel": "detections", + "message": "Detected a {class} with ID {id}" + } + ], + "branches": [ + { + "modelId": "lpr-us", + "modelFile": "lpr_model.pt", + "minConfidence": 0.7, + "triggerClasses": ["car"], + "crop": true, + "branches": [] + } + ] + } +} +``` + +## API Reference + +The `pympta` module exposes two main functions. + +### `load_pipeline_from_zip(zip_source: str, target_dir: str) -> dict` + +Loads, extracts, and parses an `.mpta` file to build a pipeline tree in memory. It also establishes a Redis connection if configured in `pipeline.json`. + +- **Parameters:** + - `zip_source` (str): The file path to the local `.mpta` zip archive. + - `target_dir` (str): A directory path where the archive's contents will be extracted. +- **Returns:** + - A dictionary representing the root node of the pipeline, ready to be used with `run_pipeline`. Returns `None` if loading fails. + +### `run_pipeline(frame, node: dict, return_bbox: bool = False)` + +Executes the inference pipeline on a single image frame. + +- **Parameters:** + - `frame`: The input image frame (e.g., a NumPy array from OpenCV). + - `node` (dict): The pipeline node to execute (typically the root node returned by `load_pipeline_from_zip`). + - `return_bbox` (bool): If `True`, the function returns a tuple `(detection, bounding_box)`. Otherwise, it returns only the `detection`. +- **Returns:** + - The final detection result from the last executed node in the chain. A detection is a dictionary like `{'class': 'car', 'confidence': 0.95, 'id': 1}`. If no detection meets the criteria, it returns `None` (or `(None, None)` if `return_bbox` is `True`). + +## Usage Example + +This snippet, inspired by `pipeline_webcam.py`, shows how to use `pympta` to load a pipeline and process an image from a webcam. + +```python +import cv2 +from siwatsystem.pympta import load_pipeline_from_zip, run_pipeline + +# 1. Define paths +MPTA_FILE = "path/to/your/pipeline.mpta" +CACHE_DIR = ".mptacache" + +# 2. Load the pipeline from the .mpta file +# This reads pipeline.json and loads the YOLO models into memory. +model_tree = load_pipeline_from_zip(MPTA_FILE, CACHE_DIR) + +if not model_tree: + print("Failed to load pipeline.") + exit() + +# 3. Open a video source +cap = cv2.VideoCapture(0) + +while True: + ret, frame = cap.read() + if not ret: + break + + # 4. Run the pipeline on the current frame + # The function will handle the entire logic tree (e.g., find a car, then find its license plate). + detection_result, bounding_box = run_pipeline(frame, model_tree, return_bbox=True) + + # 5. Display the results + if detection_result: + print(f"Detected: {detection_result['class']} with confidence {detection_result['confidence']:.2f}") + if bounding_box: + x1, y1, x2, y2 = bounding_box + cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2) + cv2.putText(frame, detection_result['class'], (x1, y1 - 10), + cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36, 255, 12), 2) + + cv2.imshow("Pipeline Output", frame) + + if cv2.waitKey(1) & 0xFF == ord('q'): + break + +cap.release() +cv2.destroyAllWindows() +``` \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 84f45cc..49ca601 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,4 +5,5 @@ torchvision ultralytics opencv-python websockets -fastapi[standard] \ No newline at end of file +fastapi[standard] +redis \ No newline at end of file diff --git a/siwatsystem/pympta.py b/siwatsystem/pympta.py index 5e32596..4514182 100644 --- a/siwatsystem/pympta.py +++ b/siwatsystem/pympta.py @@ -7,13 +7,14 @@ import requests import zipfile import shutil import traceback +import redis from ultralytics import YOLO from urllib.parse import urlparse # Create a logger specifically for this module logger = logging.getLogger("detector_worker.pympta") -def load_pipeline_node(node_config: dict, mpta_dir: str) -> dict: +def load_pipeline_node(node_config: dict, mpta_dir: str, redis_client) -> dict: # Recursively load a model node from configuration. model_path = os.path.join(mpta_dir, node_config["modelFile"]) if not os.path.exists(model_path): @@ -44,13 +45,15 @@ def load_pipeline_node(node_config: dict, mpta_dir: str) -> dict: "triggerClassIndices": trigger_class_indices, "crop": node_config.get("crop", False), "minConfidence": node_config.get("minConfidence", None), + "actions": node_config.get("actions", []), "model": model, - "branches": [] + "branches": [], + "redis_client": redis_client } logger.debug(f"Configured node {node_config['modelId']} with trigger classes: {node['triggerClasses']}") for child in node_config.get("branches", []): logger.debug(f"Loading branch for parent node {node_config['modelId']}") - node["branches"].append(load_pipeline_node(child, mpta_dir)) + node["branches"].append(load_pipeline_node(child, mpta_dir, redis_client)) return node def load_pipeline_from_zip(zip_source: str, target_dir: str) -> dict: @@ -158,7 +161,26 @@ def load_pipeline_from_zip(zip_source: str, target_dir: str) -> dict: pipeline_config = json.load(f) logger.info(f"Successfully loaded pipeline configuration from {pipeline_json_path}") logger.debug(f"Pipeline config: {json.dumps(pipeline_config, indent=2)}") - return load_pipeline_node(pipeline_config["pipeline"], mpta_dir) + + # Establish Redis connection if configured + redis_client = None + if "redis" in pipeline_config: + redis_config = pipeline_config["redis"] + try: + redis_client = redis.Redis( + host=redis_config["host"], + port=redis_config["port"], + password=redis_config.get("password"), + db=redis_config.get("db", 0), + decode_responses=True + ) + redis_client.ping() + logger.info(f"Successfully connected to Redis at {redis_config['host']}:{redis_config['port']}") + except redis.exceptions.ConnectionError as e: + logger.error(f"Failed to connect to Redis: {e}") + redis_client = None + + return load_pipeline_node(pipeline_config["pipeline"], mpta_dir, redis_client) except json.JSONDecodeError as e: logger.error(f"Error parsing pipeline.json: {str(e)}", exc_info=True) return None @@ -169,6 +191,25 @@ def load_pipeline_from_zip(zip_source: str, target_dir: str) -> dict: logger.error(f"Error loading pipeline.json: {str(e)}", exc_info=True) return None +def execute_actions(node, frame, detection_result): + if not node["redis_client"] or not node["actions"]: + return + + for action in node["actions"]: + try: + if action["type"] == "redis_save_image": + key = action["key"].format(**detection_result) + _, buffer = cv2.imencode('.jpg', frame) + node["redis_client"].set(key, buffer.tobytes()) + logger.info(f"Saved image to Redis with key: {key}") + elif action["type"] == "redis_publish": + channel = action["channel"] + message = action["message"].format(**detection_result) + node["redis_client"].publish(channel, message) + logger.info(f"Published message to Redis channel '{channel}': {message}") + except Exception as e: + logger.error(f"Error executing action {action['type']}: {e}") + def run_pipeline(frame, node: dict, return_bbox: bool=False): """ - For detection nodes (task != 'classify'): @@ -206,6 +247,7 @@ def run_pipeline(frame, node: dict, return_bbox: bool=False): "confidence": top1_conf, "id": None } + execute_actions(node, frame, det) return (det, None) if return_bbox else det @@ -254,9 +296,11 @@ def run_pipeline(frame, node: dict, return_bbox: bool=False): det2, _ = run_pipeline(sub, br, return_bbox=True) if det2: # return classification result + original bbox + execute_actions(br, sub, det2) return (det2, best_box) if return_bbox else det2 # ─── No branch matched → return this detection ───────────── + execute_actions(node, frame, best_det) return (best_det, best_box) if return_bbox else best_det except Exception as e: