kiri-ocr 0.1.0__py3-none-any.whl

kiri_ocr/__init__.py

@@ -0,0 +1,14 @@
+from .core import OCR
+from .renderer import DocumentRenderer
+from .model import LightweightOCR, CharacterSet
+from .detector import TextDetector
+
+__version__ = '0.1.0'
+
+__all__ = [
+    'OCR',
+    'DocumentRenderer',
+    'LightweightOCR',
+    'CharacterSet',
+    'TextDetector',
+]

kiri_ocr/cli.py

@@ -0,0 +1,244 @@
+import numpy as np
+import json
+import argparse
+import sys
+from pathlib import Path
+import yaml
+
+from .core import OCR
+from .renderer import DocumentRenderer
+from .training import train_command
+from .generator import generate_command
+
+DEFAULT_TRAIN_CONFIG = {
+    "height": 32,
+    "batch_size": 32,
+    "epochs": 2,
+    "hidden_size": 256,
+    "device": "cuda",
+    "output_dir": "models",
+    "train_labels": "data/train/labels.txt",
+    "val_labels": "data/val/labels.txt",
+    "lr": 0.001,
+    "weight_decay": 0.0001
+}
+
+def init_config(args):
+    path = args.output
+    # Ensure .yaml extension if default or user didn't specify
+    if path == 'config.yaml' and not path.endswith('.yaml') and not path.endswith('.yml'):
+         pass # user provided custom name
+    
+    with open(path, 'w') as f:
+        yaml.dump(DEFAULT_TRAIN_CONFIG, f, default_flow_style=False)
+    print(f"✓ Created default config at {path}")
+
+def run_inference(args):
+    # Create output directory
+    output_dir = Path(args.output)
+    output_dir.mkdir(exist_ok=True)
+    
+    print("\n" + "="*70)
+    print("  📄 Kiri OCR System")
+    print("="*70)
+    
+    try:
+        # Initialize OCR
+        ocr = OCR(
+            model_path=args.model,
+            charset_path=args.charset,
+            language=args.language,
+            padding=args.padding,
+            device=args.device,
+            verbose=True
+        )
+        
+        # Process document & Extract text
+        full_text, results = ocr.extract_text(args.image, mode=args.mode, verbose=True)
+        
+        # Save text
+        text_output = output_dir / 'extracted_text.txt'
+        with open(text_output, 'w', encoding='utf-8') as f:
+            f.write(full_text)
+        print(f"\n✓ Text saved to {text_output}")
+        
+        # Save JSON
+        json_output = output_dir / 'ocr_results.json'
+        with open(json_output, 'w', encoding='utf-8') as f:
+            json.dump(results, f, indent=2, ensure_ascii=False)
+        print(f"✓ JSON saved to {json_output}")
+        
+        # Render results
+        if not args.no_render:
+            renderer = DocumentRenderer()
+            
+            # Boxes only
+            renderer.draw_boxes(
+                args.image, 
+                results, 
+                output_path=str(output_dir / 'boxes.png')
+            )
+            
+            # Boxes with text
+            renderer.draw_results(
+                args.image,
+                results,
+                output_path=str(output_dir / 'ocr_result.png')
+            )
+            
+            # HTML report
+            renderer.create_report(
+                args.image,
+                results,
+                output_path=str(output_dir / 'report.html')
+            )
+        
+        print("\n" + "="*70)
+        print("  ✅ Processing Complete!")
+        print("="*70)
+        print(f"  Regions detected: {len(results)}")
+        if results:
+            print(f"  Average confidence: {np.mean([r['confidence'] for r in results])*100:.2f}%")
+        print(f"  Output directory: {output_dir}")
+        print("="*70 + "\n")
+        
+    except Exception as e:
+        print(f"\n❌ Error: {e}")
+        # Suggest model path if it seems to be missing
+        if "No such file" in str(e) and "model" in str(e):
+             print("\nTip: Make sure you have trained a model or specified the correct path with --model")
+             print("     Run: kiri-ocr train ...  to train a model first.")
+
+def merge_config(args, defaults):
+    """Merge defaults < config file < args"""
+    # Start with defaults
+    config = defaults.copy()
+    
+    # Update with config file if provided
+    if hasattr(args, 'config') and args.config:
+        try:
+            with open(args.config, 'r') as f:
+                if args.config.endswith('.json'):
+                    file_config = json.load(f)
+                else:
+                    file_config = yaml.safe_load(f)
+                
+                if file_config:
+                    config.update(file_config)
+            print(f"Loaded config from {args.config}")
+        except Exception as e:
+            print(f"Error loading config file: {e}")
+            sys.exit(1)
+            
+    # Update with explicit args (non-None)
+    for key, value in vars(args).items():
+        if value is not None and key in config:
+            config[key] = value
+            
+    # Update args object
+    for key, value in config.items():
+        setattr(args, key, value)
+        
+    return args
+
+def main():
+    parser = argparse.ArgumentParser(description='Kiri OCR - Command Line Tool')
+    subparsers = parser.add_subparsers(dest='command', help='Command to run')
+    
+    # === PREDICT / RUN ===
+    predict_parser = subparsers.add_parser('predict', help='Run OCR on an image')
+    predict_parser.add_argument('image', help='Path to document image')
+    predict_parser.add_argument('--mode', choices=['lines', 'words'], default='lines',
+                       help='Detection mode (default: lines)')
+    predict_parser.add_argument('--model', default='models/model.kiri',
+                       help='Path to model file')
+    predict_parser.add_argument('--charset', default='models/charset_lite.txt',
+                       help='Path to charset file')
+    predict_parser.add_argument('--language', choices=['english', 'khmer', 'mixed'], 
+                       default='mixed', help='Language mode')
+    predict_parser.add_argument('--padding', type=int, default=10,
+                       help='Padding around detected boxes in pixels (default: 10)')
+    predict_parser.add_argument('--output', '-o', default='output',
+                       help='Output directory')
+    predict_parser.add_argument('--no-render', action='store_true',
+                       help='Skip rendering (text only)')
+    predict_parser.add_argument('--device', choices=['cpu', 'cuda'], default='cpu',
+                       help='Device to use')
+    
+    # === TRAIN ===
+    train_parser = subparsers.add_parser('train', help='Train the OCR model')
+    train_parser.add_argument('--config', help='Path to config file (YAML/JSON)')
+    train_parser.add_argument('--train-labels', help='Path to training labels file')
+    train_parser.add_argument('--val-labels', help='Path to validation labels file')
+    train_parser.add_argument('--hf-dataset', help='HuggingFace dataset ID (e.g. mrrtmob/km_en_image_line)')
+    train_parser.add_argument('--hf-subset', help='HuggingFace dataset subset/config name')
+    train_parser.add_argument('--hf-train-split', default='train', help='Train split name (default: train)')
+    train_parser.add_argument('--hf-val-split', help='Validation split name')
+    train_parser.add_argument('--hf-image-col', default='image', help='Image column name')
+    train_parser.add_argument('--hf-text-col', default='text', help='Text column name')
+    train_parser.add_argument('--hf-val-percent', type=float, default=0.1, help='Val % if no split found (default: 0.1)')
+    train_parser.add_argument('--output-dir', help='Directory to save model')
+    train_parser.add_argument('--epochs', type=int, help='Number of epochs')
+    train_parser.add_argument('--batch-size', type=int, help='Batch size')
+    train_parser.add_argument('--height', type=int, help='Image height')
+    train_parser.add_argument('--hidden-size', type=int, help='Hidden size')
+    train_parser.add_argument('--device', help='Device (cuda/cpu)')
+    train_parser.add_argument('--from-model', help='Path to pretrained model for fine-tuning')
+    train_parser.add_argument('--lr', type=float, help='Learning rate')
+    train_parser.add_argument('--weight-decay', type=float, help='Weight decay')
+    
+    # === GENERATE ===
+    gen_parser = subparsers.add_parser('generate', help='Generate synthetic data')
+    gen_parser.add_argument('--train-file', '-t', required=True,
+                       help='Training text file (one line per sample)')
+    gen_parser.add_argument('--val-file', '-v', default=None,
+                       help='Validation text file (optional)')
+    gen_parser.add_argument('--output', '-o', default='data',
+                       help='Output directory')
+    gen_parser.add_argument('--language', '-l', choices=['english', 'khmer', 'mixed'],
+                       default='mixed', help='Language mode')
+    gen_parser.add_argument('--augment', '-a', type=int, default=1,
+                       help='Augmentation factor for training')
+    gen_parser.add_argument('--val-augment', type=int, default=1,
+                       help='Augmentation factor for validation')
+    gen_parser.add_argument('--height', type=int, default=32,
+                       help='Image height')
+    gen_parser.add_argument('--width', type=int, default=512,
+                       help='Image width')
+    gen_parser.add_argument('--fonts-dir', default='fonts',
+                       help='Directory containing .ttf font files')
+    gen_parser.add_argument('--font-mode', choices=['random', 'all'], default='random',
+                       help='Font selection mode: random (default) or all (iterate all fonts per line)')
+    gen_parser.add_argument('--random-augment', action='store_true',
+                       help='Apply random augmentations (noise, rotation) even if augmentation factor is 1')
+
+    # === INIT CONFIG ===
+    init_parser = subparsers.add_parser('init-config', help='Create default config file')
+    init_parser.add_argument('--output', '-o', default='config.yaml', help='Output file')
+
+    # Backward compatibility logic
+    if len(sys.argv) > 1 and sys.argv[1] not in ['predict', 'train', 'generate', 'init-config', '-h', '--help']:
+        sys.argv.insert(1, 'predict')
+        
+    args = parser.parse_args()
+    
+    if args.command == 'predict':
+        run_inference(args)
+    elif args.command == 'train':
+        # Merge config for training
+        args = merge_config(args, DEFAULT_TRAIN_CONFIG)
+        # Check required fields
+        if not args.train_labels and not args.hf_dataset:
+            print("❌ Error: --train-labels or --hf-dataset is required")
+            sys.exit(1)
+            
+        train_command(args)
+    elif args.command == 'generate':
+        generate_command(args)
+    elif args.command == 'init-config':
+        init_config(args)
+    else:
+        parser.print_help()
+
+if __name__ == '__main__':
+    main()

kiri_ocr/core.py

@@ -0,0 +1,306 @@
+import torch
+import cv2
+import numpy as np
+import sys
+from PIL import Image
+from pathlib import Path
+import json
+
+from .model import LightweightOCR, CharacterSet
+from .detector import TextDetector
+
+class OCR:
+    """Complete Document OCR System with Padding"""
+    
+    def __init__(self, model_path='models/model.kiri',
+                 charset_path='models/charset_lite.txt',
+                 language='mixed',
+                 padding=10,
+                 device='cpu',
+                 verbose=False):
+        """
+        Args:
+            model_path: Path to trained model (.kiri or .pth)
+            charset_path: Path to character set (used if model doesn't contain charset)
+            language: 'english', 'khmer', or 'mixed'
+            padding: Pixels to pad around detected boxes (default: 10)
+            device: 'cpu' or 'cuda'
+            verbose: Whether to print loading/processing info
+        """
+        self.device = device
+        self.verbose = verbose
+        self.language = language
+        self.padding = padding
+        
+        # Resolve model path
+        if not Path(model_path).exists():
+             # Try looking in package directory
+             pkg_dir = Path(__file__).parent
+             if (pkg_dir / model_path).exists():
+                 model_path = str(pkg_dir / model_path)
+             # Try looking in sibling 'models' package (if installed via setup.py with models package)
+             elif (pkg_dir.parent / 'models' / Path(model_path).name).exists():
+                 model_path = str(pkg_dir.parent / 'models' / Path(model_path).name)
+             # Fallback to .pth if .kiri not found
+             elif model_path.endswith('.kiri') and (Path(model_path).parent / 'model.pth').exists():
+                 model_path = str(Path(model_path).parent / 'model.pth')
+        
+        if self.verbose:
+            print(f"📦 Loading OCR model from {model_path}...")
+        
+        try:
+            checkpoint = torch.load(model_path, map_location=device, weights_only=False)
+            
+            # Load charset
+            if 'charset' in checkpoint:
+                if self.verbose:
+                    print(f"  ✓ Found embedded charset ({len(checkpoint['charset'])} chars)")
+                self.charset = CharacterSet.from_checkpoint(checkpoint)
+            else:
+                # Fallback to charset file
+                if not Path(charset_path).exists():
+                     # Try looking in package directory if not found locally
+                     pkg_dir = Path(__file__).parent
+                     if (pkg_dir / charset_path).exists():
+                         charset_path = str(pkg_dir / charset_path)
+                
+                if self.verbose:
+                    print(f"  ℹ️ Loading charset from file: {charset_path}")
+                self.charset = CharacterSet.load(charset_path)
+            
+            # Initialize model
+            self.model = LightweightOCR(num_chars=len(self.charset)).eval()
+            
+            # Load weights
+            if 'model_state_dict' in checkpoint:
+                self.model.load_state_dict(checkpoint['model_state_dict'])
+            else:
+                # Assume checkpoint IS state_dict (legacy/raw save)
+                self.model.load_state_dict(checkpoint)
+                
+        except RuntimeError as e:
+            if "size mismatch" in str(e):
+                print(f"\n❌ Error loading model: {e}")
+                print("\n⚠️  CRITICAL: The model weights do not match the character set.")
+                print(f"    - Charset size: {len(self.charset)}")
+                print(f"    - Model file:   {model_path}")
+                sys.exit(1)
+            else:
+                raise e
+        except Exception as e:
+            print(f"❌ Error loading model: {e}")
+            raise e
+        
+        self.model = self.model.to(device)
+        
+        if self.verbose:
+            print(f"✓ Model loaded ({len(self.charset)} characters)")
+            print(f"✓ Box padding: {padding}px")
+        
+        # Detector
+        self.detector = TextDetector()
+    
+    def _preprocess_region(self, img, box, extra_padding=5):
+        """
+        Crop and preprocess a region with extra padding
+        
+        Args:
+            img: Source image (numpy array)
+            box: Bounding box (x, y, w, h) - already has padding from detector
+            extra_padding: Additional padding when cropping (default: 5px)
+        """
+        img_h, img_w = img.shape[:2]
+        x, y, w, h = box
+        
+        # Add extra padding (with boundary checks)
+        x_extra = max(0, x - extra_padding)
+        y_extra = max(0, y - extra_padding)
+        w_extra = min(img_w - x_extra, w + 2 * extra_padding)
+        h_extra = min(img_h - y_extra, h + 2 * extra_padding)
+        
+        # Crop
+        roi = img[y_extra:y_extra+h_extra, x_extra:x_extra+w_extra]
+        
+        if roi.size == 0:
+            return None
+        
+        # Invert if dark background (Model expects Light BG / Dark Text)
+        if np.mean(roi) < 127:
+             roi = 255 - roi
+
+        # Convert to PIL
+        roi_pil = Image.fromarray(roi).convert('L')
+        
+        # Resize maintaining aspect ratio
+        orig_w, orig_h = roi_pil.size
+        new_h = 32
+        new_w = int((orig_w / orig_h) * new_h)
+        
+        # Ensure minimum width
+        if new_w < 32:
+            new_w = 32
+        
+        roi_pil = roi_pil.resize((new_w, new_h), Image.LANCZOS)
+        
+        # Normalize
+        roi_array = np.array(roi_pil) / 255.0
+        roi_tensor = torch.tensor(roi_array, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
+        
+        return roi_tensor
+    
+    def recognize_single_line_image(self, image_path):
+        """Recognize text from a single line image without detection"""
+        img = cv2.imread(str(image_path))
+        if img is None:
+            raise ValueError(f"Could not load image: {image_path}")
+            
+        # Preprocess
+        # We need to resize height to 32 and keep aspect ratio
+        if len(img.shape) == 3:
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+            
+        # Invert if dark background (Model expects Light BG / Dark Text)
+        if np.mean(img) < 127:
+             img = 255 - img
+
+        img_pil = Image.fromarray(img)
+        w, h = img_pil.size
+        new_h = 32
+        new_w = int((w / h) * new_h)
+        if new_w < 32: new_w = 32
+        
+        img_pil = img_pil.resize((new_w, new_h), Image.LANCZOS)
+        img_array = np.array(img_pil) / 255.0
+        img_tensor = torch.tensor(img_array, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
+        
+        text, confidence = self.recognize_region(img_tensor)
+        return text, confidence
+
+    def recognize_region(self, image_tensor):
+        """Recognize text in a single region"""
+        image_tensor = image_tensor.to(self.device)
+        
+        with torch.no_grad():
+            logits = self.model(image_tensor)
+            
+            # CTC decoding (LightweightOCR)
+            probs = torch.softmax(logits, dim=-1)
+            preds = torch.argmax(probs, dim=-1)
+            preds = preds.squeeze().tolist()
+            
+            # Handle single timestep
+            if not isinstance(preds, list):
+                preds = [preds]
+            
+            # CTC decode
+            char_confidences = []
+            decoded_indices = []
+            previous_idx = -1
+            
+            for i, idx in enumerate(preds):
+                if idx != previous_idx:
+                    if idx > 2:  # Skip BLANK, PAD, SOS
+                        decoded_indices.append(idx)
+                        char_confidences.append(probs[i, 0, idx].item())
+                previous_idx = idx
+            
+            confidence = np.mean(char_confidences) if char_confidences else 0.0
+            text = self.charset.decode(decoded_indices)
+        
+        return text, confidence
+    
+    def process_document(self, image_path, mode='lines', verbose=False):
+        """
+        Process entire document
+        
+        Args:
+            image_path: Path to document image
+            mode: 'lines' or 'words' for detection granularity
+            verbose: Whether to print progress
+        
+        Returns:
+            List of dicts with 'box', 'text', 'confidence'
+        """
+        if verbose:
+            print(f"\n📄 Processing document: {image_path}")
+            print(f"🔲 Box padding: {self.padding}px")
+        
+        # Detect text regions
+        if mode == 'lines':
+            boxes = self.detector.detect_lines(image_path)
+        else:
+            boxes = self.detector.detect_words(image_path)
+        
+        if verbose:
+            print(f"🔍 Detected {len(boxes)} regions")
+        
+        # Load image
+        img = cv2.imread(str(image_path))
+        
+        # Recognize each region
+        results = []
+        for i, box in enumerate(boxes, 1):
+            try:
+                # Preprocess with extra padding
+                region_tensor = self._preprocess_region(img, box, extra_padding=5)
+                if region_tensor is None:
+                    continue
+                
+                # Recognize
+                text, confidence = self.recognize_region(region_tensor)
+                
+                # Convert types for JSON serialization
+                safe_box = [int(v) for v in box]
+                safe_confidence = float(confidence)
+                
+                results.append({
+                    'box': safe_box,
+                    'text': text,
+                    'confidence': safe_confidence,
+                    'line_number': i
+                })
+                
+                if verbose:
+                    print(f"  {i:2d}. {text:50s} ({confidence*100:.1f}%)")
+                
+            except Exception as e:
+                if verbose:
+                    print(f"  {i:2d}. [Error: {e}]")
+                continue
+        
+        return results
+    
+    def extract_text(self, image_path, mode='lines', verbose=False):
+        """Extract all text from document as string"""
+        results = self.process_document(image_path, mode, verbose=verbose)
+        
+        if not results:
+            return "", results
+            
+        # Reconstruct text layout
+        # Sort by Y then X
+        results.sort(key=lambda r: (r['box'][1], r['box'][0]))
+        
+        full_text = ""
+        for i, res in enumerate(results):
+            text = res['text']
+            box = res['box']
+            _, y, _, h = box
+            
+            if i > 0:
+                prev_box = results[i-1]['box']
+                prev_y, prev_h = prev_box[1], prev_box[3]
+                
+                # Check if on the same line (vertical center is close)
+                center_y = y + h/2
+                prev_center_y = prev_y + prev_h/2
+                
+                # If vertical centers are close (within half height), assume same line
+                if abs(center_y - prev_center_y) < max(h, prev_h) / 2:
+                    full_text += " " + text
+                else:
+                    full_text += "\n" + text
+            else:
+                full_text += text
+        
+        return full_text, results