File size: 5,860 Bytes

---
license: apache-2.0
tags:
- computer-vision
- text-styling
- ocr-attributes
- edgenext
- timm
- onnx
- multi-task
pipeline_tag: image-feature-extraction
library_name: timm
base_model:
- timm/edgenext_small.usi_in1k
base_model_relation: finetune
---

# Text Styling Multi-Task Model (EdgeNeXt-Small)

This model performs **fine-grained text attribute and styling extraction** from cropped text line images. Built upon a lightweight [`edgenext_small`](https://github.com/huggingface/pytorch-image-models) backbone, the model processes an input image of fixed dimensions ($64 \times 160$) and simultaneously predicts typography styles, background modalities, and color attributes across 8 dedicated heads.

---

## 📋 Model Summary

- **Backbone**: `edgenext_small` (pretrained via `timm`)
- **Spatial Pooling**: Adaptive Average Pooling to $(2 \times 5)$ ($304 \times 2 \times 5 = 3040\text{d}$)
- **Bottleneck**: `Linear(3040, 256) -> BatchNorm1d -> Hardswish -> Dropout(0.2)`
- **Input Resolution**: $64 \times 160$ (Height $\times$ Width), 3 RGB Channels
- **Export Format**: PyTorch Checkpoints (`.pth.tar`) & ONNX (`opset 18`, dynamic batch size)

---

## 🎯 Prediction Heads & Output Specification

The model outputs **8 multi-task tensors**:

| Output Name | Shape | Range / Activation | Description |
| :--- | :--- | :--- | :--- |
| `flags` | `[B, 5]` | Raw Logits $\to \sigma(x) \in [0, 1]$ | Binary flags: `[is_bold, is_italic, has_stroke, has_shadow, has_glow]` |
| `bg_type` | `[B, 3]` | Raw Logits $\to \text{Softmax}$ | Classification: `0: Solid`, `1: Gradient`, `2: Artwork/Image` |
| `text_color` | `[B, 3]` | Clamped $[0.0, 1.0]$ | Primary text fill color (RGB) |
| `effect_color` | `[B, 3]` | Clamped $[0.0, 1.0]$ | Outer effect color (Stroke / Shadow / Glow RGB) *(Valid if effect flags $> 0.5$)* |
| `bg_color` | `[B, 3]` | Clamped $[0.0, 1.0]$ | Solid background color (RGB) *(Valid if `bg_type == 0`)* |
| `bg_color_a` | `[B, 3]` | Clamped $[0.0, 1.0]$ | Gradient start color (RGB) *(Valid if `bg_type == 1`)* |
| `bg_color_b` | `[B, 3]` | Clamped $[0.0, 1.0]$ | Gradient end color (RGB) *(Valid if `bg_type == 1`)* |
| `bg_direction` | `[B, 2]` | Unit Vector $[\sin \theta, \cos \theta]$ | Gradient angle direction *(Valid if `bg_type == 1`)* |

---

## 🚀 Quickstart Inference

### ONNX Runtime (Recommended for Deployment)

```python
import numpy as np
import onnxruntime as ort
from PIL import Image

# Initialize Session
session = ort.InferenceSession("text_styling_model.onnx", providers=["CPUExecutionProvider"])

def preprocess(img_path):
    img = Image.open(img_path).convert("RGB").resize((160, 64), Image.BICUBIC)
    arr = np.array(img).astype(np.float32) / 255.0
    # Normalize with ImageNet mean/std
    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
    arr = (arr - mean) / std
    arr = np.transpose(arr, (2, 0, 1))  # HWC to CHW
    return np.expand_dims(arr, axis=0)  # Add batch dim [1, 3, 64, 160]

# Run Inference
input_tensor = preprocess("sample_text.png")
outputs = session.run(None, {"input": input_tensor})

(pred_flags, pred_bg_type, pred_text_c, pred_effect_c, 
 pred_bg_c, pred_bg_ca, pred_bg_cb, pred_bg_dir) = outputs

# Post-process Flags
sigmoid = lambda x: 1 / (1 + np.exp(-x))
flag_probs = sigmoid(pred_flags[0])
flag_names = ["Bold", "Italic", "Stroke", "Shadow", "Glow"]
detected_flags = {name: bool(prob > 0.5) for name, prob in zip(flag_names, flag_probs)}

# Post-process Background Type
bg_types = ["Solid", "Gradient", "Artwork"]
bg_type_idx = int(np.argmax(pred_bg_type[0]))

# Post-process Gradient Angle (if gradient)
sin_val, cos_val = pred_bg_dir[0]
angle_rad = np.arctan2(sin_val, cos_val)
angle_deg = (np.degrees(angle_rad) + 360) % 360

print(f"Flags: {detected_flags}")
print(f"Background Type: {bg_types[bg_type_idx]}")
print(f"Text Color (RGB [0-255]): {(pred_text_c[0] * 255).astype(int).tolist()}")

if any([detected_flags["Stroke"], detected_flags["Shadow"], detected_flags["Glow"]]):
    print(f"Effect Color (RGB [0-255]): {(pred_effect_c[0] * 255).astype(int).tolist()}")

if bg_type_idx == 0:
    print(f"Solid BG Color (RGB): {(pred_bg_c[0] * 255).astype(int).tolist()}")
elif bg_type_idx == 1:
    print(f"Gradient Start Color (RGB): {(pred_bg_ca[0] * 255).astype(int).tolist()}")
    print(f"Gradient End Color (RGB): {(pred_bg_cb[0] * 255).astype(int).tolist()}")
    print(f"Gradient Angle: {angle_deg:.1f}°")
```

---

## 🛠️ Training Details

- **Loss Function**: Multi-Task Loss combining:
  - Binary Cross Entropy with Logits for typography flags.
  - Cross Entropy for background category.
  - Smooth L1 (Huber Loss, $\beta=0.1$) with sample masking for RGB continuous targets.
  - Cosine Distance $(1 - \cos \theta)$ for periodic gradient angle direction.
- **Optimizer**: `AdamW` ($\text{lr}=10^{-3}$, weight decay $= 10^{-4}$) with gradient clipping (`max_norm=1.0`).
- **Scheduler**: Cosine Annealing with 1 warmup epoch.
- **Weight Averaging**: Step-count normalized Exponential Moving Average (`ModelEmaV2`).
- **Mixed Precision**: Distributed FP16 training with HuggingFace Accelerate across $2\times\text{NVIDIA T4}$.

---

## ⚠️ Intended Limitations & Edge Cases

1. **Aspect Ratio & Resolution**: Optimized for horizontal text strip crops rendered or resized to $64 \times 160$. Vertical text arrangements or extreme aspect ratios may degrade spatial boundary perception.
2. **Masked Heads Dependency**:
   - `effect_color` predictions should only be evaluated when one or more of `[Stroke, Shadow, Glow]` flags are positive.
   - `bg_color` is only valid when `bg_type == Solid (0)`.
   - `bg_color_a`, `bg_color_b`, and `bg_direction` are only valid when `bg_type == Gradient (1)`.
   - For `bg_type == Artwork (2)`, no background color regression is guaranteed.
```