Accept torch.Tensor inputs, remove non-release scripts

- Add torch.Tensor support (CPU/CUDA, CHW/HWC, uint8/float)
- Add pathlib.Path support for file paths

nemotron-ocr/src/nemotron_ocr/inference/pipeline_v2.py

@@ -12,6 +12,7 @@ Extends :class:`NemotronOCR` with:
 """
 
 import logging
+import os
 import time
 
 import numpy as np
@@ -237,10 +238,30 @@ class NemotronOCRV2(NemotronOCR):
 
     @staticmethod
     def _load_image_to_tensor_uint8(image):
-        """Load image and return a uint8 CHW tensor (no dtype conversion).
+        """Load image and return a uint8 CHW tensor on CPU.
 
-        Defers the expensive uint8→float16 conversion to the GPU.
+        Accepts file paths, numpy arrays, torch tensors, base64 bytes,
+        or BytesIO objects.
         """
+        if isinstance(image, torch.Tensor):
+            t = image.detach().cpu()
+            if t.ndim == 2:
+                t = t.unsqueeze(0).expand(3, -1, -1)
+            if t.ndim == 3 and t.shape[0] in (1, 3, 4):
+                if t.shape[0] == 4:
+                    t = t[:3]
+                if t.dtype != torch.uint8:
+                    t = (t * 255).clamp(0, 255).to(torch.uint8) if t.is_floating_point() else t.to(torch.uint8)
+                return t
+            if t.ndim == 3 and t.shape[2] in (1, 3, 4):
+                t = t.permute(2, 0, 1)
+                if t.shape[0] == 4:
+                    t = t[:3]
+                if t.dtype != torch.uint8:
+                    t = (t * 255).clamp(0, 255).to(torch.uint8) if t.is_floating_point() else t.to(torch.uint8)
+                return t
+            raise ValueError(f"Unsupported tensor shape: {image.shape}")
+
         if isinstance(image, np.ndarray):
             if image.ndim == 2:
                 image = np.stack([image] * 3, axis=-1)
@@ -251,7 +272,7 @@ class NemotronOCRV2(NemotronOCR):
             return torch.from_numpy(image).permute(2, 0, 1)
 
         from torchvision.io import read_image, decode_image
-        if isinstance(image, str):
+        if isinstance(image, (str, os.PathLike)):
             return read_image(str(image), mode="RGB")
         if isinstance(image, bytes):
             import base64

run_ocr_inference.py

@@ -1,251 +0,0 @@
-#!/usr/bin/env python3
-"""
-Run OCR inference over a directory of images and write:
-  - images/          : JPG copies of processed input images
-  - visualizations/  : Annotated images with bounding boxes + text
-  - results.json     : OCR outputs for all images
-"""
-
-import argparse
-import json
-import logging
-import os
-import random
-from pathlib import Path
-
-from PIL import Image, ImageDraw, ImageFont
-from tqdm import tqdm
-
-from nemotron_ocr.inference.pipeline import NemotronOCR
-from nemotron_ocr.inference.pipeline_v2 import NemotronOCRV2
-
-SCRIPT_DIR = Path(__file__).parent.resolve()
-FONT_PATH = SCRIPT_DIR / "NotoSansCJKsc-Regular.otf"
-
-logging.getLogger("nemotron_ocr.inference.pipeline_v2").setLevel(logging.WARNING)
-logging.getLogger("nemotron_ocr.inference.models.relational").setLevel(logging.WARNING)
-
-
-def get_font(size: int = 14):
-    try:
-        return ImageFont.truetype(str(FONT_PATH), size)
-    except Exception:
-        return ImageFont.load_default()
-
-
-def find_images(images_dir, extensions=(".png", ".jpg", ".jpeg", ".tiff", ".bmp")):
-    images_dir = Path(images_dir)
-    image_files = []
-    for ext in extensions:
-        image_files.extend(images_dir.glob(f"**/*{ext}"))
-        image_files.extend(images_dir.glob(f"**/*{ext.upper()}"))
-    return sorted(set(image_files))
-
-
-def save_visualization(image_path: str, predictions: list, output_path: str):
-    pil_image = Image.open(image_path).convert("RGB")
-    draw = ImageDraw.Draw(pil_image)
-    font = get_font(14)
-
-    img_width, img_height = pil_image.size
-    colors = [
-        (255, 0, 0),
-        (0, 255, 0),
-        (0, 0, 255),
-        (255, 165, 0),
-        (128, 0, 128),
-        (0, 128, 128),
-        (255, 192, 203),
-        (165, 42, 42),
-    ]
-
-    for i, pred in enumerate(predictions):
-        if isinstance(pred.get("left"), str) and pred["left"] == "nan":
-            continue
-
-        color = colors[i % len(colors)]
-        left = int(pred["left"] * img_width)
-        right = int(pred["right"] * img_width)
-        upper = int(pred.get("upper", pred.get("bottom", 0)) * img_height)
-        lower = int(pred.get("lower", pred.get("top", 0)) * img_height)
-        top, bottom = min(upper, lower), max(upper, lower)
-
-        draw.rectangle([left, top, right, bottom], outline=color, width=2)
-
-        display_text = pred["text"][:50] + "..." if len(pred["text"]) > 50 else pred["text"]
-        text_y = max(0, top - 20)
-        try:
-            text_bbox = draw.textbbox((left, text_y), display_text, font=font)
-            draw.rectangle(
-                [text_bbox[0] - 2, text_bbox[1] - 2, text_bbox[2] + 2, text_bbox[3] + 2],
-                fill=(255, 255, 255, 200),
-                outline=color,
-            )
-            draw.text((left, text_y), display_text, fill=color, font=font)
-        except Exception:
-            draw.text((left, text_y), display_text, fill=color, font=font)
-
-    os.makedirs(os.path.dirname(output_path) if os.path.dirname(output_path) else ".", exist_ok=True)
-    pil_image.save(output_path, "JPEG", quality=90)
-
-
-def save_image_as_jpg(image_path: str, output_path: str):
-    pil_image = Image.open(image_path).convert("RGB")
-    pil_image.save(output_path, "JPEG", quality=90)
-
-
-def _store_result(image_path: Path, predictions: list, images_output_dir: Path, viz_output_dir: Path):
-    result_key = f"{image_path.stem}.jpg"
-    row = {
-        "original_path": str(image_path),
-        "predictions": predictions,
-        "num_detections": len(predictions),
-        "status": "success",
-    }
-
-    dest_image_path = images_output_dir / f"{image_path.stem}.jpg"
-    if not dest_image_path.exists():
-        save_image_as_jpg(str(image_path), str(dest_image_path))
-    if predictions:
-        vis_path = viz_output_dir / f"{image_path.stem}_viz.jpg"
-        save_visualization(str(image_path), predictions, str(vis_path))
-    return result_key, row
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Run OCR inference on image directories")
-    parser.add_argument("--model_dir", type=str, required=True, help="Path to model checkpoint directory")
-    parser.add_argument("--images_dir", type=str, required=True, help="Path to directory with input images")
-    parser.add_argument("--output_dir", type=str, required=True, help="Output directory")
-    parser.add_argument("--num_samples", type=int, default=0, help="Number of images to process (0 for all)")
-    parser.add_argument(
-        "--merge_level",
-        type=str,
-        default="sentence",
-        choices=["word", "sentence", "paragraph"],
-        help="Text merging granularity",
-    )
-    parser.add_argument("--seed", type=int, default=42, help="Random seed for sampling")
-    parser.add_argument("--include_invalid", action="store_true", help="Include regions marked invalid")
-    parser.add_argument("--infer_length", type=int, default=1024, help="Detector resolution")
-    parser.add_argument("--pipeline", type=str, default="v2", choices=("v1", "v2"), help="OCR pipeline version")
-    parser.add_argument("--batch_size", type=int, default=8, help="Batch size for v2 pipeline")
-    parser.add_argument("--verbose", action="store_true", help="Enable verbose OCR logs")
-    args = parser.parse_args()
-
-    if args.verbose:
-        logging.getLogger("nemotron_ocr.inference.pipeline_v2").setLevel(logging.INFO)
-
-    output_dir = Path(args.output_dir)
-    images_output_dir = output_dir / "images"
-    viz_output_dir = output_dir / "visualizations"
-    output_dir.mkdir(parents=True, exist_ok=True)
-    images_output_dir.mkdir(exist_ok=True)
-    viz_output_dir.mkdir(exist_ok=True)
-
-    image_files = [Path(p) for p in find_images(args.images_dir)]
-    if not image_files:
-        print("No images found.")
-        return
-
-    if 0 < args.num_samples < len(image_files):
-        random.seed(args.seed)
-        image_files = random.sample(image_files, args.num_samples)
-
-    if args.pipeline == "v2":
-        ocr = NemotronOCRV2(
-            model_dir=args.model_dir,
-            infer_length=args.infer_length,
-            detector_max_batch_size=max(1, args.batch_size),
-        )
-    else:
-        ocr = NemotronOCR(args.model_dir, infer_length=args.infer_length)
-
-    results = {}
-
-    if args.pipeline == "v2":
-        bs = max(1, args.batch_size)
-        for start in tqdm(range(0, len(image_files), bs), desc="Batches"):
-            batch_paths = image_files[start:start + bs]
-            batch_strs = [str(p) for p in batch_paths]
-            batch_preds = None
-            try:
-                batch_preds = ocr(
-                    batch_strs,
-                    merge_level=args.merge_level,
-                    include_invalid=args.include_invalid,
-                )
-            except Exception:
-                batch_preds = None
-
-            if batch_preds is None:
-                for image_path in batch_paths:
-                    try:
-                        predictions = ocr(
-                            str(image_path),
-                            merge_level=args.merge_level,
-                            include_invalid=args.include_invalid,
-                        )
-                        key, row = _store_result(image_path, predictions, images_output_dir, viz_output_dir)
-                        results[key] = row
-                    except Exception as e:
-                        import traceback
-
-                        results[f"{image_path.stem}.jpg"] = {
-                            "original_path": str(image_path),
-                            "predictions": [],
-                            "num_detections": 0,
-                            "status": "error",
-                            "error": str(e),
-                            "traceback": traceback.format_exc(),
-                        }
-                continue
-
-            for image_path, predictions in zip(batch_paths, batch_preds):
-                try:
-                    key, row = _store_result(image_path, predictions, images_output_dir, viz_output_dir)
-                    results[key] = row
-                except Exception as e:
-                    import traceback
-
-                    results[f"{image_path.stem}.jpg"] = {
-                        "original_path": str(image_path),
-                        "predictions": [],
-                        "num_detections": 0,
-                        "status": "error",
-                        "error": str(e),
-                        "traceback": traceback.format_exc(),
-                    }
-    else:
-        for image_path in tqdm(image_files, desc="Processing"):
-            try:
-                predictions = ocr(
-                    str(image_path),
-                    merge_level=args.merge_level,
-                    include_invalid=args.include_invalid,
-                )
-                key, row = _store_result(image_path, predictions, images_output_dir, viz_output_dir)
-                results[key] = row
-            except Exception as e:
-                import traceback
-
-                results[f"{image_path.stem}.jpg"] = {
-                    "original_path": str(image_path),
-                    "predictions": [],
-                    "num_detections": 0,
-                    "status": "error",
-                    "error": str(e),
-                    "traceback": traceback.format_exc(),
-                }
-
-    results_path = output_dir / "results.json"
-    with open(results_path, "w", encoding="utf-8") as f:
-        json.dump(results, f, ensure_ascii=False, indent=2)
-
-    success_count = sum(1 for r in results.values() if r["status"] == "success")
-    total_detections = sum(r["num_detections"] for r in results.values())
-    print(f"Done: {success_count}/{len(results)} images, {total_detections} detections -> {output_dir}")
-
-
-if __name__ == "__main__":
-    main()

submit_inference_viz_gpu.sh

@@ -1,87 +0,0 @@
-#!/usr/bin/env bash
-# Submit OCR + visualization (NemotronOCRV2) on a GPU node from scratch clone.
-#
-# Usage:
-#   ./submit_inference_viz_gpu.sh /path/to/images [OUTPUT_DIR]
-#
-# Environment (optional):
-#   MODEL        — checkpoint directory (default: $BASE/v2_multilingual)
-#   BATCH_SIZE   — v2 detector batch size (default: 16)
-#   MERGE_LEVEL  — word|sentence|paragraph (default: sentence)
-#   SLURM_TIME   — wall time in minutes (default: 120)
-#   JOB_NAME     — sbatch job name (default: infer_viz_scratch)
-
-set -euo pipefail
-
-BASE="/lustre/fsw/portfolios/datascience/users/rchesler/scratch/nemotron-ocr-v2"
-VENV="$BASE/nemotron-ocr/.venv"
-ACCOUNT="datascience_nemo_retriever"
-PARTITIONS="batch_block1,batch_block3,batch_block4"
-
-IMAGES="${1:?usage: $0 IMAGES_DIR [OUTPUT_DIR]}"
-TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
-OUT="${2:-$BASE/inference_output/viz_v2_${TIMESTAMP}}"
-MODEL="${MODEL:-$BASE/v2_multilingual}"
-BATCH_SIZE="${BATCH_SIZE:-16}"
-MERGE_LEVEL="${MERGE_LEVEL:-sentence}"
-SLURM_TIME="${SLURM_TIME:-120}"
-JOB_NAME="${JOB_NAME:-infer_viz_scratch}"
-
-LOGS="$OUT/logs"
-mkdir -p "$LOGS"
-
-PREAMBLE='source /etc/profile.d/modules.sh 2>/dev/null || true
-if command -v module >/dev/null 2>&1; then
-  module load cuda12.2/toolkit/12.2.2 || true
-fi
-export PATH="$HOME/.local/bin:$PATH"
-source '"$VENV"'/bin/activate'
-
-JOB="$OUT/.job_inference_${TIMESTAMP}.sh"
-cat > "$JOB" << ENDSCRIPT
-#!/bin/bash
-#SBATCH --job-name=${JOB_NAME}
-#SBATCH --account=${ACCOUNT}
-#SBATCH --partition=${PARTITIONS}
-#SBATCH --nodes=1
-#SBATCH --gpus-per-node=1
-#SBATCH --time=${SLURM_TIME}
-#SBATCH --output=${LOGS}/infer_%j.out
-#SBATCH --error=${LOGS}/infer_%j.err
-
-set -euo pipefail
-${PREAMBLE}
-cd ${BASE}
-
-python run_ocr_inference.py \\
-  --model_dir "${MODEL}" \\
-  --images_dir "${IMAGES}" \\
-  --output_dir "${OUT}" \\
-  --pipeline v2 \\
-  --batch_size ${BATCH_SIZE} \\
-  --num_samples 0 \\
-  --merge_level ${MERGE_LEVEL}
-
-echo "Done: \$(date)"
-ENDSCRIPT
-chmod +x "$JOB"
-
-echo "============================================================"
-echo "  Inference + viz (v2 pipeline) — GPU job (scratch clone)"
-echo "============================================================"
-echo "  images:   $IMAGES"
-echo "  model:    $MODEL"
-echo "  merge:    $MERGE_LEVEL"
-echo "  output:   $OUT"
-echo "  logs:     $LOGS"
-echo "  job file: $JOB"
-echo "============================================================"
-
-JID=$(sbatch "$JOB" 2>&1 | sed -n 's/.* \([0-9][0-9]*\)$/\1/p' || true)
-if [[ -n "${JID:-}" ]]; then
-  echo "  submitted job ID: $JID"
-  echo "  tail:  tail -f $LOGS/infer_${JID}.out"
-else
-  echo "  sbatch failed or did not return a job id" >&2
-  exit 1
-fi

submit_install_gpu.sh

@@ -1,79 +0,0 @@
-#!/usr/bin/env bash
-# Submit a GPU-node install/build job for scratch nemotron-ocr-v2.
-#
-# Usage:
-#   ./submit_install_gpu.sh
-#
-# Environment (optional):
-#   SLURM_TIME   — wall time in minutes (default: 45)
-#   JOB_NAME     — sbatch job name (default: install_ocr_v2)
-
-set -euo pipefail
-
-BASE="/lustre/fsw/portfolios/datascience/users/rchesler/scratch/nemotron-ocr-v2"
-PKG="$BASE/nemotron-ocr"
-ACCOUNT="datascience_nemo_retriever"
-PARTITIONS="batch_block1,batch_block3,batch_block4"
-SLURM_TIME="${SLURM_TIME:-45}"
-JOB_NAME="${JOB_NAME:-install_ocr_v2}"
-
-TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
-LOGS="$BASE/inference_output/install_logs"
-mkdir -p "$LOGS"
-
-JOB="$BASE/.job_install_${TIMESTAMP}.sh"
-cat > "$JOB" << ENDSCRIPT
-#!/bin/bash
-#SBATCH --job-name=${JOB_NAME}
-#SBATCH --account=${ACCOUNT}
-#SBATCH --partition=${PARTITIONS}
-#SBATCH --nodes=1
-#SBATCH --gpus-per-node=1
-#SBATCH --time=${SLURM_TIME}
-#SBATCH --output=${LOGS}/install_%j.out
-#SBATCH --error=${LOGS}/install_%j.err
-
-set -euo pipefail
-source /etc/profile.d/modules.sh 2>/dev/null || true
-if command -v module >/dev/null 2>&1; then
-  module load cuda12.2/toolkit/12.2.2 || true
-fi
-export PATH="\$HOME/.local/bin:\$PATH"
-
-cd "${PKG}"
-uv venv
-source .venv/bin/activate
-
-uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128
-rm -f src/nemotron_ocr_cpp/_nemotron_ocr_cpp*.so
-BUILD_CPP_FORCE=1 uv pip install -e . -v
-
-python -c "
-import torch
-print(f'torch={torch.__version__} cuda_available={torch.cuda.is_available()}')
-if torch.cuda.is_available():
-    print(f'gpu={torch.cuda.get_device_name(0)}')
-import nemotron_ocr, nemotron_ocr_cpp
-from nemotron_ocr.inference.pipeline_v2 import NemotronOCRV2
-print('nemotron_ocr import OK; nemotron_ocr_cpp OK; NemotronOCRV2 OK')
-"
-ENDSCRIPT
-
-chmod +x "$JOB"
-
-echo "============================================================"
-echo "  GPU install job (scratch clone)"
-echo "============================================================"
-echo "  package:  $PKG"
-echo "  logs:     $LOGS"
-echo "  job file: $JOB"
-echo "============================================================"
-
-JID=$(sbatch "$JOB" 2>&1 | sed -n 's/.* \([0-9][0-9]*\)$/\1/p' || true)
-if [[ -n "${JID:-}" ]]; then
-  echo "  submitted job ID: $JID"
-  echo "  tail:  tail -f $LOGS/install_${JID}.out"
-else
-  echo "  sbatch failed or did not return a job id" >&2
-  exit 1
-fi