Update README.md

README.md

@@ -1,3 +1,212 @@
 ---
-license: mit
+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
 ---
+
+# 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
+
+### 1. 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}°")
+```
+
+---
+
+### 2. PyTorch Definition & Inference
+
+```python
+import torch
+import torch.nn as nn
+from timm.models import create_model
+from timm.data.transforms_factory import transforms_imagenet_eval
+from PIL import Image
+
+class TextStylingModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.backbone = create_model('edgenext_small', pretrained=False, num_classes=0)
+        self.spatial_pool = nn.AdaptiveAvgPool2d((2, 5))
+        in_dim = 304 * 2 * 5
+
+        self.bottleneck = nn.Sequential(
+            nn.Linear(in_dim, 256),
+            nn.BatchNorm1d(256),
+            nn.Hardswish(),
+            nn.Dropout(0.2)
+        )
+
+        self.head_flags = nn.Linear(256, 5)
+        self.head_bg_type = nn.Linear(256, 3)
+        self.head_text_color = nn.Linear(256, 3)
+        self.head_effect_color = nn.Linear(256, 3)
+        self.head_bg_color = nn.Linear(256, 3)
+        self.head_bg_color_a = nn.Linear(256, 3)
+        self.head_bg_color_b = nn.Linear(256, 3)
+        self.head_bg_direction = nn.Linear(256, 2)
+
+    def forward(self, x):
+        x = self.backbone.forward_features(x)
+        x = self.spatial_pool(x)
+        x = torch.flatten(x, 1)
+        x = self.bottleneck(x)
+
+        flags = self.head_flags(x)
+        bg_type = self.head_bg_type(x)
+        text_color = torch.clamp(self.head_text_color(x), 0.0, 1.0)
+        effect_color = torch.clamp(self.head_effect_color(x), 0.0, 1.0)
+        bg_color = torch.clamp(self.head_bg_color(x), 0.0, 1.0)
+        bg_color_a = torch.clamp(self.head_bg_color_a(x), 0.0, 1.0)
+        bg_color_b = torch.clamp(self.head_bg_color_b(x), 0.0, 1.0)
+        
+        bg_direction = self.head_bg_direction(x)
+        dnorm = bg_direction.norm(dim=-1, keepdim=True).clamp(min=1e-6)
+        bg_direction = bg_direction / dnorm
+
+        return flags, bg_type, text_color, effect_color, bg_color, bg_color_a, bg_color_b, bg_direction
+
+# Load model weights
+model = TextStylingModel()
+checkpoint = torch.load("model_best.pth.tar", map_location="cpu")
+# If EMA weights exist, prioritize them:
+state_dict = checkpoint.get("state_dict_ema", checkpoint.get("state_dict"))
+model.load_state_dict(state_dict)
+model.eval()
+
+# Transform
+transform = transforms_imagenet_eval(img_size=(64, 160), crop_pct=1.0, interpolation='bicubic')
+img = Image.open("sample_text.png").convert("RGB")
+img_t = transform(img).unsqueeze(0)
+
+with torch.inference_mode():
+    preds = model(img_t)
+```
+
+---
+
+## 🛠️ 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.
+```
+
+---
+
+### Recommended Files to Upload to the Hugging Face Repository:
+1. `README.md` (The Model Card above)
+2. `text_styling_model.onnx` (The exported ONNX model)
+3. `model_best.pth.tar` (The PyTorch checkpoint containing weights and EMA states)
+4. *(Optional)* `sample_text.png` (A sample cropped image for testing)