Upload 13 files

requirements.txt

@@ -0,0 +1,13 @@
+torch>=2.0.0
+torchvision>=0.15.0
+torchaudio
+opencv-python-headless
+numpy
+pandas
+matplotlib
+tqdm
+scikit-learn
+gradio
+streamlit
+pillow
+albumentations

src/data/dataset.py

@@ -0,0 +1,77 @@
+import torch
+from torch.utils.data import Dataset, DataLoader
+import pandas as pd
+import os
+from PIL import Image
+
+class IAMDataset(Dataset):
+    def __init__(self, data_dir, csv_file, transform=None):
+        """
+        Args:
+            data_dir (str): Path to directory containing IAM word images.
+            csv_file (str): Path to CSV file containing 'filename' and 'text'.
+            transform (callable, optional): Optional transform to be applied.
+        """
+        self.data_dir = data_dir
+        # Assuming CSV has columns: 'filename' and 'text'
+        self.annotations = pd.read_csv(csv_file)
+        self.transform = transform
+        
+        # Build vocabulary
+        self.vocab = self._build_vocab()
+        self.char_to_idx = {char: idx + 1 for idx, char in enumerate(self.vocab)} # 0 is reserved for CTC blank
+        self.idx_to_char = {idx: char for char, idx in self.char_to_idx.items()}
+        self.num_classes = len(self.vocab) + 1 # +1 for CTC blank
+        
+    def _build_vocab(self):
+        chars = set()
+        for text in self.annotations['text']:
+            if pd.notna(text):
+                chars.update(list(str(text)))
+        return sorted(list(chars))
+
+    def __len__(self):
+        return len(self.annotations)
+
+    def __getitem__(self, idx):
+        if torch.is_tensor(idx):
+            idx = idx.tolist()
+
+        img_name = os.path.join(self.data_dir, str(self.annotations.iloc[idx]['filename']))
+        
+        try:
+            image = Image.open(img_name).convert('L') # Convert to grayscale
+        except FileNotFoundError:
+            # Handle missing files gracefully in a real scenario
+            image = Image.new('L', (1024, 32), color=255)
+            
+        text = str(self.annotations.iloc[idx]['text'])
+        
+        if pd.isna(text):
+            text = ""
+
+        if self.transform:
+            image = self.transform(image)
+
+        # Convert text to tensor of indices
+        encoded_text = [self.char_to_idx[char] for char in text if char in self.char_to_idx]
+        text_tensor = torch.tensor(encoded_text, dtype=torch.long)
+
+        return image, text_tensor, len(encoded_text)
+
+# Collate function for DataLoader to handle variable length sequences
+def collate_fn(batch):
+    images, texts, text_lengths = zip(*batch)
+    
+    # Stack images
+    images = torch.stack(images)
+    
+    # Pad texts to max length in batch
+    texts_padded = torch.nn.utils.rnn.pad_sequence(texts, batch_first=True, padding_value=0)
+    
+    text_lengths = torch.tensor(text_lengths, dtype=torch.long)
+    
+    return images, texts_padded, text_lengths
+
+if __name__ == "__main__":
+    print("Dataset module ready.")

src/data/download_dataset.py

@@ -0,0 +1,52 @@
+import os
+import pandas as pd
+from datasets import load_dataset
+from tqdm import tqdm
+
+def download_and_prepare_iam():
+    print("Downloading IAM-line dataset from Hugging Face...")
+    # Loading the dataset from Hugging Face (approx 266 MB)
+    dataset = load_dataset("Sj122702/IAM-line")
+    
+    data_dir = "data/iam_words"
+    os.makedirs(data_dir, exist_ok=True)
+    
+    print(f"Saving images to {data_dir} and creating labels.csv...")
+    
+    metadata = []
+    
+    # We will process the 'train' split for demonstration
+    # You can expand this to validation and test splits as well
+    split = 'train'
+    
+    for idx, item in enumerate(tqdm(dataset[split])):
+        # The dataset contains 'image' and 'text'
+        image = item['image']
+        text = item['text']
+        
+        # Save image
+        filename = f"img_{split}_{idx}.png"
+        filepath = os.path.join(data_dir, filename)
+        
+        # Some images might be in different modes, convert to grayscale
+        image = image.convert("L")
+        image.save(filepath)
+        
+        # Add to metadata
+        metadata.append({
+            "filename": filename,
+            "text": text
+        })
+        
+    # Save metadata to CSV
+    csv_path = "data/labels.csv"
+    df = pd.DataFrame(metadata)
+    df.to_csv(csv_path, index=False)
+    
+    print(f"\nDataset prepared successfully!")
+    print(f"Total images saved: {len(metadata)}")
+    print(f"Images location: {data_dir}/")
+    print(f"Labels CSV location: {csv_path}")
+
+if __name__ == "__main__":
+    download_and_prepare_iam()

src/models/crnn.py

@@ -0,0 +1,81 @@
+import torch
+import torch.nn as nn
+
+class CRNN(nn.Module):
+    def __init__(self, img_channel, img_height, img_width, num_class, 
+                 map_to_seq_hidden=64, rnn_hidden=256, leaky_relu=False):
+        super(CRNN, self).__init__()
+        
+        self.cnn, (output_channel, output_height, output_width) = \
+            self._cnn_backbone(img_channel, img_height, img_width, leaky_relu)
+        
+        self.map_to_seq = nn.Linear(output_channel * output_height, map_to_seq_hidden)
+        
+        self.rnn1 = nn.LSTM(map_to_seq_hidden, rnn_hidden, bidirectional=True, batch_first=True)
+        self.rnn2 = nn.LSTM(rnn_hidden * 2, rnn_hidden, bidirectional=True, batch_first=True)
+        
+        self.dense = nn.Linear(rnn_hidden * 2, num_class)
+        
+    def _cnn_backbone(self, img_channel, img_height, img_width, leaky_relu):
+        assert img_height % 16 == 0
+        assert img_width % 4 == 0
+        
+        channels = [img_channel, 64, 128, 256, 256, 512, 512, 512]
+        kernel_sizes = [3, 3, 3, 3, 3, 3, 2]
+        strides = [1, 1, 1, 1, 1, 1, 1]
+        paddings = [1, 1, 1, 1, 1, 1, 0]
+        
+        cnn = nn.Sequential()
+        
+        def conv_relu(i, batch_normalization=False):
+            n_in = channels[i]
+            n_out = channels[i+1]
+            cnn.add_module(f'conv{i}', nn.Conv2d(n_in, n_out, kernel_sizes[i], strides[i], paddings[i]))
+            if batch_normalization:
+                cnn.add_module(f'batchnorm{i}', nn.BatchNorm2d(n_out))
+            if leaky_relu:
+                cnn.add_module(f'relu{i}', nn.LeakyReLU(0.2, inplace=True))
+            else:
+                cnn.add_module(f'relu{i}', nn.ReLU(inplace=True))
+                
+        conv_relu(0)
+        cnn.add_module('pooling0', nn.MaxPool2d(kernel_size=2, stride=2)) # 64x16x64
+        conv_relu(1)
+        cnn.add_module('pooling1', nn.MaxPool2d(kernel_size=2, stride=2)) # 128x8x32
+        conv_relu(2, True)
+        conv_relu(3)
+        cnn.add_module('pooling2', nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 1), padding=(0, 1))) # 256x4x33
+        conv_relu(4, True)
+        conv_relu(5)
+        cnn.add_module('pooling3', nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 1), padding=(0, 1))) # 512x2x34
+        conv_relu(6, True) # 512x1x33
+        
+        output_channel, output_height, output_width = channels[-1], img_height // 16 - 1, img_width // 4 + 1
+        return cnn, (output_channel, output_height, output_width)
+        
+    def forward(self, images):
+        # shape of images: (batch, channel, height, width)
+        conv = self.cnn(images)
+        batch, channel, height, width = conv.size()
+        
+        conv = conv.view(batch, channel * height, width)
+        conv = conv.permute(0, 2, 1) # (batch, width, channel*height)
+        
+        seq = self.map_to_seq(conv)
+        
+        recurrent, _ = self.rnn1(seq)
+        recurrent, _ = self.rnn2(recurrent)
+        
+        output = self.dense(recurrent)
+        
+        # Log softmax for CTC loss
+        # Note: PyTorch's CTCLoss expects inputs of shape (input_length, batch_size, num_classes)
+        # So we permute it if we are returning it for CTC loss calculation directly
+        return output.log_softmax(2)
+
+if __name__ == '__main__':
+    # Test model
+    dummy_input = torch.randn(1, 1, 32, 1024)
+    model = CRNN(img_channel=1, img_height=32, img_width=1024, num_class=80)
+    output = model(dummy_input)
+    print(f"Output shape: {output.shape}") # Expected: (1, 33, 80)

src/models/gan.py

@@ -0,0 +1,74 @@
+import torch
+import torch.nn as nn
+
+# Simple DCGAN-style architecture for generating word images (1x32x1024)
+
+class Generator(nn.Module):
+    def __init__(self, latent_dim=100, channels=1):
+        super(Generator, self).__init__()
+        
+        # Input: latent_dim, mapping to 4x128 map initially
+        self.init_size_h = 4
+        self.init_size_w = 128
+        self.l1 = nn.Sequential(nn.Linear(latent_dim, 128 * self.init_size_h * self.init_size_w))
+
+        self.conv_blocks = nn.Sequential(
+            nn.BatchNorm2d(128),
+            nn.Upsample(scale_factor=2), # 8x256
+            nn.Conv2d(128, 128, 3, stride=1, padding=1),
+            nn.BatchNorm2d(128, 0.8),
+            nn.LeakyReLU(0.2, inplace=True),
+            
+            nn.Upsample(scale_factor=2), # 16x512
+            nn.Conv2d(128, 64, 3, stride=1, padding=1),
+            nn.BatchNorm2d(64, 0.8),
+            nn.LeakyReLU(0.2, inplace=True),
+            
+            nn.Upsample(scale_factor=2), # 32x1024
+            nn.Conv2d(64, channels, 3, stride=1, padding=1),
+            nn.Tanh(), # Output [-1, 1]
+        )
+
+    def forward(self, z):
+        out = self.l1(z)
+        out = out.view(out.shape[0], 128, self.init_size_h, self.init_size_w)
+        img = self.conv_blocks(out)
+        return img
+
+
+class Discriminator(nn.Module):
+    def __init__(self, channels=1):
+        super(Discriminator, self).__init__()
+
+        def discriminator_block(in_filters, out_filters, bn=True):
+            block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)]
+            if bn:
+                block.append(nn.BatchNorm2d(out_filters, 0.8))
+            return block
+
+        self.model = nn.Sequential(
+            *discriminator_block(channels, 16, bn=False), # 16x512
+            *discriminator_block(16, 32), # 8x256
+            *discriminator_block(32, 64), # 4x128
+            *discriminator_block(64, 128), # 2x64
+        )
+
+        # The height and width of downsampled image
+        ds_size_h = 32 // 2**4
+        ds_size_w = 1024 // 2**4
+        self.adv_layer = nn.Sequential(nn.Linear(128 * ds_size_h * ds_size_w, 1), nn.Sigmoid())
+
+    def forward(self, img):
+        out = self.model(img)
+        out = out.view(out.shape[0], -1)
+        validity = self.adv_layer(out)
+        return validity
+
+if __name__ == "__main__":
+    z = torch.randn(1, 100)
+    G = Generator()
+    D = Discriminator()
+    fake_img = G(z)
+    validity = D(fake_img)
+    print(f"Generator output shape: {fake_img.shape}")
+    print(f"Discriminator output shape: {validity.shape}")

src/training/train_crnn.py

@@ -0,0 +1,104 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from torchvision import transforms
+from PIL import Image
+import os
+import sys
+
+# Add project root to path
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
+
+from src.data.dataset import IAMDataset, collate_fn
+from src.models.crnn import CRNN
+
+# Define transforms
+transform = transforms.Compose([
+    transforms.Resize((32, 1024)),
+    transforms.ToTensor(),
+    transforms.Normalize((0.5,), (0.5,))
+])
+
+def train_baseline(model, dataloader, optimizer, criterion, device, epochs=10, start_epoch=0):
+    model.train()
+    
+    for epoch in range(start_epoch, epochs):
+        total_loss = 0
+        for i, (images, texts, text_lengths) in enumerate(dataloader):
+            images = images.to(device)
+            texts = texts.to(device)
+            
+            optimizer.zero_grad()
+            
+            # Forward pass
+            preds = model(images)
+            
+            # CTCLoss expects (sequence_length, batch_size, num_classes)
+            preds = preds.permute(1, 0, 2)
+            
+            # Calculate lengths for CTC Loss
+            input_lengths = torch.full(size=(preds.size(1),), fill_value=preds.size(0), dtype=torch.long)
+            
+            # CTCLoss expects concatenated targets, not padded 2D tensor
+            # Flatten all target sequences into 1D
+            targets_list = []
+            for i in range(texts.size(0)):
+                target_seq = texts[i][:text_lengths[i]]
+                targets_list.append(target_seq)
+            targets_concat = torch.cat(targets_list)
+            
+            loss = criterion(preds, targets_concat, input_lengths, text_lengths)
+            
+            # Backward pass
+            loss.backward()
+            optimizer.step()
+            
+            total_loss += loss.item()
+            
+            if i % 10 == 0:
+                print(f"Epoch [{epoch+1}/{epochs}], Step [{i}/{len(dataloader)}], Loss: {loss.item():.4f}")
+                
+        print(f"Epoch {epoch+1} Average Loss: {total_loss/len(dataloader):.4f}")
+        
+        # Save checkpoint
+        os.makedirs('weights', exist_ok=True)
+        torch.save(model.state_dict(), f'weights/crnn_baseline_epoch_{epoch+1}.pth')
+        
+    return model
+
+if __name__ == "__main__":
+    print("Starting CRNN Baseline Training...")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+    
+    # Setup Data
+    data_dir = 'data/iam_words'
+    csv_file = 'data/labels.csv'
+    
+    dataset = IAMDataset(data_dir=data_dir, csv_file=csv_file, transform=transform)
+    dataloader = DataLoader(dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)
+    
+    # Setup Model
+    num_classes = dataset.num_classes
+    model = CRNN(img_channel=1, img_height=32, img_width=1024, num_class=num_classes).to(device)
+    
+    # Resume from checkpoint if exists
+    start_epoch = 0
+    # Find the latest checkpoint
+    import glob
+    checkpoints = glob.glob('weights/crnn_baseline_epoch_*.pth')
+    if checkpoints:
+        checkpoints.sort(key=lambda x: int(os.path.basename(x).split('_')[-1].split('.')[0]))
+        latest_checkpoint = checkpoints[-1]
+        start_epoch = int(os.path.basename(latest_checkpoint).split('_')[-1].split('.')[0])
+        print(f"Resuming training from {latest_checkpoint} (epoch {start_epoch})")
+        model.load_state_dict(torch.load(latest_checkpoint, map_location=device))
+    
+    # Setup Optimizer & Loss
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+    criterion = nn.CTCLoss(blank=0, zero_infinity=True)
+    
+    # Train
+    train_baseline(model, dataloader, optimizer, criterion, device, epochs=30, start_epoch=start_epoch)
+    print("Training complete!")

src/training/train_gan.py

@@ -0,0 +1,112 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from torchvision import transforms
+import os
+import sys
+
+# Add project root to path
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
+
+from src.data.dataset import IAMDataset, collate_fn
+from src.models.gan import Generator, Discriminator
+
+# Define transforms for GAN (needs to be slightly different, just standard normalization)
+transform = transforms.Compose([
+    transforms.Resize((32, 1024)),
+    transforms.ToTensor(),
+    transforms.Normalize((0.5,), (0.5,))
+])
+
+def train_gan(generator, discriminator, dataloader, epochs, device, latent_dim=100, start_epoch=0):
+    criterion = nn.BCELoss()
+    
+    optimizer_G = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))
+    optimizer_D = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))
+    
+    generator.train()
+    discriminator.train()
+    
+    for epoch in range(start_epoch, epochs):
+        for i, (imgs, _, _) in enumerate(dataloader):
+            
+            batch_size = imgs.size(0)
+            # Adversarial ground truths
+            valid = torch.ones(batch_size, 1, requires_grad=False).to(device)
+            fake = torch.zeros(batch_size, 1, requires_grad=False).to(device)
+            
+            # Configure input
+            real_imgs = imgs.to(device)
+            
+            # -----------------
+            #  Train Generator
+            # -----------------
+            optimizer_G.zero_grad()
+            
+            # Sample noise as generator input
+            z = torch.randn(batch_size, latent_dim).to(device)
+            
+            # Generate a batch of images
+            gen_imgs = generator(z)
+            
+            # Loss measures generator's ability to fool the discriminator
+            g_loss = criterion(discriminator(gen_imgs), valid)
+            
+            g_loss.backward()
+            optimizer_G.step()
+            
+            # ---------------------
+            #  Train Discriminator
+            # ---------------------
+            optimizer_D.zero_grad()
+            
+            # Measure discriminator's ability to classify real from generated samples
+            real_loss = criterion(discriminator(real_imgs), valid)
+            fake_loss = criterion(discriminator(gen_imgs.detach()), fake)
+            d_loss = (real_loss + fake_loss) / 2
+            
+            d_loss.backward()
+            optimizer_D.step()
+            
+            if i % 50 == 0:
+                print(f"[Epoch {epoch+1}/{epochs}] [Batch {i}/{len(dataloader)}] [D loss: {d_loss.item():.4f}] [G loss: {g_loss.item():.4f}]")
+                
+        # Save checkpoints
+        os.makedirs('weights', exist_ok=True)
+        torch.save(generator.state_dict(), f'weights/gan_generator_epoch_{epoch+1}.pth')
+        torch.save(discriminator.state_dict(), f'weights/gan_discriminator_epoch_{epoch+1}.pth')
+
+if __name__ == "__main__":
+    print("Starting GAN Training...")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+    
+    # Setup Data
+    data_dir = 'data/iam_words'
+    csv_file = 'data/labels.csv'
+    
+    dataset = IAMDataset(data_dir=data_dir, csv_file=csv_file, transform=transform)
+    dataloader = DataLoader(dataset, batch_size=64, shuffle=True, collate_fn=collate_fn)
+    
+    # Setup Models
+    generator = Generator().to(device)
+    discriminator = Discriminator().to(device)
+    
+    # Resume from checkpoint if exists
+    start_epoch = 0
+    import glob
+    checkpoints = glob.glob('weights/gan_generator_epoch_*.pth')
+    if checkpoints:
+        checkpoints.sort(key=lambda x: int(os.path.basename(x).split('_')[-1].split('.')[0]))
+        latest_gen_checkpoint = checkpoints[-1]
+        start_epoch = int(os.path.basename(latest_gen_checkpoint).split('_')[-1].split('.')[0])
+        latest_disc_checkpoint = f'weights/gan_discriminator_epoch_{start_epoch}.pth'
+        
+        print(f"Resuming GAN training from epoch {start_epoch}")
+        generator.load_state_dict(torch.load(latest_gen_checkpoint, map_location=device))
+        discriminator.load_state_dict(torch.load(latest_disc_checkpoint, map_location=device))
+    
+    # Train
+    train_gan(generator, discriminator, dataloader, epochs=50, device=device, start_epoch=start_epoch)
+    print("GAN Training complete!")

src/training/train_ssl.py

@@ -0,0 +1,183 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from torchvision import transforms
+import os
+import sys
+import glob
+
+# Add project root to path
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
+
+from src.data.dataset import IAMDataset, collate_fn
+from src.models.crnn import CRNN
+from src.models.gan import Generator
+
+# Define transforms matching training exactly
+transform = transforms.Compose([
+    transforms.Resize((32, 1024)),
+    transforms.ToTensor(),
+    transforms.Normalize((0.5,), (0.5,))
+])
+
+def decode_pseudo_labels(preds):
+    # preds: (seq_len, batch, classes)
+    _, max_preds = torch.max(preds, 2)
+    max_preds = max_preds.permute(1, 0) # (batch, seq_len)
+    
+    targets_list = []
+    target_lengths = []
+    
+    for batch_idx in range(max_preds.size(0)):
+        pred_seq = max_preds[batch_idx]
+        decoded_seq = []
+        for i in range(len(pred_seq)):
+            if pred_seq[i] != 0 and (i == 0 or pred_seq[i] != pred_seq[i-1]):
+                decoded_seq.append(pred_seq[i].item())
+        
+        target_tensor = torch.tensor(decoded_seq, dtype=torch.long)
+        targets_list.append(target_tensor)
+        target_lengths.append(len(decoded_seq))
+        
+    return targets_list, target_lengths
+
+def train_ssl(model, generator, dataloader, optimizer, criterion, device, epochs=5, threshold=0.8, latent_dim=100):
+    """
+    Pseudo-labeling approach for Semi-Supervised Learning.
+    Combines real labeled data with synthetic unlabeled data generated dynamically by the GAN.
+    """
+    model.train()
+    generator.eval() # Generator is fixed during this phase
+    
+    for epoch in range(epochs):
+        total_loss_real = 0
+        total_loss_fake = 0
+        
+        for step, (labeled_imgs, labeled_texts, labeled_lengths) in enumerate(dataloader):
+            labeled_imgs = labeled_imgs.to(device)
+            labeled_texts = labeled_texts.to(device)
+            batch_size = labeled_imgs.size(0)
+            
+            optimizer.zero_grad()
+            
+            # ==============================================================
+            # 1. Train on Real Labeled Data
+            # ==============================================================
+            preds_l = model(labeled_imgs)
+            preds_l = preds_l.permute(1, 0, 2) # (seq_len, batch, classes)
+            
+            input_lengths_l = torch.full(size=(preds_l.size(1),), fill_value=preds_l.size(0), dtype=torch.long)
+            
+            targets_list_l = []
+            for i in range(labeled_texts.size(0)):
+                targets_list_l.append(labeled_texts[i][:labeled_lengths[i]])
+            targets_concat_l = torch.cat(targets_list_l)
+            
+            loss_real = criterion(preds_l, targets_concat_l, input_lengths_l, labeled_lengths)
+            
+            # ==============================================================
+            # 2. Train on Synthetic GAN Data (Pseudo-Labeling)
+            # ==============================================================
+            # Generate fake images
+            with torch.no_grad():
+                z = torch.randn(batch_size, latent_dim).to(device)
+                fake_imgs = generator(z) # Shape: (batch, 1, 32, 1024), range [-1, 1]
+                
+                # Get pseudo-labels
+                model.eval()
+                preds_fake_eval = model(fake_imgs)
+                probs = torch.exp(preds_fake_eval) # Softmax probs
+                max_probs, _ = torch.max(probs, dim=2)
+                avg_confidence = max_probs.mean(dim=1)
+                
+                # Mask confident predictions
+                mask = avg_confidence > threshold
+                
+            model.train()
+            loss_fake = torch.tensor(0.0).to(device)
+            
+            if mask.sum() > 0:
+                confident_imgs = fake_imgs[mask]
+                
+                preds_fake = model(confident_imgs)
+                preds_fake_perm = preds_fake.permute(1, 0, 2)
+                
+                # Decode the pseudo-labels into CTC targets
+                targets_list_u, target_lengths_u = decode_pseudo_labels(preds_fake_perm.detach())
+                
+                # Filter out empty pseudo-labels
+                valid_idx = [i for i, length in enumerate(target_lengths_u) if length > 0]
+                
+                if valid_idx:
+                    valid_preds_fake_perm = preds_fake_perm[:, valid_idx, :]
+                    valid_targets_list = [targets_list_u[i].to(device) for i in valid_idx]
+                    valid_target_lengths = torch.tensor([target_lengths_u[i] for i in valid_idx], dtype=torch.long).to(device)
+                    
+                    valid_targets_concat = torch.cat(valid_targets_list)
+                    input_lengths_u = torch.full(size=(valid_preds_fake_perm.size(1),), fill_value=valid_preds_fake_perm.size(0), dtype=torch.long).to(device)
+                    
+                    loss_fake = criterion(valid_preds_fake_perm, valid_targets_concat, input_lengths_u, valid_target_lengths)
+                    # Scale down the fake loss slightly so it doesn't overwhelm real data
+                    loss_fake = loss_fake * 0.5 
+
+            # Total loss
+            total_loss = loss_real + loss_fake
+            total_loss.backward()
+            optimizer.step()
+            
+            total_loss_real += loss_real.item()
+            total_loss_fake += loss_fake.item() if loss_fake > 0 else 0
+            
+            if step % 20 == 0:
+                print(f"Epoch [{epoch+1}/{epochs}], Step [{step}/{len(dataloader)}], Real Loss: {loss_real.item():.4f}, Fake Loss: {loss_fake.item() if loss_fake > 0 else 0:.4f}, Confident Fakes: {mask.sum().item()}/{batch_size}")
+                
+        print(f"Epoch {epoch+1} Average Real Loss: {total_loss_real/len(dataloader):.4f}, Average Fake Loss: {total_loss_fake/len(dataloader):.4f}")
+        
+        # Save checkpoints
+        os.makedirs('weights', exist_ok=True)
+        torch.save(model.state_dict(), f'weights/crnn_ssl_epoch_{epoch+1}.pth')
+
+if __name__ == "__main__":
+    print("Starting Semi-Supervised Learning (SSL) Training Phase...")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+    
+    # 1. Load Dataset
+    data_dir = 'data/iam_words'
+    csv_file = 'data/labels.csv'
+    dataset = IAMDataset(data_dir=data_dir, csv_file=csv_file, transform=transform)
+    dataloader = DataLoader(dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)
+    
+    # 2. Load the Baseline CRNN Model
+    num_classes = dataset.num_classes
+    crnn_model = CRNN(img_channel=1, img_height=32, img_width=1024, num_class=num_classes).to(device)
+    
+    checkpoints_crnn = glob.glob('weights/crnn_baseline_epoch_*.pth')
+    if not checkpoints_crnn:
+        print("Error: Could not find baseline CRNN weights.")
+        sys.exit(1)
+    checkpoints_crnn.sort(key=lambda x: int(os.path.basename(x).split('_')[-1].split('.')[0]))
+    latest_crnn = checkpoints_crnn[-1]
+    print(f"Loading Baseline CRNN from {latest_crnn}")
+    crnn_model.load_state_dict(torch.load(latest_crnn, map_location=device))
+    
+    # 3. Load the Trained GAN Generator
+    generator = Generator(latent_dim=100).to(device)
+    checkpoints_gan = glob.glob('weights/gan_generator_epoch_*.pth')
+    if not checkpoints_gan:
+        print("Error: Could not find GAN Generator weights.")
+        sys.exit(1)
+    checkpoints_gan.sort(key=lambda x: int(os.path.basename(x).split('_')[-1].split('.')[0]))
+    latest_gan = checkpoints_gan[-1]
+    print(f"Loading GAN Generator from {latest_gan}")
+    generator.load_state_dict(torch.load(latest_gan, map_location=device))
+    
+    # 4. Setup Optimizer & Loss
+    # Use a smaller learning rate for fine-tuning
+    optimizer = optim.Adam(crnn_model.parameters(), lr=0.0001)
+    criterion = nn.CTCLoss(blank=0, zero_infinity=True)
+    
+    # 5. Start SSL Training Loop
+    train_ssl(crnn_model, generator, dataloader, optimizer, criterion, device, epochs=5, threshold=0.8)
+    print("SSL Training complete!")

src/utils/preprocessing.py

@@ -0,0 +1,75 @@
+import cv2
+import numpy as np
+
+def preprocess_image(image_path_or_array, target_size=(1024, 32)):
+    """
+    Preprocess the image for handwritten text recognition.
+    1. Read image as grayscale
+    2. Resize while maintaining aspect ratio (padding with white)
+    3. Apply binarization / normalization
+    """
+    if isinstance(image_path_or_array, str):
+        img = cv2.imread(image_path_or_array, cv2.IMREAD_GRAYSCALE)
+        if img is None:
+            raise FileNotFoundError(f"Could not read image at {image_path_or_array}")
+    else:
+        if len(image_path_or_array.shape) == 3:
+            img = cv2.cvtColor(image_path_or_array, cv2.COLOR_BGR2GRAY)
+        else:
+            img = image_path_or_array.copy()
+
+    # Enhance contrast (CLAHE - Contrast Limited Adaptive Histogram Equalization)
+    # We do NOT want to do this if the image is already aggressively thresholded/binarized
+    # However, for smooth grayscale training images, CLAHE is great. 
+    # Let's keep it but recognize it might amplify noise if not careful.
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
+    img = clahe.apply(img)
+
+    # Resize keeping aspect ratio
+    h, w = img.shape
+    target_w, target_h = target_size
+    
+    # Calculate ratio
+    ratio_w = target_w / w
+    ratio_h = target_h / h
+    ratio = min(ratio_w, ratio_h)
+    
+    new_w = int(w * ratio)
+    new_h = int(h * ratio)
+    
+    # Check to prevent 0 width/height
+    if new_w == 0 or new_h == 0:
+        return np.ones((target_h, target_w), dtype=np.uint8) * 255
+    
+    img_resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
+    
+    # Create target blank (white) image
+    target_img = np.ones((target_h, target_w), dtype=np.uint8) * 255
+    
+    # Calculate padding to center it vertically, but align LEFT horizontally
+    # (Aligning left is usually better for sequence models like CTC)
+    pad_y = (target_h - new_h) // 2
+    pad_x = 0 # Align left instead of center
+    
+    # Paste resized image into target
+    target_img[pad_y:pad_y+new_h, pad_x:pad_x+new_w] = img_resized
+    
+    # Return as uint8 array without inverting, to match training behavior (white background)
+    return target_img
+
+def deskew(img):
+    """
+    Deskew the image using image moments.
+    """
+    m = cv2.moments(img)
+    if abs(m['mu02']) < 1e-2:
+        return img.copy()
+    
+    skew = m['mu11'] / m['mu02']
+    M = np.float32([[1, skew, -0.5 * img.shape[0] * skew], [0, 1, 0]])
+    img_deskewed = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
+    return img_deskewed
+
+if __name__ == "__main__":
+    # Simple test
+    print("Preprocessing module ready.")

src/web/app.py

@@ -0,0 +1,328 @@
+import gradio as gr
+import cv2
+import numpy as np
+import torch
+from torchvision import transforms
+from PIL import Image
+import pandas as pd
+import sys
+import os
+import matplotlib.pyplot as plt
+
+# Import preprocessing and model
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
+from src.utils.preprocessing import preprocess_image, deskew
+from src.models.crnn import CRNN
+
+# Define device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Build vocabulary directly from labels.csv without loading images
+try:
+    df = pd.read_csv('data/labels.csv')
+    chars = set()
+    for text in df['text']:
+        if pd.notna(text):
+            chars.update(list(str(text)))
+    vocab = sorted(list(chars))
+    idx_to_char = {i+1: c for i, c in enumerate(vocab)}
+    num_classes = len(vocab) + 1
+    print(f"Loaded vocabulary with {len(vocab)} characters")
+except Exception as e:
+    print(f"Could not load vocabulary from labels.csv: {e}")
+    # Fallback to standard IAM vocab if dataset not available
+    vocab = list("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789.,!? ")
+    idx_to_char = {i+1: c for i, c in enumerate(vocab)}
+    num_classes = len(vocab) + 1
+
+# Load Model
+model = CRNN(img_channel=1, img_height=32, img_width=1024, num_class=num_classes).to(device)
+
+import glob
+def get_latest_checkpoint(weights_dir='weights'):
+    checkpoints = glob.glob(os.path.join(weights_dir, 'crnn_baseline_epoch_*.pth'))
+    if not checkpoints:
+        return None
+    # Sort by epoch number
+    checkpoints.sort(key=lambda x: int(os.path.basename(x).split('_')[-1].split('.')[0]))
+    return checkpoints[-1]
+
+weights_path = get_latest_checkpoint()
+if weights_path and os.path.exists(weights_path):
+    print(f"Loading trained weights from {weights_path}...")
+    try:
+        model.load_state_dict(torch.load(weights_path, map_location=device))
+    except Exception as e:
+        print(f"Error loading weights perfectly (might be minor mismatch): {e}")
+        model.load_state_dict(torch.load(weights_path, map_location=device), strict=False)
+else:
+    print(f"Warning: Could not find any weights in weights/. Model will output random predictions.")
+
+model.eval()
+
+# Transform matching training exactly
+transform = transforms.Compose([
+    transforms.Resize((32, 1024)),
+    transforms.ToTensor(),
+    transforms.Normalize((0.5,), (0.5,))
+])
+
+def decode_predictions(preds, idx_to_char):
+    _, max_preds = torch.max(preds, 2)
+    max_preds = max_preds.permute(1, 0)
+    
+    decoded_texts = []
+    for batch_idx in range(max_preds.size(0)):
+        pred_seq = max_preds[batch_idx]
+        decoded_seq = []
+        for i in range(len(pred_seq)):
+            if pred_seq[i] != 0 and (i == 0 or pred_seq[i] != pred_seq[i-1]):
+                char_idx = pred_seq[i].item()
+                if char_idx in idx_to_char:
+                    decoded_seq.append(idx_to_char[char_idx])
+        decoded_texts.append("".join(decoded_seq))
+    return decoded_texts
+
+def auto_crop_image(gray_img):
+    # Apply Gaussian blur to reduce noise
+    blurred = cv2.GaussianBlur(gray_img, (5, 5), 0)
+    
+    # Apply Otsu's thresholding to separate dark ink from white background
+    _, thresh = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
+    
+    # Find contours (shapes) in the image
+    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    
+    if not contours:
+        return gray_img
+        
+    # Filter contours to exclude tiny noise and giant objects (like the pen)
+    img_area = gray_img.shape[0] * gray_img.shape[1]
+    valid_contours = []
+    for c in contours:
+        area = cv2.contourArea(c)
+        # Keep contours that are larger than a speck of dust but smaller than half the image
+        if 20 < area < (img_area * 0.4): 
+            valid_contours.append(c)
+            
+    if not valid_contours:
+        return gray_img # Fallback to original if filtering removes everything
+        
+    # Find the bounding box that encompasses all valid text contours
+    x_min, y_min = float('inf'), float('inf')
+    x_max, y_max = 0, 0
+    
+    for c in valid_contours:
+        x, y, w, h = cv2.boundingRect(c)
+        x_min = min(x_min, x)
+        y_min = min(y_min, y)
+        x_max = max(x_max, x + w)
+        y_max = max(y_max, y + h)
+        
+    # Add a generous padding around the text
+    pad_y = int((y_max - y_min) * 0.2)
+    pad_x = int((x_max - x_min) * 0.05)
+    
+    x_min = max(0, x_min - pad_x)
+    y_min = max(0, y_min - pad_y)
+    x_max = min(gray_img.shape[1], x_max + pad_x)
+    y_max = min(gray_img.shape[0], y_max + pad_y)
+    
+    # Crop the image
+    cropped = gray_img[y_min:y_max, x_min:x_max]
+    
+    # CRITICAL FIX for Out-of-Distribution aspect ratios:
+    # The training data (IAM dataset) has an average aspect ratio of ~16:1.
+    # The training pipeline blindly squashes images to 32x1024 (32:1 ratio).
+    # If a user uploads a short word (like a 3:1 ratio "THANK YOU"), 
+    # it gets stretched 10x horizontally, destroying the letters!
+    # To fix this, we pad the cropped image with white space on the right 
+    # so its aspect ratio matches the training average (16:1) BEFORE squashing.
+    
+    h, w = cropped.shape
+    target_aspect_ratio = 16.0
+    if w / h < target_aspect_ratio:
+        target_w = int(h * target_aspect_ratio)
+        pad_width = target_w - w
+        # Pad with white (255) on the right
+        cropped = cv2.copyMakeBorder(cropped, 0, 0, 0, pad_width, cv2.BORDER_CONSTANT, value=255)
+    
+    return cropped
+
+def process_and_predict(image, apply_auto_crop=True):
+    if image is None:
+        return None, "Please upload an image.", None, None, None
+
+    # Convert Gradio Image (which is a PIL Image by default) to grayscale
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+        
+    gray_image = image.convert('L')
+    
+    # For display purposes (Gradio output image)
+    img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    gray_cv = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
+    
+    # CRITICAL: Binarization (Otsu's thresholding) to force pure black text on pure white background
+    # This removes shadows, lighting gradients, and colored paper backgrounds
+    # that the model was never trained on.
+    blurred = cv2.GaussianBlur(gray_cv, (5, 5), 0)
+    _, binarized = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+    
+    if not apply_auto_crop:
+        # If auto-crop is disabled, we bypass all fancy preprocessing to precisely 
+        # match the dataset loading behavior. This ensures dataset images work perfectly.
+        gray_image_pil = Image.fromarray(gray_cv)
+        img_tensor = transform(gray_image_pil).unsqueeze(0).to(device)
+        # For display, just show what the network sees (squashed)
+        display_processed_img = np.array(gray_image_pil.resize((1024, 32), Image.BILINEAR))
+    else:
+        # Auto-crop if requested (using the binarized image for cleaner crops)
+        processed_base = auto_crop_image(binarized)
+            
+        deskewed_img = deskew(processed_base)
+        processed_img_np = preprocess_image(deskewed_img, target_size=(1024, 32))
+        display_processed_img = processed_img_np
+        
+        # Convert cropped numpy array back to PIL for tensor transform
+        gray_image_cropped = Image.fromarray(display_processed_img)
+        
+        # For Model Prediction
+        # We must use exactly the same transform as training, and pass a PIL image
+        img_tensor = transform(gray_image_cropped).unsqueeze(0).to(device)
+    
+    # Predict and extract features
+    with torch.no_grad():
+        # Get CNN features for activation map
+        cnn_features = model.cnn(img_tensor) # shape: (1, 512, 1, seq_len)
+        
+        preds = model(img_tensor)
+        preds = preds.permute(1, 0, 2) # (seq_len, batch, num_classes)
+        decoded_text = decode_predictions(preds, idx_to_char)[0]
+        
+        # Calculate probabilities from LogSoftmax output
+        probs = torch.exp(preds[:, 0, :]) # shape: (seq_len, num_classes)
+    
+    if not decoded_text.strip():
+        decoded_text = "[Model returned blank - Needs more training epochs]"
+
+    # 1. Generate CTC Probability Matrix Heatmap
+    probs_np = probs.cpu().numpy().T # shape: (num_classes, seq_len)
+    fig_heatmap, ax1 = plt.subplots(figsize=(10, 4))
+    cax = ax1.imshow(probs_np, aspect='auto', cmap='viridis')
+    ax1.set_title("CTC Probability Matrix Heatmap")
+    ax1.set_xlabel("Time Frame (Sequence Steps)")
+    ax1.set_ylabel("Vocabulary Character Index")
+    fig_heatmap.colorbar(cax, ax=ax1, fraction=0.046, pad=0.04, label="Probability")
+    plt.tight_layout()
+
+    # 2. Generate Character Confidence Bar Chart
+    max_probs, max_idx = torch.max(probs, dim=1)
+    chars = []
+    confidences = []
+    
+    for i in range(len(max_idx)):
+        if max_idx[i] != 0 and (i == 0 or max_idx[i] != max_idx[i-1]):
+            char_idx = max_idx[i].item()
+            if char_idx in idx_to_char:
+                chars.append(idx_to_char[char_idx])
+                confidences.append(max_probs[i].item())
+
+    # Adjust width based on number of characters
+    fig_bar, ax2 = plt.subplots(figsize=(max(8, len(chars)*0.4), 4))
+    if chars:
+        bars = ax2.bar(range(len(chars)), confidences, color='#FF9900')
+        ax2.set_xticks(range(len(chars)))
+        ax2.set_xticklabels(chars)
+        ax2.set_ylim(0, 1.1)
+        ax2.set_title("Character Confidence Scores")
+        ax2.set_ylabel("Confidence Probability")
+        
+        # Add percentage labels above bars
+        for bar in bars:
+            yval = bar.get_height()
+            ax2.text(bar.get_x() + bar.get_width()/2.0, yval + 0.02, 
+                     f'{yval*100:.0f}%', va='bottom', ha='center', fontsize=8, rotation=45)
+    else:
+        ax2.text(0.5, 0.5, "No characters predicted", ha='center', va='center')
+        
+    plt.tight_layout()
+    
+    # 3. Generate CNN Feature Activation Overlay
+    # Average the CNN features across all channels to get a 1D activation map
+    activation = torch.mean(cnn_features, dim=1).squeeze().cpu().numpy()
+    
+    # Normalize activation to 0-255
+    activation = (activation - activation.min()) / (activation.max() - activation.min() + 1e-8)
+    activation = (activation * 255).astype(np.uint8)
+    
+    # Resize to match the original image dimensions
+    heatmap_img = cv2.resize(activation, (processed_img_np.shape[1], processed_img_np.shape[0]))
+    
+    # Apply color map
+    heatmap_color = cv2.applyColorMap(heatmap_img, cv2.COLORMAP_JET)
+    
+    # Convert grayscale original image to BGR so we can blend it
+    original_bgr = cv2.cvtColor(display_processed_img, cv2.COLOR_GRAY2BGR)
+    
+    # Overlay heatmap on original image (50% alpha blend)
+    overlay_img = cv2.addWeighted(heatmap_color, 0.5, original_bgr, 0.5, 0)
+    # Convert BGR to RGB for Gradio display
+    overlay_img = cv2.cvtColor(overlay_img, cv2.COLOR_BGR2RGB)
+        
+    return display_processed_img, decoded_text, fig_heatmap, fig_bar, overlay_img
+
+# Redesign UI with Gradio Blocks for a proper Dashboard layout
+with gr.Blocks(title="Handwritten Text Recognition (HTR)", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("<h1 style='text-align: center;'>Handwritten Text Recognition (HTR) Dashboard</h1>")
+    gr.Markdown("Upload an image of handwritten text. The system will preprocess it and extract the text using our trained custom CRNN model.")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            # Editor tool allows manual cropping in UI before sending
+            input_image = gr.Image(type="pil", label="Upload Handwritten Text Image")
+            auto_crop_checkbox = gr.Checkbox(label="✨ Auto-Crop Background (Smart Vision)", value=True, info="Automatically zooms in on the text and removes giant background objects/pens.")
+            with gr.Row():
+                clear_btn = gr.Button("Clear")
+                submit_btn = gr.Button("Submit", variant="primary")
+            
+        with gr.Column(scale=1):
+            output_image = gr.Image(type="numpy", label="Preprocessed (1024 x 32)")
+            gr.Markdown("<p style='font-size: 12px; color: gray;'>Grayscale, aspect-ratio preserved, padded to 32x1024</p>")
+            output_text = gr.Textbox(label="Predicted Text", lines=2)
+            
+    gr.Markdown("---")
+    gr.Markdown("### 📊 Model Insights & Analytics (Explainable AI)")
+    
+    with gr.Accordion("📖 How to read these graphs (Interpretation Guide)", open=False):
+        gr.Markdown("""
+        **1. CNN Feature Activation Overlay:** Shows exactly where the model's 'eyes' are focusing on the image. Red/hot areas indicate regions with strong visual features (like complex curves or sharp lines) that the Convolutional Neural Network detected.
+        
+        **2. CTC Probability Matrix Heatmap:** Shows *when* the model made a decision. The X-axis is the timeline (reading left-to-right), and the Y-axis contains all possible characters. Yellow dots indicate the exact moment the AI identified a specific letter.
+        
+        **3. Character Confidence Scores:** Shows *how sure* the model is about each letter it predicted. If the model misreads a word, this chart usually shows a low confidence score for the incorrect letter, proving it was uncertain.
+        """)
+        
+    with gr.Row():
+        cnn_activation_image = gr.Image(type="numpy", label="1. CNN Feature Activation Overlay")
+        
+    with gr.Row():
+        heatmap_plot = gr.Plot(label="2. CTC Probability Heatmap")
+        
+    with gr.Row():
+        confidence_plot = gr.Plot(label="3. Character Confidence Scores")
+        
+    submit_btn.click(
+        fn=process_and_predict,
+        inputs=[input_image, auto_crop_checkbox],
+        outputs=[output_image, output_text, heatmap_plot, confidence_plot, cnn_activation_image]
+    )
+    
+    clear_btn.click(
+        fn=lambda: [None, True, None, "", None, None, None],
+        inputs=[],
+        outputs=[input_image, auto_crop_checkbox, output_image, output_text, heatmap_plot, confidence_plot, cnn_activation_image]
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=True)

weights/crnn_baseline_epoch_30.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f203e852eb08710b520beed65b3bbf0edb5c8fb66ac34e61936bb9660ed2dec7
+size 31473673