ablation_study_experiments.py

# -*- coding: utf-8 -*-
"""Ablation Study Experiments.ipynb

Automatically generated by Colab.

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

# ===================================================================
# Cell 1: Setup & Mount Drive
# ===================================================================
from google.colab import drive
drive.mount('/content/drive')
from google.colab import files
import os

print("📤 Upload the 3 pre-split JSON files from train_test_0 folder:")
print("   - ocr_train.json")
print("   - ocr_val.json")
print("   - ocr_test.json")

uploaded = files.upload()

# Verify uploads
print("\n✅ Upload verification:")
for filename in ['ocr_train.json', 'ocr_val.json', 'ocr_test.json']:
    if filename in uploaded:
        print(f"   ✅ {filename} ({len(uploaded[filename])} bytes)")
    else:
        raise FileNotFoundError(f"   ❌ Missing: {filename} - Please upload it!")

print("\n✅ All files uploaded successfully!")

# ===================================================================
# Cell 2: Install Dependencies & Main Imports
# ===================================================================
!pip install -q transformers datasets accelerate pillow editdistance jiwer opencv-python matplotlib seaborn scipy pandas

import os
import json
import random
from pathlib import Path
from datetime import datetime
import numpy as np
import pandas as pd
import torch
import cv2
from PIL import Image
import matplotlib.pyplot as plt
import seaborn as sns
import time
from scipy.stats import pearsonr, spearmanr, linregress

# Configuration
DRIVE_BASE = "/content/drive/MyDrive"
IMG_DIR = Path(f"{DRIVE_BASE}/OMR/images/Dataset cortona")
TRAIN_JSON = Path("/content/ocr_train.json")
VAL_JSON = Path("/content/ocr_val.json")
TEST_JSON = Path("/content/ocr_test.json")
PROCESSED_DIR = Path("/content/processed_data")
OUTPUT_BASE = Path("/content/trocr_ablation_results")

PROCESSED_DIR.mkdir(exist_ok=True)
OUTPUT_BASE.mkdir(exist_ok=True)

# Experiment selection
EXPERIMENT = "freeze_none"  # Options: baseline, no_clahe, no_aug, freeze_none

# Training hyperparameters
NUM_EPOCHS = 50
TRAIN_BATCH = 4
EVAL_BATCH = 8
GRAD_ACCUM = 4
LR = 3e-5

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# ===================================================================
# Cell 3: Load Pre-Split Data from train_test_2 folder
# ===================================================================

print("📂 Loading pre-split data from train_test_2 folder...")

# Check files exist
for json_path in [TRAIN_JSON, VAL_JSON, TEST_JSON]:
    if not json_path.exists():
        raise FileNotFoundError(f"❌ Missing file: {json_path}")

# Load the splits
with open(TRAIN_JSON, "r", encoding="utf-8") as f:
    train_data = json.load(f)

with open(VAL_JSON, "r", encoding="utf-8") as f:
    val_data = json.load(f)

with open(TEST_JSON, "r", encoding="utf-8") as f:
    test_data = json.load(f)

# Extract annotations
train_anns = train_data["annotations"]
val_anns = val_data["annotations"]
test_anns = test_data["annotations"]

# Create unified image map
all_images = train_data.get("images", []) + val_data.get("images", []) + test_data.get("images", [])
image_map = {img["id"]: img for img in all_images}

# Remove duplicates in image_map (in case images appear in multiple files)
unique_images = {}
for img in all_images:
    if img["id"] not in unique_images:
        unique_images[img["id"]] = img
image_map = unique_images

print(f"✅ Loaded pre-split data:")
print(f"   Train: {len(train_anns)} annotations")
print(f"   Val:   {len(val_anns)} annotations")
print(f"   Test:  {len(test_anns)} annotations")
print(f"   Total unique images: {len(image_map)}")

# Filter for text-line category (category_id = 6)
TEXT_LINE_CATEGORY = 6

train_anns = [ann for ann in train_anns
              if ann.get("category_id") == TEXT_LINE_CATEGORY and
              ann.get("description", "").strip()]
val_anns = [ann for ann in val_anns
            if ann.get("category_id") == TEXT_LINE_CATEGORY and
            ann.get("description", "").strip()]
test_anns = [ann for ann in test_anns
             if ann.get("category_id") == TEXT_LINE_CATEGORY and
             ann.get("description", "").strip()]

print(f"\n✅ After filtering for text-line category:")
print(f"   Train: {len(train_anns)} text lines")
print(f"   Val:   {len(val_anns)} text lines")
print(f"   Test:  {len(test_anns)} text lines")

# Verification: Check for data leakage
train_ids = set([ann["id"] for ann in train_anns])
val_ids = set([ann["id"] for ann in val_anns])
test_ids = set([ann["id"] for ann in test_anns])

overlaps = {
    "train-val": len(train_ids & val_ids),
    "train-test": len(train_ids & test_ids),
    "val-test": len(val_ids & test_ids)
}

print(f"\n✅ Data leakage check:")
print(f"   Train-Val overlap: {overlaps['train-val']} (should be 0)")
print(f"   Train-Test overlap: {overlaps['train-test']} (should be 0)")
print(f"   Val-Test overlap: {overlaps['val-test']} (should be 0)")

if sum(overlaps.values()) == 0:
    print("   ✅ PASSED: No data leakage detected")
else:
    print("   ⚠️ WARNING: Data leakage detected!")

# Store filtered data
filtered = {
    "info": train_data.get("info", {}),
    "licenses": train_data.get("licenses", []),
    "categories": train_data.get("categories", []),
    "images": list(image_map.values()),
    "annotations": train_anns + val_anns + test_anns
}

# ===================================================================
# Cell 4: Data Augmentation Class
# ===================================================================
class ManuscriptAugmentation:
    def __call__(self, image: Image.Image) -> Image.Image:
        img_np = np.array(image)

        # Rotation
        if random.random() < 0.4:
            angle = random.uniform(-5, 5)
            h, w = img_np.shape[:2]
            M = cv2.getRotationMatrix2D((w/2, h/2), angle, 1.0)
            img_np = cv2.warpAffine(img_np, M, (w, h), borderMode=cv2.BORDER_REPLICATE)

        # Elastic deformation
        if random.random() < 0.25:
            h, w = img_np.shape[:2]
            if h > 20 and w > 20:
                dx = cv2.resize(np.random.randn(max(h//10, 3), max(w//10, 3)), (w, h)) * 3
                dy = cv2.resize(np.random.randn(max(h//10, 3), max(w//10, 3)), (w, h)) * 3
                x, y = np.meshgrid(np.arange(w), np.arange(h))
                map_x = np.clip(x + dx, 0, w-1).astype(np.float32)
                map_y = np.clip(y + dy, 0, h-1).astype(np.float32)
                img_np = cv2.remap(img_np, map_x, map_y, cv2.INTER_LINEAR)

        # Gaussian blur
        if random.random() < 0.35:
            kernel_size = random.choice([3,5,7])
            sigma = random.uniform(0.5,1.2)
            img_np = cv2.GaussianBlur(img_np, (kernel_size,kernel_size), sigma)

        # Brightness
        if random.random() < 0.4:
            factor = random.uniform(0.7,1.3)
            img_np = np.clip(img_np * factor, 0, 255).astype(np.uint8)

        # Contrast
        if random.random() < 0.3:
            factor = random.uniform(0.8,1.2)
            mean = img_np.mean()
            img_np = np.clip((img_np-mean)*factor + mean, 0, 255).astype(np.uint8)

        # Speckle noise
        if random.random() < 0.25:
            noise = np.random.randn(*img_np.shape) * 6
            img_np = np.clip(img_np + noise, 0, 255).astype(np.uint8)

        # Morphological operations
        if random.random() < 0.15:
            kernel = np.ones((2,2), np.uint8)
            if random.random() < 0.5:
                img_np = cv2.erode(img_np, kernel, iterations=1)
            else:
                img_np = cv2.dilate(img_np, kernel, iterations=1)

        # Shadow/staining
        if random.random() < 0.15:
            h, w = img_np.shape[:2]
            if h > 10 and w > 10:
                shadow = np.random.randn(max(h//5,2), max(w//5,2))
                shadow = cv2.resize(shadow, (w,h))
                shadow = (shadow - shadow.min()) / (shadow.max()-shadow.min()+1e-8)
                shadow = 1 - shadow * 0.25
                if len(img_np.shape) == 3:
                    shadow = np.expand_dims(shadow, axis=-1)
                img_np = np.clip(img_np * shadow, 0, 255).astype(np.uint8)

        return Image.fromarray(img_np)

# ===================================================================
# Cell 5: Dataset Class
# ===================================================================
from torch.utils.data import Dataset
from transformers import TrOCRProcessor

class TrOCRDataset(Dataset):
    def __init__(self, annotations, images_info, image_root, processor: TrOCRProcessor,
                 max_target_length=128, augment_transform=None, use_clahe=True):
        self.annotations = annotations
        self.image_root = Path(image_root)
        self.processor = processor
        self.max_target_length = max_target_length
        self.augment_transform = augment_transform
        self.use_clahe = use_clahe
        self.image_id_to_info = images_info
        print(f"Dataset: {len(self.annotations)} samples, CLAHE={self.use_clahe}, AUG={augment_transform is not None}")

    def preprocess_image(self, img_np):
        if len(img_np.shape) == 3:
            gray = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY)
        else:
            gray = img_np

        if self.use_clahe:
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
            enhanced = clahe.apply(gray)
        else:
            enhanced = cv2.equalizeHist(gray)

        rgb = cv2.cvtColor(enhanced, cv2.COLOR_GRAY2RGB)
        return rgb

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

    def __getitem__(self, idx):
        ann = self.annotations[idx]
        text = ann.get("description", "") or ann.get("text", "") or " "
        img_info = self.image_id_to_info[ann["image_id"]]
        image_path = self.image_root / img_info["file_name"]

        img = cv2.imread(str(image_path))
        if img is None:
            raise FileNotFoundError(f"Image not found: {image_path}")

        x,y,w,h = [int(v) for v in ann["bbox"]]
        margin = 8
        x = max(0, x-margin)
        y = max(0, y-margin)
        w = min(img.shape[1] - x, w + 2*margin)
        h = min(img.shape[0] - y, h + 2*margin)

        crop = img[y:y+h, x:x+w]
        crop = self.preprocess_image(crop)
        pil = Image.fromarray(crop)

        if self.augment_transform:
            pil = self.augment_transform(pil)

        encoding = self.processor(
            images=pil,
            text=text,
            padding="max_length",
            max_length=self.max_target_length,
            truncation=True,
            return_tensors="pt"
        )

        encoding = {k: v.squeeze(0) for k, v in encoding.items()}
        return encoding

# ===================================================================
# Cell 6: Model Setup
# ===================================================================
from transformers import (
    TrOCRProcessor,
    VisionEncoderDecoderModel,
    Seq2SeqTrainingArguments,
    Seq2SeqTrainer,
    EarlyStoppingCallback
)

MODEL_NAME = "microsoft/trocr-base-handwritten"
processor = TrOCRProcessor.from_pretrained(MODEL_NAME)
model = VisionEncoderDecoderModel.from_pretrained(MODEL_NAME)

model.config.decoder_start_token_id = processor.tokenizer.cls_token_id
model.config.eos_token_id = processor.tokenizer.eos_token_id
model.config.pad_token_id = processor.tokenizer.pad_token_id
model.config.vocab_size = len(processor.tokenizer)

print("Model loaded")

# ===================================================================
# Cell 7: Metrics
# ===================================================================
import editdistance
from jiwer import wer as jiwer_wer

def calculate_wer_cer(pred_strs, label_strs):
    if not pred_strs:
        return 1.0, 1.0

    total_cer = 0.0
    for p, l in zip(pred_strs, label_strs):
        total_cer += editdistance.eval(p, l) / max(len(l), 1)
    cer = total_cer / len(pred_strs)

    try:
        wer = jiwer_wer(label_strs, pred_strs)
    except:
        wer = 1.0

    return wer, cer

def compute_metrics(pred):
    labels_ids = pred.label_ids
    pred_ids = pred.predictions

    pred_ids[pred_ids < 0] = processor.tokenizer.pad_token_id
    labels_ids[labels_ids == -100] = processor.tokenizer.pad_token_id

    pred_strs = processor.batch_decode(pred_ids, skip_special_tokens=True)
    label_strs = processor.batch_decode(labels_ids, skip_special_tokens=True)

    wer, cer = calculate_wer_cer(pred_strs, label_strs)
    return {"wer": wer, "cer": cer}

# ===================================================================
# Cell 8: Data Collator
# ===================================================================
class FixedDataCollator:
    def __init__(self, processor):
        self.processor = processor

    def __call__(self, features):
        pixel_values = torch.stack([f["pixel_values"] for f in features]).float()
        labels = torch.stack([f["labels"] for f in features]).long()

        decoder_input_ids = labels.new_zeros(labels.shape)
        decoder_input_ids[:, 1:] = labels[:, :-1].clone()
        decoder_input_ids[:, 0] = self.processor.tokenizer.cls_token_id

        labels[labels == self.processor.tokenizer.pad_token_id] = -100

        return {
            "pixel_values": pixel_values,
            "decoder_input_ids": decoder_input_ids,
            "labels": labels
        }

# ===================================================================
# Cell 9: Experiment Configurations
# ===================================================================
def configure_experiment(exp_name, model):
    configs = {
        "baseline": {
            "use_clahe": True,
            "use_aug": True,
            "freeze_layers": 10,
            "description": "Full pipeline with CLAHE, augmentation, and 10 layers frozen"
        },
        "no_clahe": {
            "use_clahe": False,
            "use_aug": True,
            "freeze_layers": 10,
            "description": "Without CLAHE preprocessing"
        },
        "no_aug": {
            "use_clahe": True,
            "use_aug": False,
            "freeze_layers": 10,
            "description": "Without data augmentation"
        },
        "freeze_none": {
            "use_clahe": True,
            "use_aug": True,
            "freeze_layers": 0,
            "description": "Full fine-tuning (no frozen layers)"
        },
        "freeze_all": {
            "use_clahe": True,
            "use_aug": True,
            "freeze_layers": 12,
            "description": "Freeze entire encoder"
        }
    }

    if exp_name not in configs:
        raise ValueError(f"Unknown experiment: {exp_name}. Choose from {list(configs.keys())}")

    config = configs[exp_name]

    num_layers = len(model.encoder.encoder.layer)
    freeze_layers = config["freeze_layers"]

    for i, layer in enumerate(model.encoder.encoder.layer):
        for param in layer.parameters():
            param.requires_grad = (i >= freeze_layers)

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())

    print(f"\n{'='*70}")
    print(f"Experiment: {exp_name}")
    print(f"Description: {config['description']}")
    print(f"Frozen layers: {freeze_layers}/{num_layers}")
    print(f"Trainable parameters: {trainable:,} / {total:,} ({100*trainable/total:.1f}%)")
    print(f"{'='*70}\n")

    return config

# ===================================================================
# Cell 10: Main Training Function
# ===================================================================
def run_experiment(exp_name):
    config = configure_experiment(exp_name, model)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    out_dir = OUTPUT_BASE / f"{exp_name}_{timestamp}"
    out_dir.mkdir(parents=True, exist_ok=True)

    with open(out_dir / "config.json", "w") as f:
        json.dump(config, f, indent=2)

    aug = ManuscriptAugmentation() if config["use_aug"] else None

    train_dataset = TrOCRDataset(
        train_anns, image_map, IMG_DIR, processor,
        max_target_length=128,
        augment_transform=aug,
        use_clahe=config["use_clahe"]
    )

    val_dataset = TrOCRDataset(
        val_anns, image_map, IMG_DIR, processor,
        max_target_length=128,
        augment_transform=None,
        use_clahe=config["use_clahe"]
    )

    test_dataset = TrOCRDataset(
        test_anns, image_map, IMG_DIR, processor,
        max_target_length=128,
        augment_transform=None,
        use_clahe=config["use_clahe"]
    )

    training_args = Seq2SeqTrainingArguments(
        output_dir=str(out_dir / "checkpoints"),
        predict_with_generate=True,
        eval_strategy="epoch",
        save_strategy="epoch",
        per_device_train_batch_size=TRAIN_BATCH,
        per_device_eval_batch_size=EVAL_BATCH,
        gradient_accumulation_steps=GRAD_ACCUM,
        num_train_epochs=NUM_EPOCHS,
        learning_rate=LR,
        warmup_ratio=0.1,
        weight_decay=0.01,
        label_smoothing_factor=0.1,
        fp16=torch.cuda.is_available(),
        load_best_model_at_end=True,
        metric_for_best_model="cer",
        greater_is_better=False,
        logging_strategy="steps",
        logging_steps=50,
        save_total_limit=3,
        dataloader_num_workers=2,
        report_to="none"
    )

    trainer = Seq2SeqTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        data_collator=FixedDataCollator(processor),
        compute_metrics=compute_metrics,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=5)]
    )

    print("\n🚀 Starting training...")
    train_result = trainer.train()

    print("\n📊 Evaluating on test set...")
    test_result = trainer.predict(test_dataset)

    results = {
        "experiment": exp_name,
        "config": config,
        "train_metrics": train_result.metrics,
        "test_metrics": test_result.metrics
    }

    with open(out_dir / "results.json", "w") as f:
        json.dump(results, f, indent=2)

    print(f"\n✅ Test CER: {test_result.metrics['test_cer']:.4f}")
    print(f"✅ Test WER: {test_result.metrics['test_wer']:.4f}")

    model_dir = out_dir / "final_model"
    trainer.save_model(model_dir)
    processor.save_pretrained(model_dir)
    print(f"\n💾 Model saved to {model_dir}")

    print("\n🎨 Generating attention visualizations...")
    vis_dir = out_dir / "attention_maps"
    vis_dir.mkdir(parents=True, exist_ok=True)

    n_vis = min(10, len(test_dataset))
    for i in range(n_vis):
        try:
            item = test_dataset[i]
            save_path = vis_dir / f"attention_{i:03d}.png"
            visualize_decoder_cross_attention_fixed(model, item, processor, save_path, device)
        except Exception as e:
            print(f"❌ Failed to visualize sample {i}: {e}")

    print(f"\n✅ Experiment complete! Results saved to {out_dir}")
    return out_dir, results

# ===================================================================
# Cell 11: RUN MAIN EXPERIMENT
# ===================================================================
print(f"\n{'#'*70}")
print(f"# RUNNING EXPERIMENT: {EXPERIMENT}")
print(f"{'#'*70}\n")

output_dir, results = run_experiment(EXPERIMENT)

print(f"\n\n{'#'*70}")
print(f"# EXPERIMENT COMPLETED")
print(f"# Output directory: {output_dir}")
print(f"{'#'*70}\n")