kiri-ocr 0.1.0__py3-none-any.whl

kiri_ocr/model.py

@@ -0,0 +1,159 @@
+import torch
+import torch.nn as nn
+from pathlib import Path
+
+# ========== CHARACTER SET ==========
+class CharacterSet:
+    """Manages character to index mapping"""
+    def __init__(self):
+        self.chars = ['BLANK', 'PAD', 'SOS', ' ']  # Special tokens + Space
+        self.char2idx = {'BLANK': 0, 'PAD': 1, 'SOS': 2, ' ': 3}
+        self.idx2char = {0: 'BLANK', 1: 'PAD', 2: 'SOS', 3: ' '}
+        
+    def add_chars(self, text):
+        """Add new characters from text"""
+        for char in text:
+            if char not in self.char2idx:
+                idx = len(self.chars)
+                self.chars.append(char)
+                self.char2idx[char] = idx
+                self.idx2char[idx] = char
+    
+    def encode(self, text):
+        """Convert text to indices"""
+        return [self.char2idx.get(char, 0) for char in text]
+    
+    def decode(self, indices):
+        """Convert indices to text (CTC decode)"""
+        chars = []
+        prev_idx = None
+        for idx in indices:
+            if idx > 2 and idx != prev_idx:  # Skip blank, pad, sos
+                chars.append(self.idx2char.get(idx, ''))
+            prev_idx = idx
+        return ''.join(chars)
+    
+    def __len__(self):
+        return len(self.chars)
+    
+    def save(self, path):
+        with open(path, 'w', encoding='utf-8') as f:
+            for char in self.chars:
+                f.write(char + '\n')
+    
+    @classmethod
+    def load(cls, path):
+        charset = cls()
+        with open(path, 'r', encoding='utf-8') as f:
+            chars = [line.rstrip('\n') for line in f]
+            charset.chars = chars
+            charset.char2idx = {char: idx for idx, char in enumerate(chars)}
+            charset.idx2char = {idx: char for idx, char in enumerate(chars)}
+        return charset
+    
+    @classmethod
+    def from_checkpoint(cls, checkpoint):
+        """Load charset from model checkpoint dictionary"""
+        charset = cls()
+        if 'charset' in checkpoint:
+            charset.chars = checkpoint['charset']
+            charset.char2idx = {char: idx for idx, char in enumerate(charset.chars)}
+            charset.idx2char = {idx: char for idx, char in enumerate(charset.chars)}
+        return charset
+
+def save_checkpoint(model, charset, optimizer, epoch, val_loss, accuracy, path):
+    """Save model checkpoint with charset included"""
+    torch.save({
+        'epoch': epoch,
+        'model_state_dict': model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+        'val_loss': val_loss,
+        'accuracy': accuracy,
+        'charset': charset.chars
+    }, path)
+
+# ========== LIGHTWEIGHT OCR MODEL ==========
+class LightweightOCR(nn.Module):
+    """
+    Lightweight OCR model (~13MB)
+    Similar to the reference model you showed
+    """
+    def __init__(self, num_chars, hidden_size=256):
+        super(LightweightOCR, self).__init__()
+        
+        # CNN backbone (lighter than before)
+        self.cnn = nn.Sequential(
+            # Block 1: [B, 1, 32, W] -> [B, 32, 32, W]
+            self._conv_block(1, 32, kernel_size=3, stride=1, padding=1),
+            nn.MaxPool2d(kernel_size=2, stride=2),  # -> [B, 32, 16, W/2]
+            
+            # Block 2: [B, 32, 16, W/2] -> [B, 64, 16, W/2]
+            self._conv_block(32, 64, kernel_size=3, stride=1, padding=1),
+            nn.MaxPool2d(kernel_size=2, stride=2),  # -> [B, 64, 8, W/4]
+            
+            # Block 3: [B, 64, 8, W/4] -> [B, 128, 8, W/4]
+            self._conv_block(64, 128, kernel_size=3, stride=1, padding=1),
+            self._conv_block(128, 128, kernel_size=3, stride=1, padding=1),
+            nn.MaxPool2d(kernel_size=(2, 1), stride=(2, 1)),  # -> [B, 128, 4, W/4]
+            
+            # Block 4: [B, 128, 4, W/4] -> [B, 256, 4, W/4]
+            self._conv_block(128, 256, kernel_size=3, stride=1, padding=1),
+            self._conv_block(256, 256, kernel_size=3, stride=1, padding=1),
+            nn.MaxPool2d(kernel_size=(4, 1), stride=(4, 1)),  # -> [B, 256, 1, W/4]
+        )
+        
+        # Two-layer LSTM with intermediate projection
+        self.lstm1 = nn.LSTM(
+            256, hidden_size, 
+            bidirectional=True, 
+            batch_first=True
+        )
+        
+        self.intermediate_linear = nn.Linear(hidden_size * 2, hidden_size)
+        
+        self.lstm2 = nn.LSTM(
+            hidden_size, hidden_size, 
+            bidirectional=True, 
+            batch_first=True
+        )
+        
+        # Output layer
+        self.fc = nn.Linear(hidden_size * 2, num_chars)
+    
+    def _conv_block(self, in_ch, out_ch, **kwargs):
+        """Convolutional block with BatchNorm and ReLU"""
+        return nn.Sequential(
+            nn.Conv2d(in_ch, out_ch, bias=False, **kwargs),
+            nn.BatchNorm2d(out_ch),
+            nn.ReLU(inplace=True),
+        )
+    
+    def forward(self, x):
+        """
+        Args:
+            x: [B, 1, H, W]
+        Returns:
+            output: [T, B, num_chars]
+        """
+        # CNN features
+        features = self.cnn(x)  # [B, 256, 1, W']
+        
+        # Remove height dimension
+        features = features.squeeze(2)  # [B, 256, W']
+        features = features.permute(0, 2, 1)  # [B, W', 256]
+        
+        # First LSTM layer
+        rnn_out, _ = self.lstm1(features)  # [B, W', hidden*2]
+        
+        # Intermediate projection
+        rnn_out = self.intermediate_linear(rnn_out)  # [B, W', hidden]
+        rnn_out = torch.relu(rnn_out)
+        
+        # Second LSTM layer
+        rnn_out, _ = self.lstm2(rnn_out)  # [B, W', hidden*2]
+        
+        # Fully connected
+        logits = self.fc(rnn_out)  # [B, W', num_chars]
+        
+        # Permute for CTC
+        return logits.permute(1, 0, 2)  # [W', B, num_chars]

kiri_ocr/renderer.py

@@ -0,0 +1,193 @@
+import cv2
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+from pathlib import Path
+
+class DocumentRenderer:
+    """Render OCR results on image"""
+    
+    def __init__(self, font_path=None, font_size=12):
+        self.font_size = font_size
+        
+        # Try to load font for text rendering
+        self.font = None
+        if font_path and Path(font_path).exists():
+            try:
+                self.font = ImageFont.truetype(font_path, font_size)
+            except:
+                pass
+        
+        # Fallback fonts
+        if self.font is None:
+            # Priority list including Khmer fonts
+            candidate_fonts = [
+                'fonts/KhmerOSbattambang.ttf',
+                'fonts/Battambang-Regular.ttf',
+                'fonts/NotoSansKhmer-Regular.ttf',
+                'Arial.ttf',
+                'DejaVuSans.ttf',
+                'NotoSans-Regular.ttf'
+            ]
+            
+            # Check for any TTF in fonts/ directory
+            if Path('fonts').exists():
+                candidate_fonts = [str(f) for f in Path('fonts').glob('*.ttf')] + candidate_fonts
+            
+            for font_name in candidate_fonts:
+                try:
+                    self.font = ImageFont.truetype(font_name, font_size)
+                    break
+                except:
+                    continue
+    
+    def draw_boxes(self, image_path, results, output_path='output_boxes.png'):
+        """Draw bounding boxes only"""
+        img = cv2.imread(str(image_path))
+        
+        for result in results:
+            x, y, w, h = result['box']
+            conf = result['confidence']
+            
+            # Color based on confidence
+            if conf > 0.9:
+                color = (0, 255, 0)  # Green
+            elif conf > 0.7:
+                color = (0, 165, 255)  # Orange
+            else:
+                color = (0, 0, 255)  # Red
+            
+            # Draw box
+            cv2.rectangle(img, (x, y), (x+w, y+h), color, 2)
+            
+            # Draw line number
+            if 'line_number' in result:
+                cv2.putText(img, str(result['line_number']), (x-5, y-5),
+                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+        
+        cv2.imwrite(output_path, img)
+        print(f"\n✓ Boxes saved to {output_path}")
+        
+        return img
+    
+    def draw_results(self, image_path, results, output_path='output_ocr.png',
+                    show_text=True, show_confidence=True):
+        """Draw boxes with recognized text"""
+        img = cv2.imread(str(image_path))
+        
+        # Convert to PIL for better text rendering
+        img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+        draw = ImageDraw.Draw(img_pil)
+        
+        for result in results:
+            x, y, w, h = result['box']
+            text = result['text']
+            conf = result['confidence']
+            
+            # Color based on confidence
+            if conf > 0.9:
+                color = (0, 255, 0)
+            elif conf > 0.7:
+                color = (255, 165, 0)
+            else:
+                color = (255, 0, 0)
+            
+            # Draw box
+            draw.rectangle([x, y, x+w, y+h], outline=color, width=2)
+            
+            if show_text:
+                # Prepare label
+                label = text[:50]  # Limit length
+                if show_confidence:
+                    label += f" ({conf*100:.0f}%)"
+                
+                # Background for text
+                if self.font:
+                    try:
+                        # Calculate text position
+                        left, top, right, bottom = draw.textbbox((0, 0), label, font=self.font)
+                        text_height = bottom - top
+                        text_y = y - text_height - 5
+                        
+                        bbox = draw.textbbox((x, text_y), label, font=self.font)
+                        # Add padding
+                        bbox = (bbox[0]-2, bbox[1]-2, bbox[2]+2, bbox[3]+2)
+                        
+                        draw.rectangle(bbox, fill=color)
+                        draw.text((x, text_y), label, fill=(255, 255, 255), font=self.font)
+                    except:
+                        # Fallback
+                        img_cv = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
+                        cv2.putText(img_cv, label, (x, y-10),
+                                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
+                        img_pil = Image.fromarray(cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB))
+                        draw = ImageDraw.Draw(img_pil)
+        
+        # Save
+        img_pil.save(output_path)
+        print(f"✓ Results saved to {output_path}")
+        
+        return img_pil
+    
+    def create_report(self, image_path, results, output_path='ocr_report.html'):
+        """Create HTML report"""
+        html = f"""
+<!DOCTYPE html>
+<html>
+<head>
+    <meta charset="UTF-8">
+    <title>OCR Report</title>
+    <style>
+        body {{ font-family: Arial, sans-serif; margin: 20px; }}
+        .header {{ background: #4CAF50; color: white; padding: 20px; border-radius: 5px; }}
+        .result {{ padding: 10px; margin: 10px 0; border-left: 4px solid #4CAF50; }}
+        .conf-high {{ border-color: #4CAF50; background: #f1f8f4; }}
+        .conf-medium {{ border-color: #FF9800; background: #fff8f1; }}
+        .conf-low {{ border-color: #F44336; background: #fef1f1; }}
+        .text {{ font-size: 18px; font-weight: bold; margin-bottom: 5px; }}
+        .confidence {{ color: #666; font-size: 14px; }}
+        .stats {{ background: #f5f5f5; padding: 15px; border-radius: 5px; margin: 20px 0; }}
+        .full-text {{ background: #f5f5f5; padding: 20px; white-space: pre-wrap; 
+                     font-family: monospace; border-radius: 5px; }}
+    </style>
+</head>
+<body>
+    <div class="header">
+        <h1>📄 OCR Report</h1>
+        <p>Document: {Path(image_path).name}</p>
+    </div>
+    
+    <div class="stats">
+        <strong>Statistics:</strong><br>
+        Total Regions: {len(results)}<br>
+        Average Confidence: {np.mean([r['confidence'] for r in results])*100:.2f}%<br>
+        High Confidence (>90%): {sum(1 for r in results if r['confidence'] > 0.9)}<br>
+        Medium Confidence (70-90%): {sum(1 for r in results if 0.7 < r['confidence'] <= 0.9)}<br>
+        Low Confidence (<70%): {sum(1 for r in results if r['confidence'] <= 0.7)}
+    </div>
+    
+    <h2>📝 Full Text</h2>
+    <div class="full-text">{"<br>".join([r['text'] for r in results])}</div>
+    
+    <h2>📋 Detailed Results</h2>
+"""
+        
+        for i, result in enumerate(results, 1):
+            conf = result['confidence']
+            conf_class = 'conf-high' if conf > 0.9 else ('conf-medium' if conf > 0.7 else 'conf-low')
+            
+            html += f"""
+    <div class="result {conf_class}">
+        <div class="text">{i}. {result['text']}</div>
+        <div class="confidence">Confidence: {conf*100:.2f}% | Box: {result['box']}</div>
+    </div>
+"""
+        
+        html += """
+</body>
+</html>
+"""
+        
+        with open(output_path, 'w', encoding='utf-8') as f:
+            f.write(html)
+        
+        print(f"✓ Report saved to {output_path}")