compat: drop imgaug, fix PL2/NumPy2/Unicode, add Marathi support

configs/charset/marathi.yaml

@@ -0,0 +1,5 @@
+# @package _global_
+
+model:
+  charset_train: ",-./:|ँंःअआइईउऊएऐऑओऔकखगघचछजझञटठडढणतथदधनपफबभमयरऱलळवशषसह़ािीुूृॅेैॉोौ्०१२३४५६७८९ॲ"
+  charset_test: ",-./:|ँंःअआइईउऊएऐऑओऔकखगघचछजझञटठडढणतथदधनपफबभमयरऱलळवशषसह़ािीुूृॅेैॉोौ्०१२३४५६७८९ॲ"

configs/dataset/marathi.yaml

@@ -0,0 +1,24 @@
+# @package _global_
+# Marathi scene-text dataset configuration.
+#
+# Expected LMDB layout (relative to data.root_dir):
+#   train/
+#     data.mdb        <- single LMDB  OR
+#     <subset>/       <- one LMDB per subset (build_tree_dataset handles both)
+#       data.mdb
+#   val/
+#     data.mdb
+#   test/
+#     data.mdb
+#
+# Override normalize_unicode=false if your LMDB labels are already NFC-normalised
+# (the dataset loader applies NFC automatically for non-ASCII charsets, so this
+#  is safe to leave at true).
+model:
+  charset_train: ",-./:|ँंःअआइईउऊएऐऑओऔकखगघचछजझञटठडढणतथदधनपफबभमयरऱलळवशषसह़ािीुूृॅेैॉोौ्०१२३४५६७८९ॲ"
+  charset_test: ",-./:|ँंःअआइईउऊएऐऑओऔकखगघचछजझञटठडढणतथदधनपफबभमयरऱलळवशषसह़ािीुूृॅेैॉोौ्०१२३४५६७८९ॲ"
+
+data:
+  train_dir: "."
+  remove_whitespace: true
+  normalize_unicode: true

configs/experiment/finetune_marathi.yaml

@@ -0,0 +1,40 @@
+# @package _global_
+# Staged fine-tuning configuration for Marathi (Devanagari) OCR.
+#
+# Experiment: freeze the visual encoder, train everything else with differential LRs.
+# To run:
+#   python train.py \
+#     experiment=finetune_marathi \
+#     charset=marathi \
+#     dataset=marathi \
+#     pretrained=parseq
+#
+# Adjust freeze flags and learning rates per experiment:
+#
+#   Experiment 1 — freeze encoder + decoder, train head only:
+#     model.freeze.encoder=true  model.freeze.decoder=true
+#     model.backbone_lr=1e-5     model.head_lr=5e-4
+#
+#   Experiment 2 — freeze encoder only:
+#     model.freeze.encoder=true
+#
+#   Experiment 3 — train everything (default behaviour, identical to no freeze):
+#     (just omit any freeze/lr overrides)
+defaults:
+  - override /model: parseq
+
+model:
+  # --- Layer freezing ---
+  # Set a component to true to freeze it (requires_grad=False).
+  # Omitting this section entirely preserves the original trainable behaviour.
+  freeze:
+    encoder:    true   # freeze the ViT backbone — reuse English visual features
+    decoder:    false  # keep decoder trainable for Marathi sequence modelling
+    head:       false  # always train the head (it was re-initialised for Marathi)
+    text_embed: false  # always train embeddings for the new charset
+
+  # --- Differential learning rates ---
+  # encoder + decoder use backbone_lr; head + text_embed use head_lr.
+  # Set both to null (or omit entirely) to use the single global lr instead.
+  backbone_lr: 1.0e-5
+  head_lr:     5.0e-4

configs/main.yaml

@@ -3,6 +3,7 @@ defaults:
   - model: parseq
   - charset: 94_full
   - dataset: real
+  - experiment: null
 
 model:
   _convert_: all

requirements/constraints.txt

@@ -84,13 +84,7 @@ idna==3.6
     #   requests
     #   yarl
 imageio==2.34.0
-    # via
-    #   imgaug
-    #   scikit-image
-imgaug==0.4.0
-    # via
-    #   -r requirements/train.in
-    #   -r requirements/tune.in
+    # via scikit-image
 importlib-resources==6.1.1
     # via matplotlib
 iopath==0.1.10
@@ -137,7 +131,7 @@ lmdb==1.4.1
 markupsafe==2.1.5
     # via jinja2
 matplotlib==3.8.3
-    # via imgaug
+    # via ax-platform
 matplotlib-inline==0.1.6
     # via ipython
 mpmath==1.3.0
@@ -153,9 +147,7 @@ multipledispatch==1.0.0
 mypy-extensions==1.0.0
     # via typing-inspect
 networkx==3.2.1
-    # via
-    #   scikit-image
-    #   torch
+    # via torch
 nltk==3.8.1
     # via -r requirements/core.in
 nodeenv==1.8.0
@@ -164,28 +156,22 @@ numpy==1.26.4
     # via
     #   contourpy
     #   fvcore
-    #   imageio
-    #   imgaug
     #   jaxtyping
     #   matplotlib
-    #   opencv-python
     #   opt-einsum
     #   pandas
     #   pyarrow
     #   pyro-ppl
     #   pytorch-lightning
-    #   scikit-image
     #   scikit-learn
     #   scipy
-    #   shapely
     #   tensorboardx
-    #   tifffile
     #   torchmetrics
     #   torchvision
 omegaconf==2.3.0
     # via hydra-core
 opencv-python==4.9.0.80
-    # via imgaug
+    # via -r requirements/bench.in
 opt-einsum==3.3.0
     # via pyro-ppl
 packaging==23.2
@@ -197,7 +183,6 @@ packaging==23.2
     #   plotly
     #   pytorch-lightning
     #   ray
-    #   scikit-image
     #   tensorboardx
     #   torchmetrics
 pandas==2.2.0
@@ -214,10 +199,7 @@ pillow==10.2.0
     #   -r requirements/train.in
     #   -r requirements/tune.in
     #   fvcore
-    #   imageio
-    #   imgaug
     #   matplotlib
-    #   scikit-image
     #   torchvision
 platformdirs==4.2.0
     # via virtualenv
@@ -286,8 +268,10 @@ rpds-py==0.18.0
     #   referencing
 safetensors==0.4.2
     # via timm
+lazy-loader==0.3
+    # via scikit-image
 scikit-image==0.22.0
-    # via imgaug
+    # via -r requirements/bench.in
 scikit-learn==1.4.1.post1
     # via
     #   ax-platform
@@ -296,16 +280,11 @@ scipy==1.12.0
     # via
     #   ax-platform
     #   botorch
-    #   imgaug
     #   linear-operator
-    #   scikit-image
     #   scikit-learn
-shapely==2.0.3
-    # via imgaug
 six==1.16.0
     # via
     #   asttokens
-    #   imgaug
     #   python-dateutil
 stack-data==0.6.3
     # via ipython

requirements/train.in

@@ -2,6 +2,5 @@
 
 lmdb
 Pillow
-imgaug
 hydra-core >=1.2.0
 tensorboardx

requirements/train.txt

@@ -3,27 +3,20 @@ contourpy==1.2.0
 cycler==0.12.1
 fonttools==4.49.0
 hydra-core==1.3.2
-imageio==2.34.0
-imgaug==0.4.0
 importlib-resources==6.1.1
 kiwisolver==1.4.5
-lazy-loader==0.3
 lmdb==1.4.1
 matplotlib==3.8.3
 networkx==3.2.1
 numpy==1.26.4
 omegaconf==2.3.0
-opencv-python==4.9.0.80
 packaging==23.2
 pillow==10.2.0
 protobuf==4.25.3
 pyparsing==3.1.1
 python-dateutil==2.8.2
 pyyaml==6.0.1
-scikit-image==0.22.0
 scipy==1.12.0
-shapely==2.0.3
 six==1.16.0
 tensorboardx==2.6.2.2
-tifffile==2024.2.12
 zipp==3.17.0

requirements/tune.in

@@ -2,7 +2,6 @@
 
 lmdb
 Pillow
-imgaug
 hydra-core >=1.2.0
 ray[tune] >=2.0.0
 ax-platform

requirements/tune.txt

@@ -20,8 +20,6 @@ fsspec==2024.2.0
 gpytorch==1.11
 hydra-core==1.3.2
 idna==3.6
-imageio==2.34.0
-imgaug==0.4.0
 importlib-resources==6.1.1
 ipython==8.18.1
 ipywidgets==8.1.2
@@ -33,7 +31,6 @@ jsonschema==4.21.1
 jsonschema-specifications==2023.12.1
 jupyterlab-widgets==3.0.10
 kiwisolver==1.4.5
-lazy-loader==0.3
 linear-operator==0.5.1
 lmdb==1.4.1
 markupsafe==2.1.5
@@ -46,7 +43,6 @@ mypy-extensions==1.0.0
 networkx==3.2.1
 numpy==1.26.4
 omegaconf==2.3.0
-opencv-python==4.9.0.80
 opt-einsum==3.3.0
 packaging==23.2
 pandas==2.2.0
@@ -71,17 +67,14 @@ ray==2.9.2
 referencing==0.33.0
 requests==2.31.0
 rpds-py==0.18.0
-scikit-image==0.22.0
 scikit-learn==1.4.1.post1
 scipy==1.12.0
-shapely==2.0.3
 six==1.16.0
 stack-data==0.6.3
 sympy==1.12
 tenacity==8.2.3
 tensorboardx==2.6.2.2
 threadpoolctl==3.3.0
-tifffile==2024.2.12
 tqdm==4.66.2
 traitlets==5.14.1
 typeguard==2.13.3

strhub/data/augment.py

@@ -12,10 +12,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+#
+# NOTE: imgaug has been replaced with NumPy / SciPy equivalents so that this
+# module is compatible with NumPy 2.x.  The public API (rand_augment_transform)
+# is unchanged.
 
 from functools import partial
 
-import imgaug.augmenters as iaa
 import numpy as np
 from PIL import Image, ImageFilter
 
@@ -42,6 +45,10 @@ def _get_param(level, img, max_dim_factor, min_level=1):
     return round(min(level, max_level))
 
 
+# ---------------------------------------------------------------------------
+# Blur ops
+# ---------------------------------------------------------------------------
+
 def gaussian_blur(img, radius, **__):
     radius = _get_param(radius, img, 0.02)
     key = 'gaussian_blur_' + str(radius)
@@ -50,26 +57,47 @@ def gaussian_blur(img, radius, **__):
 
 
 def motion_blur(img, k, **__):
+    """Horizontal motion blur implemented with scipy.ndimage (no imgaug)."""
     k = _get_param(k, img, 0.08, 3) | 1  # bin to odd values
-    key = 'motion_blur_' + str(k)
-    op = _get_op(key, lambda: iaa.MotionBlur(k))
-    return Image.fromarray(op(image=np.asarray(img)))
-
+    from scipy.ndimage import uniform_filter1d
+    arr = np.asarray(img, dtype=np.float32)
+    if arr.ndim == 2:
+        result = uniform_filter1d(arr, size=k, axis=1, mode='reflect')
+    else:
+        result = np.stack(
+            [uniform_filter1d(arr[..., c], size=k, axis=1, mode='reflect')
+             for c in range(arr.shape[2])],
+            axis=2,
+        )
+    return Image.fromarray(np.clip(result, 0, 255).astype(np.uint8))
+
+
+# ---------------------------------------------------------------------------
+# Noise ops
+# ---------------------------------------------------------------------------
 
 def gaussian_noise(img, scale, **_):
+    """Additive Gaussian noise implemented with NumPy (no imgaug)."""
     scale = _get_param(scale, img, 0.25) | 1  # bin to odd values
-    key = 'gaussian_noise_' + str(scale)
-    op = _get_op(key, lambda: iaa.AdditiveGaussianNoise(scale=scale))
-    return Image.fromarray(op(image=np.asarray(img)))
+    arr = np.asarray(img, dtype=np.float32)
+    rng = np.random.default_rng()
+    noise = rng.normal(0.0, float(scale), arr.shape).astype(np.float32)
+    return Image.fromarray(np.clip(arr + noise, 0, 255).astype(np.uint8))
 
 
 def poisson_noise(img, lam, **_):
+    """Additive Poisson noise implemented with NumPy (no imgaug)."""
     lam = _get_param(lam, img, 0.2) | 1  # bin to odd values
-    key = 'poisson_noise_' + str(lam)
-    op = _get_op(key, lambda: iaa.AdditivePoissonNoise(lam))
-    return Image.fromarray(op(image=np.asarray(img)))
+    arr = np.asarray(img, dtype=np.float32)
+    rng = np.random.default_rng()
+    noise = rng.poisson(float(lam), arr.shape).astype(np.float32)
+    return Image.fromarray(np.clip(arr + noise, 0, 255).astype(np.uint8))
 
 
+# ---------------------------------------------------------------------------
+# RandAugment integration
+# ---------------------------------------------------------------------------
+
 def _level_to_arg(level, _hparams, max):
     level = max * level / auto_augment._LEVEL_DENOM
     return (level,)
@@ -98,6 +126,10 @@ def _level_to_arg(level, _hparams, max):
 
 
 def rand_augment_transform(magnitude=5, num_layers=3):
+    """Return a RandAugment transform compatible with torchvision's Compose.
+
+    The signature is identical to the original so all call-sites are unchanged.
+    """
     # These are tuned for magnitude=5, which means that effective magnitudes are half of these values.
     hparams = {
         'rotate_deg': 30,