Liiesl/bubble-segment-onnx

Koharu YOLO26s (1024x1024 ONNX Export)

This repository provides an optimized 1024×1024 static ONNX export of mayocream/koharu-yolo26s, an instance segmentation model designed for manga page element extraction (panels, dialogue text, speech bubbles, and onomatopoeia).

Re-exporting the model at $1024 \times 1024$ (instead of native $1280 \times 1280$) lowers activation memory by ~37% (peak memory dropping from ~600 MB to ~385 MB), making it significantly faster and lighter for desktop scanlation pipelines, edge engines, and memory-constrained environments (such as Rust/WASM/CPU runtimes).

---

📊 Comparison: 1280 vs 1024

Metric	Original (1280)	Re-export (1024)	Delta
Input Shape	[1, 3, 1280, 1280]	[1, 3, 1024, 1024]	-36% spatial pixels
Proto Mask Output	[1, 32, 320, 320] (50.0 MB)	[1, 32, 256, 256] (32.0 MB)	-36% activation size
ONNX Opset	Opset 12	Opset 17	Modern node support
File Size	40.29 MB	~40.29 MB	Identical (FP32 weights unchanged)
Peak Activation RAM	~600 MB	~385 MB	-37% RAM reduction
Parameter Count	11.437M	11.437M	Exact original weights

---

🏷️ Classes & Detection Targets

The model detects and segments 4 manga page element classes:

Class ID	Class Name	Description
0	frame	Comic panels / page framing borders
1	dialogue_text	Vertical and horizontal spoken dialogue text lines
2	balloon	Speech / thought bubble contours
3	onomatopoeia_text	Sound effects (SFX) / stylized sound text

---

📐 Model Inputs & Outputs

Input Tensor

• images: [1, 3, 1024, 1024] (Float32, RGB channel order, normalized to [0.0, 1.0], padded with fill value 114/255).

Output Tensors

• Output 0 (output0): [1, 300, 38] (Bounding boxes, confidence scores, 4 class probabilities, and 32 prototype mask coefficients).
• Output 1 (output1): [1, 32, 256, 256] (Prototype masks at $1/4$ input resolution: $256 \times 256$).

---

🚀 Quickstart Inference (Python + ONNX Runtime)

import cv2
import numpy as np
import onnxruntime as ort

# 1. Load Session
session = ort.InferenceSession("koharu-yolo26s-1024.onnx", providers=["CPUExecutionProvider"])
input_name = session.get_inputs()[0].name

# 2. Letterbox Preprocessing (Padded Resize to 1024x1024)
def letterbox(img, size=1024):
    h, w = img.shape[:2]
    r = min(size / w, size / h)
    nw, nh = int(round(w * r)), int(round(h * r))
    resized = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_LINEAR)
    
    # Fill canvas with 114 (YOLO standard gray)
    canvas = np.full((size, size, 3), 114, dtype=np.uint8)
    dx = (size - nw) // 2
    dy = (size - nh) // 2
    canvas[dy:dy + nh, dx:dx + nw] = resized
    return canvas, r, dx, dy

# 3. Prepare Image
raw_img = cv2.imread("manga_page.jpg")
raw_rgb = cv2.cvtColor(raw_img, cv2.COLOR_BGR2RGB)
canvas, ratio, dx, dy = letterbox(raw_rgb, size=1024)

# Normalize & transpose to NCHW
inp = np.transpose(canvas.astype(np.float32) / 255.0, (2, 0, 1))[None, ...]

# 4. Inference
outputs = session.run(None, {input_name: inp})
dets, proto = outputs[0], outputs[1]  # dets: [1, 300, 38], proto: [1, 32, 256, 256]

# 5. Parse Detections
CLASS_NAMES = ["frame", "dialogue_text", "balloon", "onomatopoeia_text"]
CONF_THRESHOLD = 0.25

detections = dets[0]  # shape: [300, 38]
for det in detections:
    conf = det[4]
    if conf < CONF_THRESHOLD:
        continue
    
    class_id = int(np.argmax(det[5:9]))
    class_name = CLASS_NAMES[class_id]
    
    # Extract Bounding Box (cx, cy, w, h in 1024 space -> unpad back to original coords)
    cx, cy, w, h = det[0], det[1], det[2], det[3]
    x1 = (cx - w / 2 - dx) / ratio
    y1 = (cy - h / 2 - dy) / ratio
    x2 = (cx + w / 2 - dx) / ratio
    y2 = (cy + h / 2 - dy) / ratio
    
    print(f"Detected {class_name} [{conf:.2f}]: ({x1:.1f}, {y1:.1f}, {x2:.1f}, {y2:.1f})")

---

⚙️ Integration Notes for Rust / Custom C++ Engines

If migrating downstream engines (such as scanlateit-segment) from 1280 to 1024:

1. Input Size Constant: Update IMG_SIZE from 1280 to 1024.
2. Prototype Grid Dimensions: Expect proto dimensions $256 \times 256$ ($1024 / 4$) instead of $320 \times 320$ ($1280 / 4$).
3. Letterbox Scale: Adjust target center offsets (dx, dy) and aspect scaling factor $r = \min(1024/w, 1024/h)$.

---

📚 Acknowledgements & References

• Upstream Model Weights: mayocream/koharu-yolo26s
• Resolution Scaling Concept: mednasserallah/manga109-segmentation-bubble-onnx
• Underlying YOLO Framework: Ultralytics YOLO