add src

.gitignore

@@ -0,0 +1 @@
+.idea

app.py

@@ -0,0 +1,279 @@
+#!/usr/bin/env python3
+"""
+Simple Gradio app for testing an EyeQ QC model.
+
+Example
+-------
+python app_eyeq.py \
+    --checkpoint ./checkpoints/eyeq_vit_base/best.pt
+
+Then open the printed local URL in your browser.
+"""
+
+import argparse
+from pathlib import Path
+
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image
+from torchvision import transforms
+import timm
+
+
+ID_TO_LABEL = {0: "Good", 1: "Usable", 2: "Reject"}
+
+
+def build_transform(img_size: int):
+    return transforms.Compose([
+        transforms.Resize((img_size, img_size)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+    ])
+
+
+def load_model(checkpoint_path: str, device: torch.device):
+    ckpt = torch.load(checkpoint_path, map_location="cpu")
+
+    args = ckpt.get("args", {})
+    model_name = args.get("model", "vit_base_patch16_224")
+    img_size = int(args.get("img_size", 224))
+
+    id_to_label = ckpt.get("id_to_label", ID_TO_LABEL)
+    id_to_label = {int(k): v for k, v in id_to_label.items()}
+
+    model = timm.create_model(
+        model_name,
+        pretrained=False,
+        num_classes=len(id_to_label),
+    )
+    model.load_state_dict(ckpt["model"], strict=True)
+    model.to(device)
+    model.eval()
+
+    tfm = build_transform(img_size)
+    return model, tfm, id_to_label, model_name, img_size
+
+
+def get_eyeq_class_ids(id_to_label):
+    """Return class IDs for Good, Usable, Reject.
+
+    Falls back to the standard EyeQ ordering if the checkpoint does not store
+    string labels in the expected form.
+    """
+    label_to_id = {str(v).lower(): int(k) for k, v in id_to_label.items()}
+
+    good_id = label_to_id.get("good", 0)
+    usable_id = label_to_id.get("usable", 1)
+    reject_id = label_to_id.get("reject", 2)
+
+    return good_id, usable_id, reject_id
+
+
+def soft_eyeq_decision(probs, id_to_label, reject_threshold=0.60, reject_margin=0.15):
+    """Apply a conservative Reject rule.
+
+    Reject is only returned when:
+      1. P(Reject) >= reject_threshold, and
+      2. P(Reject) beats the best non-Reject class by reject_margin.
+
+    Otherwise, the prediction is forced to Good vs Usable.
+    """
+    good_id, usable_id, reject_id = get_eyeq_class_ids(id_to_label)
+
+    prob_good = float(probs[good_id])
+    prob_usable = float(probs[usable_id])
+    prob_reject = float(probs[reject_id])
+
+    best_non_reject_id = good_id if prob_good >= prob_usable else usable_id
+    best_non_reject_prob = max(prob_good, prob_usable)
+
+    if (
+        prob_reject >= reject_threshold
+        and (prob_reject - best_non_reject_prob) >= reject_margin
+    ):
+        pred_id = reject_id
+        decision = "Soft rule: Reject threshold and margin were both satisfied."
+    else:
+        pred_id = best_non_reject_id
+        decision = "Soft rule: Reject was not confident enough, so prediction was forced to Good/Usable."
+
+    return pred_id, id_to_label[pred_id], decision
+
+
+def update_margin_slider(reject_threshold, reject_margin):
+    """Keep reject_margin within a sensible range for the current threshold."""
+    max_margin = min(0.50, float(reject_threshold))
+    reject_margin = min(float(reject_margin), max_margin)
+
+    return gr.update(
+        maximum=max_margin,
+        value=reject_margin,
+    )
+
+
+@torch.no_grad()
+def predict_quality(
+    image: Image.Image,
+    model,
+    tfm,
+    id_to_label,
+    device,
+    reject_threshold=0.60,
+    reject_margin=0.15,
+):
+    if image is None:
+        return None, {}, "Upload an image to run QC."
+
+    image = image.convert("RGB")
+    x = tfm(image).unsqueeze(0).to(device)
+
+    logits = model(x)
+    probs = torch.softmax(logits, dim=1)[0].detach().cpu().numpy()
+
+    raw_pred_id = int(np.argmax(probs))
+    raw_pred_label = id_to_label[raw_pred_id]
+
+    soft_pred_id, soft_pred_label, decision = soft_eyeq_decision(
+        probs=probs,
+        id_to_label=id_to_label,
+        reject_threshold=reject_threshold,
+        reject_margin=reject_margin,
+    )
+
+    prob_dict = {
+        id_to_label[i]: float(probs[i])
+        for i in range(len(probs))
+    }
+
+    detail = (
+        f"Raw argmax: {raw_pred_label}\n"
+        f"Soft decision: {soft_pred_label}\n"
+        f"Reject threshold: {reject_threshold:.2f} | Reject margin: {reject_margin:.2f}\n"
+        f"{decision}"
+    )
+
+    return soft_pred_label, prob_dict, detail
+
+
+def make_app(checkpoint_path: str):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model, tfm, id_to_label, model_name, img_size = load_model(checkpoint_path, device)
+
+    def run(image, reject_threshold, reject_margin):
+        pred_label, prob_dict, detail = predict_quality(
+            image=image,
+            model=model,
+            tfm=tfm,
+            id_to_label=id_to_label,
+            device=device,
+            reject_threshold=reject_threshold,
+            reject_margin=reject_margin,
+        )
+        return pred_label, prob_dict, detail
+
+    with gr.Blocks(title="EyeQ CFP Quality Control") as demo:
+        gr.Markdown("# EyeQ CFP Quality Control")
+        gr.Markdown(
+            f"Model: `{model_name}`  \n"
+            f"Input size: `{img_size} × {img_size}`  \n"
+            f"Device: `{device}`  \n"
+            f"Checkpoint: `{checkpoint_path}`"
+        )
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                image_input = gr.Image(
+                    label="Input CFP",
+                    type="pil",
+                    height=520,
+                )
+                with gr.Accordion("Soft Reject rule", open=True):
+                    reject_threshold = gr.Slider(
+                        minimum=0.40,
+                        maximum=0.95,
+                        value=0.60,
+                        step=0.01,
+                        label="Reject threshold",
+                        info="Minimum Reject probability required before an image can be called Reject.",
+                    )
+                    reject_margin = gr.Slider(
+                        minimum=0.00,
+                        maximum=0.50,
+                        value=0.15,
+                        step=0.01,
+                        label="Reject margin",
+                        info="Reject must beat both Good and Usable by at least this much.",
+                    )
+
+                run_button = gr.Button("Run QC", variant="primary")
+
+            with gr.Column(scale=1):
+                pred_output = gr.Label(label="Predicted quality")
+                prob_output = gr.Label(label="Class probabilities", num_top_classes=3)
+                decision_output = gr.Textbox(
+                    label="Decision details",
+                    lines=4,
+                    interactive=False,
+                )
+
+        run_inputs = [image_input, reject_threshold, reject_margin]
+        run_outputs = [pred_output, prob_output, decision_output]
+
+        run_button.click(
+            fn=run,
+            inputs=run_inputs,
+            outputs=run_outputs,
+        )
+
+        image_input.change(
+            fn=run,
+            inputs=run_inputs,
+            outputs=run_outputs,
+        )
+
+        reject_threshold.change(
+            fn=update_margin_slider,
+            inputs=[reject_threshold, reject_margin],
+            outputs=reject_margin,
+        ).then(
+            fn=run,
+            inputs=run_inputs,
+            outputs=run_outputs,
+        )
+
+        reject_margin.change(
+            fn=run,
+            inputs=run_inputs,
+            outputs=run_outputs,
+        )
+
+    return demo
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", type=str, default="./checkpoints/eyeq_vit_base/eyeq_deploy.pt")
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    parser.add_argument("--port", type=int, default=7860)
+    parser.add_argument("--share", action="store_true")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    checkpoint_path = Path(args.checkpoint)
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
+
+    demo = make_app(str(checkpoint_path))
+    demo.launch(
+        # server_name=args.host,
+        # server_port=args.port,
+        # share=args.share,
+    )
+
+
+if __name__ == "__main__":
+    main()

checkpoints/eyeq_vit_base/best_report.txt

@@ -0,0 +1,17 @@
+Best epoch: 3
+Best test balanced accuracy: 0.8573
+
+              precision    recall  f1-score   support
+
+        Good     0.9262    0.9337    0.9299      8471
+      Usable     0.7829    0.7760    0.7794      4558
+      Reject     0.8697    0.8621    0.8659      3220
+
+    accuracy                         0.8753     16249
+   macro avg     0.8596    0.8573    0.8584     16249
+weighted avg     0.8748    0.8753    0.8750     16249
+
+Confusion matrix rows=true cols=pred, labels=[Good, Usable, Reject]
+[[7909  556    6]
+ [ 611 3537  410]
+ [  19  425 2776]]

checkpoints/eyeq_vit_base/eyeq_deploy.pt

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

checkpoints/eyeq_vit_base/test_eval/test_confusion_matrix.csv

@@ -0,0 +1,4 @@
+,pred_Good,pred_Usable,pred_Reject
+true_Good,7908,557,6
+true_Usable,611,3537,410
+true_Reject,19,426,2775

checkpoints/eyeq_vit_base/test_eval/test_report.txt

@@ -0,0 +1,25 @@
+Checkpoint: checkpoints/eyeq_vit_base/best.pt
+Test CSV: /data/MIDS/datasets/retina/EyeQ/data/Label_EyeQ_test.csv
+Test images: /data/MIDS/datasets/retina/EyePACS/test
+Model: vit_base_patch16_224
+Image size: 224
+Device: cuda
+
+test_loss=0.312007
+test_acc=0.875131
+test_bal_acc=0.857112
+
+              precision    recall  f1-score   support
+
+        Good     0.9262    0.9335    0.9299      8471
+      Usable     0.7825    0.7760    0.7792      4558
+      Reject     0.8696    0.8618    0.8657      3220
+
+    accuracy                         0.8751     16249
+   macro avg     0.8595    0.8571    0.8583     16249
+weighted avg     0.8747    0.8751    0.8749     16249
+
+Confusion matrix rows=true cols=pred, labels=[Good, Usable, Reject]
+[[7908  557    6]
+ [ 611 3537  410]
+ [  19  426 2775]]

requirements.txt

@@ -0,0 +1,13 @@
+albumentations
+gradio
+huggingface_hub
+numpy
+opencv-python
+pandas
+pillow
+pydantic
+timm
+torch
+torchvision
+torchaudio
+tqdm

test.py

@@ -0,0 +1,349 @@
+#!/usr/bin/env python3
+"""
+Evaluate an EyeQ CFP image-quality-control model on Label_EyeQ_test.csv.
+
+Example
+-------
+python EyeQ_test.py \
+  --images_dir /data/MIDS/datasets/retina/EyePACS \
+  --csv_dir /data/MIDS/datasets/retina/EyeQ/data \
+  --checkpoint ./checkpoints/eyeq_vit_base/best.pt \
+  --output_dir ./checkpoints/eyeq_vit_base/test_eval \
+  --batch_size 32 \
+  --num_workers 24
+"""
+
+import argparse
+from pathlib import Path
+from typing import Dict, Tuple
+
+import numpy as np
+import pandas as pd
+from PIL import Image
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+
+import timm
+from sklearn.metrics import (
+    accuracy_score,
+    balanced_accuracy_score,
+    classification_report,
+    confusion_matrix,
+)
+from tqdm import tqdm
+
+
+ID_TO_LABEL = {0: "Good", 1: "Usable", 2: "Reject"}
+LABEL_TO_ID: Dict[str, int] = {
+    "good": 0,
+    "usable": 1,
+    "reject": 2,
+    "0": 0,
+    "1": 1,
+    "2": 2,
+}
+
+
+class EyeQDataset(Dataset):
+    def __init__(self, df: pd.DataFrame, images_dir: str, transform=None):
+        self.df = df.reset_index(drop=True)
+        self.images_dir = Path(images_dir)
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, idx):
+        row = self.df.iloc[idx]
+
+        image_name = str(row["image"])
+        image_path = self.images_dir / image_name
+
+        image = Image.open(image_path).convert("RGB")
+        label = int(row["quality"])
+
+        if self.transform is not None:
+            image = self.transform(image)
+
+        return image, label, image_name
+
+
+def normalize_quality_label(x) -> int:
+    key = str(x).strip().lower()
+
+    if key in LABEL_TO_ID:
+        return LABEL_TO_ID[key]
+
+    try:
+        value = int(float(key))
+        if value in [0, 1, 2]:
+            return value
+    except ValueError:
+        pass
+
+    raise ValueError(f"Unknown quality label: {x}. Expected 0/1/2 or Good/Usable/Reject.")
+
+
+def load_eyeq_csv(csv_path: str, images_dir: str) -> pd.DataFrame:
+    df = pd.read_csv(csv_path)
+
+    if "image" not in df.columns:
+        raise ValueError(f"CSV must contain an 'image' column. Found columns: {list(df.columns)}")
+    if "quality" not in df.columns:
+        raise ValueError(f"CSV must contain a 'quality' column. Found columns: {list(df.columns)}")
+
+    # Keep DR_grade if present for optional downstream inspection.
+    keep_cols = ["image", "quality"]
+    if "DR_grade" in df.columns:
+        keep_cols.append("DR_grade")
+
+    df = df[keep_cols].copy()
+    df["image"] = df["image"].astype(str)
+    df["quality"] = df["quality"].apply(normalize_quality_label)
+
+    images_dir = Path(images_dir)
+    exists = df["image"].apply(lambda x: (images_dir / x).exists())
+
+    missing = int((~exists).sum())
+    if missing > 0:
+        print(f"Warning: dropping {missing} rows with missing image files from {csv_path}")
+        print(f"         searched in: {images_dir}")
+
+    df = df.loc[exists].reset_index(drop=True)
+
+    if len(df) == 0:
+        raise RuntimeError(f"No valid images found for {csv_path}. Searched in: {images_dir}")
+
+    return df
+
+
+def build_transform(img_size: int):
+    return transforms.Compose([
+        transforms.Resize((img_size, img_size)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=(0.485, 0.456, 0.406),
+            std=(0.229, 0.224, 0.225),
+        ),
+    ])
+
+
+def load_model(checkpoint_path: str, device: torch.device):
+    ckpt = torch.load(checkpoint_path, map_location="cpu")
+
+    ckpt_args = ckpt.get("args", {})
+    model_name = ckpt_args.get("model", "vit_base_patch16_224")
+    img_size = int(ckpt_args.get("img_size", 224))
+
+    id_to_label = ckpt.get("id_to_label", ID_TO_LABEL)
+    id_to_label = {int(k): str(v) for k, v in id_to_label.items()}
+
+    model = timm.create_model(
+        model_name,
+        pretrained=False,
+        num_classes=len(id_to_label),
+    )
+
+    model.load_state_dict(ckpt["model"], strict=True)
+    model.to(device)
+    model.eval()
+
+    return model, id_to_label, model_name, img_size, ckpt
+
+
+@torch.no_grad()
+def evaluate(model, loader, criterion, device, amp=False):
+    model.eval()
+
+    running_loss = 0.0
+    all_labels = []
+    all_preds = []
+    all_probs = []
+    all_images = []
+
+    for images, labels, image_names in tqdm(loader, desc="Test"):
+        images = images.to(device, non_blocking=True)
+        labels = labels.to(device, non_blocking=True)
+
+        with torch.cuda.amp.autocast(enabled=amp and device.type == "cuda"):
+            logits = model(images)
+            loss = criterion(logits, labels)
+            probs = torch.softmax(logits, dim=1)
+
+        preds = probs.argmax(dim=1)
+
+        running_loss += loss.item() * images.size(0)
+
+        all_labels.extend(labels.detach().cpu().numpy().tolist())
+        all_preds.extend(preds.detach().cpu().numpy().tolist())
+        all_probs.extend(probs.detach().cpu().numpy().tolist())
+        all_images.extend(list(image_names))
+
+    test_loss = running_loss / len(loader.dataset)
+
+    y_true = np.array(all_labels)
+    y_pred = np.array(all_preds)
+    probs = np.array(all_probs)
+
+    acc = accuracy_score(y_true, y_pred)
+    bal_acc = balanced_accuracy_score(y_true, y_pred)
+
+    return test_loss, acc, bal_acc, y_true, y_pred, probs, all_images
+
+
+def print_label_counts(name: str, df: pd.DataFrame):
+    print(f"{name}: {len(df)}")
+    for label_id in [0, 1, 2]:
+        count = int((df["quality"] == label_id).sum())
+        print(f"  {ID_TO_LABEL[label_id]} ({label_id}): {count}")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--images_dir", type=str, required=True,
+                        help="EyePACS root containing train/ and test/ folders.")
+    parser.add_argument("--csv_dir", type=str, required=True,
+                        help="Directory containing Label_EyeQ_test.csv.")
+    parser.add_argument("--checkpoint", type=str, default="./checkpoints/eyeq_vit_base/best.pt")
+    parser.add_argument("--output_dir", type=str, default=None)
+
+    parser.add_argument("--batch_size", type=int, default=32)
+    parser.add_argument("--num_workers", type=int, default=8)
+    parser.add_argument("--amp", action="store_true", default=True)
+    parser.add_argument("--no_amp", dest="amp", action="store_false")
+    parser.add_argument("--cpu", action="store_true")
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    images_root = Path(args.images_dir)
+    csv_root = Path(args.csv_dir)
+    checkpoint_path = Path(args.checkpoint)
+
+    test_images_dir = images_root / "test"
+    test_csv = csv_root / "Label_EyeQ_test.csv"
+
+    if args.output_dir is None:
+        output_dir = checkpoint_path.parent / "test_eval"
+    else:
+        output_dir = Path(args.output_dir)
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
+    if not test_images_dir.exists():
+        raise FileNotFoundError(f"Test image directory not found: {test_images_dir}")
+    if not test_csv.exists():
+        raise FileNotFoundError(f"Test CSV not found: {test_csv}")
+
+    device = torch.device("cuda" if torch.cuda.is_available() and not args.cpu else "cpu")
+
+    model, id_to_label, model_name, img_size, ckpt = load_model(str(checkpoint_path), device)
+    transform = build_transform(img_size)
+
+    test_df = load_eyeq_csv(str(test_csv), str(test_images_dir))
+    test_ds = EyeQDataset(test_df, str(test_images_dir), transform)
+
+    test_loader = DataLoader(
+        test_ds,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        pin_memory=(device.type == "cuda"),
+        persistent_workers=(args.num_workers > 0),
+    )
+
+    criterion = nn.CrossEntropyLoss()
+
+    print("Evaluation summary")
+    print(f"Checkpoint:  {checkpoint_path}")
+    print(f"Test CSV:    {test_csv}")
+    print(f"Test images: {test_images_dir}")
+    print(f"Output dir:  {output_dir}")
+    print(f"Model:       {model_name}")
+    print(f"Image size:  {img_size}")
+    print(f"Device:      {device}")
+    print(f"Labels:      {id_to_label}")
+    print_label_counts("Test", test_df)
+
+    test_loss, acc, bal_acc, y_true, y_pred, probs, image_names = evaluate(
+        model=model,
+        loader=test_loader,
+        criterion=criterion,
+        device=device,
+        amp=args.amp,
+    )
+
+    target_names = [id_to_label[i] for i in [0, 1, 2]]
+
+    report = classification_report(
+        y_true,
+        y_pred,
+        labels=[0, 1, 2],
+        target_names=target_names,
+        digits=4,
+    )
+
+    cm = confusion_matrix(y_true, y_pred, labels=[0, 1, 2])
+
+    print()
+    print(f"test_loss={test_loss:.4f}")
+    print(f"test_acc={acc:.4f}")
+    print(f"test_bal_acc={bal_acc:.4f}")
+    print()
+    print(report)
+    print("Confusion matrix rows=true cols=pred, labels=[Good, Usable, Reject]")
+    print(cm)
+
+    # Save text report
+    with open(output_dir / "test_report.txt", "w") as f:
+        f.write(f"Checkpoint: {checkpoint_path}\n")
+        f.write(f"Test CSV: {test_csv}\n")
+        f.write(f"Test images: {test_images_dir}\n")
+        f.write(f"Model: {model_name}\n")
+        f.write(f"Image size: {img_size}\n")
+        f.write(f"Device: {device}\n\n")
+        f.write(f"test_loss={test_loss:.6f}\n")
+        f.write(f"test_acc={acc:.6f}\n")
+        f.write(f"test_bal_acc={bal_acc:.6f}\n\n")
+        f.write(report)
+        f.write("\nConfusion matrix rows=true cols=pred, labels=[Good, Usable, Reject]\n")
+        f.write(str(cm))
+        f.write("\n")
+
+    # Save confusion matrix CSV
+    cm_df = pd.DataFrame(
+        cm,
+        index=[f"true_{name}" for name in target_names],
+        columns=[f"pred_{name}" for name in target_names],
+    )
+    cm_df.to_csv(output_dir / "test_confusion_matrix.csv")
+
+    # Save per-image predictions
+    pred_df = test_df.copy()
+    pred_df["pred_quality"] = y_pred
+    pred_df["true_label"] = [id_to_label[int(x)] for x in y_true]
+    pred_df["pred_label"] = [id_to_label[int(x)] for x in y_pred]
+    pred_df["prob_good"] = probs[:, 0]
+    pred_df["prob_usable"] = probs[:, 1]
+    pred_df["prob_reject"] = probs[:, 2]
+    pred_df["correct"] = pred_df["quality"].values == pred_df["pred_quality"].values
+
+    pred_df.to_csv(output_dir / "test_predictions.csv", index=False)
+
+    print()
+    print(f"Saved report:      {output_dir / 'test_report.txt'}")
+    print(f"Saved confusion:   {output_dir / 'test_confusion_matrix.csv'}")
+    print(f"Saved predictions: {output_dir / 'test_predictions.csv'}")
+
+
+if __name__ == "__main__":
+    main()

train.py

@@ -0,0 +1,397 @@
+#!/usr/bin/env python3
+"""
+Train a CFP image-quality-control model on EyeQ / EyePACS-style data.
+
+Expected dataset format
+-----------------------
+EyePACS/
+    train/
+        10009_left.jpeg
+        10009_right.jpeg
+        ...
+    test/
+        ...
+
+data/
+    Label_EyeQ_train.csv
+    Label_EyeQ_test.csv
+
+Label CSV format:
+    ,image,quality,DR_grade
+    0,10009_left.jpeg,0,0
+    1,10009_right.jpeg,0,0
+    2,10014_left.jpeg,2,0
+
+For EyeQ, this script assumes:
+    quality = 0 -> Good
+    quality = 1 -> Usable
+    quality = 2 -> Reject
+
+DR_grade is ignored because this script trains only the image-quality model.
+
+Example
+-------
+python EyeQ_train.py \
+    --images_dir /path/to/EyePACS \
+    --csv_dir /path/to/data \
+    --output_dir ./runs/eyeq_vit_base \
+    --epochs 30 \
+    --batch_size 32 \
+    --lr 3e-5
+"""
+
+import argparse
+import random
+from pathlib import Path
+from typing import Dict, Tuple
+
+import numpy as np
+import pandas as pd
+from PIL import Image
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+
+import timm
+from sklearn.metrics import accuracy_score, balanced_accuracy_score, classification_report, confusion_matrix
+from tqdm import tqdm
+
+
+ID_TO_LABEL = {0: "Good", 1: "Usable", 2: "Reject"}
+LABEL_TO_ID: Dict[str, int] = {
+    "good": 0,
+    "usable": 1,
+    "reject": 2,
+    "0": 0,
+    "1": 1,
+    "2": 2,
+}
+
+
+class EyeQDataset(Dataset):
+    def __init__(self, df: pd.DataFrame, images_dir: str, transform=None):
+        self.df = df.reset_index(drop=True)
+        self.images_dir = Path(images_dir)
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, idx):
+        row = self.df.iloc[idx]
+        image_path = self.images_dir / str(row["image"])
+        image = Image.open(image_path).convert("RGB")
+        label = int(row["quality"])
+
+        if self.transform is not None:
+            image = self.transform(image)
+
+        return image, label
+
+
+def seed_everything(seed: int):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.benchmark = True
+
+
+def normalize_quality_label(x) -> int:
+    key = str(x).strip().lower()
+    if key in LABEL_TO_ID:
+        return LABEL_TO_ID[key]
+    try:
+        value = int(float(key))
+        if value in [0, 1, 2]:
+            return value
+    except ValueError:
+        pass
+    raise ValueError(f"Unknown quality label: {x}. Expected 0/1/2 or Good/Usable/Reject.")
+
+
+def load_eyeq_csv(csv_path: str, images_dir: str) -> pd.DataFrame:
+    df = pd.read_csv(csv_path)
+
+    if "image" not in df.columns:
+        raise ValueError(f"CSV must contain an 'image' column. Found columns: {list(df.columns)}")
+    if "quality" not in df.columns:
+        raise ValueError(f"CSV must contain a 'quality' column. Found columns: {list(df.columns)}")
+
+    df = df[["image", "quality"]].copy()
+    df["image"] = df["image"].astype(str)
+    df["quality"] = df["quality"].apply(normalize_quality_label)
+
+    images_dir = Path(images_dir)
+    exists = df["image"].apply(lambda x: (images_dir / x).exists())
+    missing = int((~exists).sum())
+    if missing > 0:
+        print(f"Warning: dropping {missing} rows with missing image files from {csv_path}")
+        print(f"         searched in: {images_dir}")
+    df = df.loc[exists].reset_index(drop=True)
+
+    if len(df) == 0:
+        raise RuntimeError(f"No valid images found for {csv_path}. Searched in: {images_dir}")
+
+    return df
+
+
+def build_transforms(img_size: int) -> Tuple[transforms.Compose, transforms.Compose]:
+    train_tfms = transforms.Compose([
+        transforms.Resize((img_size, img_size)),
+        transforms.RandomHorizontalFlip(p=0.5),
+        transforms.RandomApply([
+            transforms.ColorJitter(
+                brightness=0.15,
+                contrast=0.15,
+                saturation=0.10,
+                hue=0.02,
+            )
+        ], p=0.8),
+        transforms.RandomApply([
+            transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 1.0))
+        ], p=0.15),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+    ])
+
+    test_tfms = transforms.Compose([
+        transforms.Resize((img_size, img_size)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+    ])
+
+    return train_tfms, test_tfms
+
+
+def build_model(model_name: str, num_classes: int, pretrained: bool):
+    return timm.create_model(
+        model_name,
+        pretrained=pretrained,
+        num_classes=num_classes,
+    )
+
+
+def train_one_epoch(model, loader, criterion, optimizer, scaler, device, epoch):
+    model.train()
+    running_loss = 0.0
+    all_preds = []
+    all_labels = []
+
+    pbar = tqdm(loader, desc=f"Train {epoch}", leave=False)
+    for images, labels in pbar:
+        images = images.to(device, non_blocking=True)
+        labels = labels.to(device, non_blocking=True)
+
+        optimizer.zero_grad(set_to_none=True)
+
+        with torch.cuda.amp.autocast(enabled=scaler is not None):
+            logits = model(images)
+            loss = criterion(logits, labels)
+
+        if scaler is not None:
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            loss.backward()
+            optimizer.step()
+
+        running_loss += loss.item() * images.size(0)
+        preds = logits.argmax(dim=1)
+        all_preds.extend(preds.detach().cpu().numpy().tolist())
+        all_labels.extend(labels.detach().cpu().numpy().tolist())
+
+        pbar.set_postfix(loss=f"{loss.item():.4f}")
+
+    epoch_loss = running_loss / len(loader.dataset)
+    acc = accuracy_score(all_labels, all_preds)
+    bal_acc = balanced_accuracy_score(all_labels, all_preds)
+    return epoch_loss, acc, bal_acc
+
+
+@torch.no_grad()
+def evaluate(model, loader, criterion, device, split_name="Test"):
+    model.eval()
+    running_loss = 0.0
+    all_preds = []
+    all_labels = []
+
+    pbar = tqdm(loader, desc=split_name, leave=False)
+    for images, labels in pbar:
+        images = images.to(device, non_blocking=True)
+        labels = labels.to(device, non_blocking=True)
+
+        logits = model(images)
+        loss = criterion(logits, labels)
+
+        running_loss += loss.item() * images.size(0)
+        preds = logits.argmax(dim=1)
+        all_preds.extend(preds.detach().cpu().numpy().tolist())
+        all_labels.extend(labels.detach().cpu().numpy().tolist())
+
+    val_loss = running_loss / len(loader.dataset)
+    acc = accuracy_score(all_labels, all_preds)
+    bal_acc = balanced_accuracy_score(all_labels, all_preds)
+    return val_loss, acc, bal_acc, np.array(all_labels), np.array(all_preds)
+
+
+def save_checkpoint(path, model, optimizer, scheduler, epoch, best_metric, args):
+    torch.save({
+        "epoch": epoch,
+        "model": model.state_dict(),
+        "optimizer": optimizer.state_dict(),
+        "scheduler": scheduler.state_dict() if scheduler is not None else None,
+        "best_metric": best_metric,
+        "args": vars(args),
+        "id_to_label": ID_TO_LABEL,
+    }, path)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--images_dir", type=str, required=True, help="EyePACS root containing train/ and test/ folders.")
+    parser.add_argument("--csv_dir", type=str, required=True, help="Directory containing Label_EyeQ_train.csv and Label_EyeQ_test.csv.")
+    parser.add_argument("--output_dir", type=str, default="./runs/eyeq_vit_base")
+
+    parser.add_argument("--model", type=str, default="vit_base_patch16_224")
+    parser.add_argument("--img_size", type=int, default=224)
+    parser.add_argument("--pretrained", action="store_true", default=True)
+    parser.add_argument("--no_pretrained", dest="pretrained", action="store_false")
+
+    parser.add_argument("--epochs", type=int, default=30)
+    parser.add_argument("--batch_size", type=int, default=32)
+    parser.add_argument("--num_workers", type=int, default=8)
+    parser.add_argument("--lr", type=float, default=3e-5)
+    parser.add_argument("--weight_decay", type=float, default=1e-4)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--amp", action="store_true", default=True)
+    parser.add_argument("--no_amp", dest="amp", action="store_false")
+    parser.add_argument("--class_weights", action="store_true", help="Use inverse-frequency class weights.")
+
+    return parser.parse_args()
+
+
+def print_label_counts(name: str, df: pd.DataFrame):
+    print(f"{name}: {len(df)}")
+    for label_id in [0, 1, 2]:
+        count = int((df["quality"] == label_id).sum())
+        print(f"  {ID_TO_LABEL[label_id]} ({label_id}): {count}")
+
+
+def main():
+    args = parse_args()
+    seed_everything(args.seed)
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    images_root = Path(args.images_dir)
+    csv_root = Path(args.csv_dir)
+
+    train_images_dir = images_root / "train"
+    test_images_dir = images_root / "test"
+    train_csv = csv_root / "Label_EyeQ_train.csv"
+    test_csv = csv_root / "Label_EyeQ_test.csv"
+
+    train_df = load_eyeq_csv(str(train_csv), str(train_images_dir))
+    test_df = load_eyeq_csv(str(test_csv), str(test_images_dir))
+
+    train_tfms, test_tfms = build_transforms(args.img_size)
+
+    train_ds = EyeQDataset(train_df, str(train_images_dir), train_tfms)
+    test_ds = EyeQDataset(test_df, str(test_images_dir), test_tfms)
+
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=True,
+    )
+    test_loader = DataLoader(
+        test_ds,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        pin_memory=True,
+    )
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = build_model(args.model, num_classes=3, pretrained=args.pretrained).to(device)
+
+    if args.class_weights:
+        counts = train_df["quality"].value_counts().sort_index().reindex([0, 1, 2], fill_value=1).values
+        weights = counts.sum() / (len(counts) * counts)
+        weights = torch.tensor(weights, dtype=torch.float32, device=device)
+        criterion = nn.CrossEntropyLoss(weight=weights)
+        print(f"Using class weights: {weights.detach().cpu().numpy().round(3).tolist()}")
+    else:
+        criterion = nn.CrossEntropyLoss()
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
+    scaler = torch.cuda.amp.GradScaler() if args.amp and device.type == "cuda" else None
+
+    print("Dataset summary")
+    print(f"Train CSV:    {train_csv}")
+    print(f"Test CSV:     {test_csv}")
+    print(f"Train images: {train_images_dir}")
+    print(f"Test images:  {test_images_dir}")
+    print_label_counts("Train", train_df)
+    print_label_counts("Test", test_df)
+    print(f"Model: {args.model}")
+    print(f"Device: {device}")
+
+    best_bal_acc = -1.0
+
+    for epoch in range(1, args.epochs + 1):
+        train_loss, train_acc, train_bal_acc = train_one_epoch(
+            model, train_loader, criterion, optimizer, scaler, device, epoch
+        )
+        test_loss, test_acc, test_bal_acc, y_true, y_pred = evaluate(
+            model, test_loader, criterion, device, split_name="Test"
+        )
+        scheduler.step()
+
+        print(
+            f"Epoch {epoch:03d}/{args.epochs} | "
+            f"train_loss={train_loss:.4f} train_acc={train_acc:.4f} train_bal_acc={train_bal_acc:.4f} | "
+            f"test_loss={test_loss:.4f} test_acc={test_acc:.4f} test_bal_acc={test_bal_acc:.4f}"
+        )
+
+        save_checkpoint(output_dir / "last.pt", model, optimizer, scheduler, epoch, best_bal_acc, args)
+
+        if test_bal_acc > best_bal_acc:
+            best_bal_acc = test_bal_acc
+            best_path = output_dir / "best.pt"
+            save_checkpoint(best_path, model, optimizer, scheduler, epoch, best_bal_acc, args)
+
+            report = classification_report(
+                y_true,
+                y_pred,
+                labels=[0, 1, 2],
+                target_names=[ID_TO_LABEL[i] for i in [0, 1, 2]],
+                digits=4,
+            )
+            cm = confusion_matrix(y_true, y_pred, labels=[0, 1, 2])
+
+            with open(output_dir / "best_report.txt", "w") as f:
+                f.write(f"Best epoch: {epoch}\n")
+                f.write(f"Best test balanced accuracy: {best_bal_acc:.4f}\n\n")
+                f.write(report)
+                f.write("\nConfusion matrix rows=true cols=pred, labels=[Good, Usable, Reject]\n")
+                f.write(str(cm))
+                f.write("\n")
+
+            print(f"  Saved new best checkpoint: {best_path}")
+
+    print(f"Training complete. Best test balanced accuracy: {best_bal_acc:.4f}")
+    print(f"Outputs saved to: {output_dir}")
+
+
+if __name__ == "__main__":
+    main()