Liiesl/text-styling-classificationv1

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 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)

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.