Upload train_thyroid.py

train_thyroid.py

@@ -0,0 +1,204 @@
+#!/usr/bin/env python3
+"""
+Thyroid Ultrasound Nodule Malignancy Classification
+Dataset: BTX24/thyroid-cancer-classification-ultrasound-dataset
+Binary classification: benign (0) vs malignant (1)
+"""
+
+import os
+import sys
+import numpy as np
+from collections import Counter
+
+from datasets import load_dataset
+from transformers import (
+    AutoImageProcessor,
+    AutoModelForImageClassification,
+    TrainingArguments,
+    Trainer,
+    DefaultDataCollator,
+    EarlyStoppingCallback,
+)
+import evaluate
+import torch
+from torchvision.transforms import (
+    Compose, Resize, RandomRotation, RandomHorizontalFlip,
+    RandomVerticalFlip, ColorJitter, ToTensor, Normalize
+)
+
+# ------------------------------------------------------------------
+# Config
+# ------------------------------------------------------------------
+DATASET_NAME = "BTX24/thyroid-cancer-classification-ultrasound-dataset"
+MODEL_NAME = "microsoft/swinv2-base-patch4-window8-256"
+OUTPUT_DIR = "./thyroid-swinv2-model"
+HUB_MODEL_ID = "Johnyquest7/ML-Inter_thyroid"
+
+NUM_LABELS = 2
+ID2LABEL = {0: "benign", 1: "malignant"}
+LABEL2ID = {"benign": 0, "malignant": 1}
+
+# ------------------------------------------------------------------
+# Metrics
+# ------------------------------------------------------------------
+accuracy = evaluate.load("accuracy")
+f1 = evaluate.load("f1")
+precision = evaluate.load("precision")
+recall = evaluate.load("recall")
+roc_auc = evaluate.load("roc_auc")
+
+def compute_metrics(eval_pred):
+    logits, labels = eval_pred
+    preds = np.argmax(logits, axis=1)
+    probs = torch.softmax(torch.tensor(logits), dim=1)[:, 1].numpy()
+
+    result = {}
+    result.update(accuracy.compute(predictions=preds, references=labels))
+    result.update(f1.compute(predictions=preds, references=labels, average="binary"))
+    result.update(precision.compute(predictions=preds, references=labels, average="binary"))
+    result.update(recall.compute(predictions=preds, references=labels, average="binary"))
+    try:
+        result.update(roc_auc.compute(prediction_scores=probs, references=labels))
+    except Exception:
+        result["roc_auc"] = 0.0
+    return result
+
+# ------------------------------------------------------------------
+# Load dataset
+# ------------------------------------------------------------------
+print("Loading dataset...")
+train_ds = load_dataset(DATASET_NAME, split="train")
+test_ds = load_dataset(DATASET_NAME, split="test")
+
+# Create validation split from train
+train_val = train_ds.train_test_split(test_size=0.15, stratify_by_column="label", seed=42)
+train_ds = train_val["train"]
+val_ds = train_val["test"]
+
+print(f"Train: {len(train_ds)} | Val: {len(val_ds)} | Test: {len(test_ds)}")
+print(f"Train labels: {Counter(train_ds['label'])}")
+print(f"Val labels: {Counter(val_ds['label'])}")
+print(f"Test labels: {Counter(test_ds['label'])}")
+
+# ------------------------------------------------------------------
+# Image processor & transforms
+# ------------------------------------------------------------------
+print(f"Loading image processor from {MODEL_NAME}...")
+image_processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
+
+# Ultrasound images are grayscale (mode 'L') — convert to RGB for SwinV2
+image_mean = image_processor.image_mean
+image_std = image_processor.image_std
+size = (
+    image_processor.size["shortest_edge"]
+    if "shortest_edge" in image_processor.size
+    else (image_processor.size["height"], image_processor.size["width"])
+)
+
+train_transforms = Compose([
+    Resize(size),
+    RandomRotation(degrees=10),
+    RandomHorizontalFlip(p=0.5),
+    RandomVerticalFlip(p=0.3),
+    ColorJitter(brightness=0.2, contrast=0.2),
+    ToTensor(),
+    Normalize(mean=image_mean, std=image_std),
+])
+
+val_transforms = Compose([
+    Resize(size),
+    ToTensor(),
+    Normalize(mean=image_mean, std=image_std),
+])
+
+def preprocess_train(examples):
+    # Convert grayscale to RGB
+    examples["pixel_values"] = [
+        train_transforms(img.convert("RGB")) for img in examples["image"]
+    ]
+    del examples["image"]
+    return examples
+
+def preprocess_val(examples):
+    examples["pixel_values"] = [
+        val_transforms(img.convert("RGB")) for img in examples["image"]
+    ]
+    del examples["image"]
+    return examples
+
+print("Applying transforms...")
+train_ds = train_ds.with_transform(preprocess_train)
+val_ds = val_ds.with_transform(preprocess_val)
+test_ds = test_ds.with_transform(preprocess_val)
+
+# ------------------------------------------------------------------
+# Model
+# ------------------------------------------------------------------
+print(f"Loading model {MODEL_NAME}...")
+model = AutoModelForImageClassification.from_pretrained(
+    MODEL_NAME,
+    num_labels=NUM_LABELS,
+    id2label=ID2LABEL,
+    label2id=LABEL2ID,
+    ignore_mismatched_sizes=True,
+)
+
+# ------------------------------------------------------------------
+# Training arguments
+# ------------------------------------------------------------------
+training_args = TrainingArguments(
+    output_dir=OUTPUT_DIR,
+    remove_unused_columns=False,
+    eval_strategy="epoch",
+    save_strategy="epoch",
+    learning_rate=2e-5,
+    per_device_train_batch_size=16,
+    per_device_eval_batch_size=16,
+    gradient_accumulation_steps=2,
+    num_train_epochs=30,
+    warmup_steps=100,
+    weight_decay=0.01,
+    logging_strategy="steps",
+    logging_steps=10,
+    logging_first_step=True,
+    disable_tqdm=True,
+    load_best_model_at_end=True,
+    metric_for_best_model="eval_roc_auc",
+    greater_is_better=True,
+    push_to_hub=True,
+    hub_model_id=HUB_MODEL_ID,
+    report_to="trackio",
+    run_name="thyroid-swinv2-binary",
+    project="thyroid-malignancy",
+    seed=42,
+    bf16=True,
+    dataloader_num_workers=4,
+)
+
+# ------------------------------------------------------------------
+# Trainer
+# ------------------------------------------------------------------
+data_collator = DefaultDataCollator()
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    data_collator=data_collator,
+    train_dataset=train_ds,
+    eval_dataset=val_ds,
+    processing_class=image_processor,
+    compute_metrics=compute_metrics,
+    callbacks=[EarlyStoppingCallback(early_stopping_patience=5)],
+)
+
+print("Starting training...")
+trainer.train()
+
+print("Evaluating on test set...")
+test_results = trainer.evaluate(test_ds, metric_key_prefix="test")
+print("Test results:", test_results)
+
+print("Pushing to Hub...")
+trainer.push_to_hub()
+
+print("Done!")