# -*- coding: utf-8 -*-
"""crnn.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1Sg_6ocKWq5Yhjc2jxwCjbSU5oi5-d2YM
"""

import re  # Regular expressions for extracting numerical values
import cv2  # OpenCV for image processing
import pytesseract  # OCR engine
from pytesseract import Output  # Output type for pytesseract data extraction
import requests  # To fetch the image from a URL
from io import BytesIO  # To handle image byte streams
from PIL import Image  # To open images in Python using PIL
import numpy as np  # NumPy for array manipulations
pytesseract.pytesseract.tesseract_cmd = '/opt/homebrew/bin/tesseract'

import torch
import torch.nn as nn

# Define the character set (including space and special characters)
charset = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ .,%'
char_to_idx = {char: idx + 1 for idx, char in enumerate(charset)}  # Start indices from 1
idx_to_char = {idx: char for char, idx in char_to_idx.items()}
class CRNN(nn.Module):
    def __init__(self, num_classes):
        super(CRNN, self).__init__()

        # Convolutional layers
        self.cnn = nn.Sequential(
            # Conv Layer 1
            nn.Conv2d(1, 64, kernel_size=3, padding=1),  # Grayscale images have 1 channel
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),  # Reduces height and width by half

            # Conv Layer 2
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),

            # Conv Layer 3
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),

            # Conv Layer 4
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d((2, 1)),  # Only halves height

            # Conv Layer 5
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),

            # Conv Layer 6
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d((2, 1)),

            # Conv Layer 7
            nn.Conv2d(512, 512, kernel_size=2),  # KxK = 2x2
            nn.ReLU(inplace=True)
        )

        # LSTM layers
        self.lstm = nn.LSTM(
            input_size=512,     # Should match the output channels of CNN
            hidden_size=256,
            num_layers=2,
            bidirectional=True,
            batch_first=True
        )

        # Fully connected layer
        self.fc = nn.Linear(512, num_classes)  # 512 = 256 * 2 (bidirectional)

    def forward(self, x):
        # Pass through CNN layers
        conv = self.cnn(x)

        # Prepare for LSTM
        b, c, h, w = conv.size()
        assert h == 1, f"Expected height == 1 after CNN layers, got {h}"
        conv = conv.squeeze(2)  # [batch, channels, width]
        conv = conv.permute(0, 2, 1)  # [batch, width, channels]

        # Pass through LSTM
        rnn_out, _ = self.lstm(conv)  # rnn_out: [batch, width, hidden_size*2]

        # Pass through fully connected layer
        output = self.fc(rnn_out)  # [batch, width, num_classes]

        # Transpose to match CTC Loss requirements
        output = output.permute(1, 0, 2)  # [width, batch, num_classes]

        return output

import pandas as pd
from sklearn.model_selection import train_test_split

# Load the full dataset
full_data = pd.read_csv('/Users/kartik/Downloads/student_resource 3/dataset/train.csv')

# Split into training and test sets
train_data, test_data = train_test_split(full_data, test_size=0.2, random_state=42)

# Save the split datasets (optional)
# train_data.to_csv('train_split.csv', index=False)
# test_data.to_csv('test_split.csv', index=False)

import pandas as pd
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import requests

class OCRDataset(Dataset):
    def __init__(self, data_frame, transform=None):
        self.data_frame = data_frame.reset_index(drop=True)
        self.transform = transform

    def __len__(self):
        return len(self.data_frame)

    def __getitem__(self, idx):
        img_url = self.data_frame.iloc[idx]['image_link']
        label = self.data_frame.iloc[idx]['entity_value']
        # Fetch the image
        try:
            response = requests.get(img_url)
            img = Image.open(BytesIO(response.content)).convert('L')  # Grayscale
        except Exception as e:
            # Handle failed image download
            img = Image.new('L', (128, 32))  # Placeholder image
            print(f"Failed to load image at index {idx}: {e}")

        if self.transform:
            img = self.transform(img)
        # Encode the label
        label_idx = [char_to_idx.get(char, 0) for char in label]  # 0 is for unknown chars
        label_tensor = torch.tensor(label_idx, dtype=torch.long)
        return img, label_tensor, len(label_idx)

from torchvision import transforms

transform = transforms.Compose([
    transforms.Resize((32, 128)),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

def collate_fn(batch):
    images, labels, label_lengths = zip(*batch)
    images = torch.stack(images)
    labels = torch.cat(labels)
    label_lengths = torch.tensor(label_lengths, dtype=torch.long)
    batch_size = images.size(0)
    target_lengths = label_lengths
    return images, labels, target_lengths

from symspellpy import SymSpell, Verbosity

# Spell correction setup
sym_spell = SymSpell(max_dictionary_edit_distance=2)
sym_spell.create_dictionary_entry('gram', 1)
sym_spell.create_dictionary_entry('kilogram', 1)

# Correct OCR output using spell correction
def correct_text(text):
    suggestions = sym_spell.lookup_compound(text, max_edit_distance=2)
    return suggestions[0].term if suggestions else text

# Instantiate the datasets
train_dataset = OCRDataset(data_frame=train_data, transform=transform)
test_dataset = OCRDataset(data_frame=test_data, transform=transform)

# Create DataLoaders
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, collate_fn=collate_fn)

from tqdm import tqdm
import Levenshtein
import torch.optim as optim
# Initialize model, loss function, and optimizer
num_classes = len(charset) + 1  # +1 for the blank label required for CTC
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = CRNN(num_classes).to(device)
criterion = nn.CTCLoss(blank=0, zero_infinity=True)
optimizer = optim.Adam(model.parameters(), lr=0.0001)

def cer(prediction, ground_truth):
    prediction = prediction.replace(' ', '')
    ground_truth = ground_truth.replace(' ', '')
    distance = Levenshtein.distance(prediction, ground_truth)
    return distance / max(len(ground_truth), 1)

def wer(prediction, ground_truth):
    pred_words = prediction.split()
    gt_words = ground_truth.split()
    distance = Levenshtein.distance(' '.join(pred_words), ' '.join(gt_words))
    return distance / max(len(gt_words), 1)

def train_model(model, train_loader, test_loader, criterion, optimizer, num_epochs=10):
    for epoch in range(num_epochs):
        model.train()
        total_loss = 0
        progress_bar = tqdm(train_loader, desc=f"Epoch [{epoch+1}/{num_epochs}]", leave=False)
        for images, labels, label_lengths in progress_bar:
            images = images.to(device)
            labels = labels.to(device)
            label_lengths = label_lengths.to(device)
            optimizer.zero_grad()
            outputs = model(images)

            batch_size = images.size(0)
            seq_len = outputs.size(0)
            input_lengths = torch.full(size=(batch_size,), fill_value=seq_len, dtype=torch.long).to(device)

            loss = criterion(outputs, labels, input_lengths, label_lengths)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

            progress_bar.set_postfix(loss=loss.item())

        avg_loss = total_loss / len(train_loader)
        print(f"Epoch [{epoch+1}/{num_epochs}], Training Loss: {avg_loss:.4f}")

        # Validation
        model.eval()
        with torch.no_grad():
            cer_total = 0
            wer_total = 0
            num_samples = 0
            for images, labels, label_lengths in test_loader:
                images = images.to(device)
                labels = labels.to(device)
                label_lengths = label_lengths.to(device)
                outputs = model(images)

                # Decode predictions
                decoded_preds = decode_predictions(outputs)
                # Prepare ground truth texts
                ground_truths = []
                index = 0
                labels = labels.cpu()
                for length in label_lengths.cpu():
                    gt_text = ''.join([idx_to_char.get(idx.item(), '') for idx in labels[index:index+length]])
                    ground_truths.append(gt_text)
                    index += length

                # Calculate CER and WER
                for pred, gt in zip(decoded_preds, ground_truths):
                    cer_score = cer(pred, gt)
                    wer_score = wer(pred, gt)
                    cer_total += cer_score
                    wer_total += wer_score
                    num_samples += 1

            avg_cer = cer_total / num_samples
            avg_wer = wer_total / num_samples
            print(f"Validation CER: {avg_cer:.4f}, WER: {avg_wer:.4f}")

    print("Training completed.")

train_model(model, train_loader, test_loader, criterion, optimizer, num_epochs=10)
torch.save(model.state_dict(), 'crnn_model.pth')

def decode_predictions(outputs):
    # Outputs: [seq_len, batch_size, num_classes]
    outputs = outputs.permute(1, 0, 2)  # [batch_size, seq_len, num_classes]
    batch_size = outputs.size(0)
    _, max_indices = torch.max(outputs, dim=2)
    # Convert to text
    texts = []
    for i in range(batch_size):
        indices = max_indices[i].detach().cpu().numpy()
        # Collapse repeated characters and remove blanks
        chars = []
        prev_idx = -1
        for idx in indices:
            if idx != prev_idx and idx != 0:
                chars.append(idx_to_char.get(idx, ''))
            prev_idx = idx
        text = ''.join(chars)
        texts.append(text)
    return texts

import Levenshtein
def cer(prediction, ground_truth):
    # Simplified CER calculation
    prediction = prediction.replace(' ', '')
    ground_truth = ground_truth.replace(' ', '')
    distance = Levenshtein.distance(prediction, ground_truth)
    return distance / max(len(ground_truth), 1)

def wer(prediction, ground_truth):
    # Simplified WER calculation
    pred_words = prediction.split()
    gt_words = ground_truth.split()
    distance = Levenshtein.distance(pred_words, gt_words)
    return distance / max(len(gt_words), 1)


i want to make the process of training faster bring logging and stuff whatever makes it faster 
i am using a amc m1