Update src/train_vision_model.py

src/train_vision_model.py

@@ -1,172 +1,292 @@
-"""Train and persist the local vehicle-class vision model.
+"""Train a transfer-learning vehicle brand classifier using ResNet-18.
 
-This uses a lightweight local classifier on handcrafted image features from the
-expanded dataset stored under data/raw/Cars Dataset/train and data/raw/Cars Dataset/test.
+This script uses:
+- datasets.load_dataset("imagefolder") for data loading
+- AutoImageProcessor for image preprocessing
+- AutoModelForImageClassification for transfer learning
+- transformers.Trainer for training
+- Proper data augmentation and evaluation
+
+The trained model is saved in Hugging Face format under models/car-image-classifier/
 """
 
 from __future__ import annotations
 
+import argparse
 import json
+import random
 import sys
 from datetime import datetime, timezone
 from pathlib import Path
 
+import evaluate
+import numpy as np
+from datasets import load_dataset
+from sklearn.metrics import classification_report
+from transformers import (
+    AutoImageProcessor,
+    AutoModelForImageClassification,
+    Trainer,
+    TrainingArguments,
+    set_seed,
+)
+
 PROJECT_ROOT = Path(__file__).resolve().parent.parent
 if str(PROJECT_ROOT) not in sys.path:
     sys.path.insert(0, str(PROJECT_ROOT))
 
-import joblib
-import numpy as np
-from PIL import Image
-from sklearn.ensemble import ExtraTreesClassifier
-from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
-
-from src.config import DATA_RAW_DIR, MODEL_DIR, VISION_LABELS_PATH, VISION_METADATA_PATH, VISION_MODEL_PATH
-from src.vision_features import extract_image_features
-
-
-def _find_image_root() -> Path:
-    candidates = [
-        DATA_RAW_DIR / "Cars Dataset",
-        DATA_RAW_DIR / "car_images",
-    ]
-    for candidate in candidates:
-        if candidate.exists():
-            return candidate
-    raise FileNotFoundError("No image dataset found in data/raw/Cars Dataset or data/raw/car_images.")
-
-
-def _collect_image_paths(image_root: Path) -> tuple[list[Path], list[Path], list[str]]:
-    train_dir = image_root / "train"
-    test_dir = image_root / "test"
-    if not train_dir.exists() or not test_dir.exists():
-        raise FileNotFoundError("Expected train/ and test/ folders inside the image dataset root.")
-
-    classes = sorted([path.name for path in train_dir.iterdir() if path.is_dir()])
-    train_paths: list[Path] = []
-    test_paths: list[Path] = []
-
-    for class_name in classes:
-        for split_paths, split_dir in ((train_paths, train_dir), (test_paths, test_dir)):
-            class_dir = split_dir / class_name
-            for pattern in ("*.jpg", "*.jpeg", "*.png", "*.webp"):
-                split_paths.extend(sorted(class_dir.glob(pattern)))
-
-    return train_paths, test_paths, classes
-
-
-def _load_grouped_features(paths: list[Path], classes: list[str]) -> dict[str, list[np.ndarray]]:
-    grouped_features: dict[str, list[np.ndarray]] = {class_name: [] for class_name in classes}
-
-    for image_path in paths:
-        class_name = image_path.parent.name
-        try:
-            image = Image.open(image_path)
-            grouped_features[class_name].append(extract_image_features(image))
-        except Exception:
-            continue
-
-    return grouped_features
-
-
-def _rebalance_training_data(
-    grouped_features: dict[str, list[np.ndarray]],
-    classes: list[str],
-    target_per_class: int = 200,
-) -> tuple[np.ndarray, np.ndarray]:
-    rng = np.random.default_rng(42)
-    features = []
-    labels = []
-
-    for class_index, class_name in enumerate(classes):
-        class_features = grouped_features.get(class_name, [])
-        if not class_features:
-            continue
-
-        class_array = np.asarray(class_features, dtype=np.float32)
-        sample_size = min(target_per_class, max(len(class_array), 1))
-        replace = len(class_array) < sample_size
-        selected_indices = rng.choice(len(class_array), size=sample_size, replace=replace)
-        selected_features = class_array[selected_indices]
-
-        features.append(selected_features)
-        labels.extend([class_index] * sample_size)
-
-    if not features:
-        return np.empty((0, 0), dtype=np.float32), np.empty((0,), dtype=np.int64)
-
-    return np.vstack(features), np.asarray(labels, dtype=np.int64)
-
-
-def main() -> None:
-    MODEL_DIR.mkdir(parents=True, exist_ok=True)
-    image_root = _find_image_root()
-    train_paths, test_paths, classes = _collect_image_paths(image_root)
-    train_grouped = _load_grouped_features(train_paths, classes)
-    test_grouped = _load_grouped_features(test_paths, classes)
-
-    X_train, y_train = _rebalance_training_data(train_grouped, classes, target_per_class=200)
-    X_test, y_test = _rebalance_training_data(test_grouped, classes, target_per_class=200)
-
-    if len(X_train) == 0 or len(X_test) == 0:
-        raise RuntimeError("Could not extract any usable image features from the dataset.")
-
-    epochs = int(sys.argv[1]) if len(sys.argv) > 1 else 1
-    model = ExtraTreesClassifier(
-        n_estimators=max(800, 200 * epochs),
-        max_depth=30,                    # Prevent overfitting on noise
-        min_samples_split=5,              # Require 5+ samples to split
-        min_samples_leaf=2,               # Require 2+ samples at leaf
-        class_weight="balanced",          # Handle class imbalance (especially Audi)
-        random_state=42,
-        n_jobs=-1,
+from src.config import DATA_RAW_DIR, MODEL_DIR
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Train a transfer-learning vehicle brand classifier.")
+    parser.add_argument(
+        "--data_dir",
+        type=str,
+        default="data/raw/Cars Dataset",
+        help="Folder with train/ and test/ subfolders containing class folders.",
+    )
+    parser.add_argument(
+        "--base_model",
+        type=str,
+        default="microsoft/resnet-18",
+        help="Base model identifier from Hugging Face.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="models/car-image-classifier",
+        help="Output directory for trained model.",
+    )
+    parser.add_argument("--epochs", type=int, default=3, help="Number of training epochs.")
+    parser.add_argument("--batch_size", type=int, default=64, help="Batch size for training and evaluation.")
+    parser.add_argument("--learning_rate", type=float, default=5e-5, help="Learning rate.")
+    parser.add_argument("--weight_decay", type=float, default=0.01, help="Weight decay.")
+    parser.add_argument("--warmup_ratio", type=float, default=0.1, help="Warmup ratio.")
+    parser.add_argument("--label_smoothing", type=float, default=0.1, help="Label smoothing factor.")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
+    parser.add_argument("--freeze_backbone", action="store_true", help="Freeze backbone and only train head.")
+    parser.add_argument("--push_to_hub", action="store_true", help="Push model to Hugging Face Hub.")
+    parser.add_argument("--hub_model_id", type=str, default="", help="Hub model ID for pushing.")
+    return parser.parse_args()
+
+
+def build_transforms(processor):
+    """Build training and validation transforms based on processor config."""
+    image_mean = processor.image_mean
+    image_std = processor.image_std
+    size_cfg = processor.size
+
+    # Extract image size from processor config
+    if isinstance(size_cfg, dict):
+        size = size_cfg.get("shortest_edge") or size_cfg.get("height") or size_cfg.get("width") or 224
+    else:
+        size = int(size_cfg) if size_cfg else 224
+
+    from torchvision.transforms import (
+        CenterCrop,
+        ColorJitter,
+        Compose,
+        Normalize,
+        RandomHorizontalFlip,
+        RandomResizedCrop,
+        RandomRotation,
+        Resize,
+        ToTensor,
     )
-    model.fit(X_train, y_train)
 
-    predictions = model.predict(X_test)
-    accuracy = float(accuracy_score(y_test, predictions))
+    train_tfm = Compose(
+        [
+            RandomResizedCrop(size),
+            RandomHorizontalFlip(),
+            RandomRotation(15),
+            ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
+            ToTensor(),
+            Normalize(mean=image_mean, std=image_std),
+        ]
+    )
 
-    # Only include labels that actually appear in the test set to avoid
-    # mismatches between `classes` and present test labels (can happen when
-    # a class has no extracted test samples after filtering).
-    present_label_indices = sorted(np.unique(y_test).tolist())
-    present_class_names = [classes[idx] for idx in present_label_indices]
+    val_tfm = Compose(
+        [
+            Resize(size),
+            CenterCrop(size),
+            ToTensor(),
+            Normalize(mean=image_mean, std=image_std),
+        ]
+    )
 
-    report = classification_report(
-        y_test,
-        predictions,
-        labels=present_label_indices,
-        target_names=present_class_names,
-        output_dict=True,
-        zero_division=0,
+    return train_tfm, val_tfm
+
+
+def main():
+    args = parse_args()
+    set_seed(args.seed)
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+
+    # Load dataset
+    data_dir = Path(args.data_dir)
+    if not data_dir.exists():
+        raise FileNotFoundError(f"Dataset folder not found: {data_dir}")
+
+    print(f"Loading dataset from {data_dir}...")
+    ds = load_dataset("imagefolder", data_dir=str(data_dir))
+
+    # Ensure we have train and test (or validation)
+    if "train" not in ds:
+        raise ValueError("Dataset must have a 'train' split (in train/ folder).")
+
+    if "test" not in ds:
+        if "validation" in ds:
+            ds["test"] = ds["validation"]
+        else:
+            # Create test split if only train exists
+            split = ds["train"].train_test_split(test_size=0.2, seed=42)
+            ds["train"] = split["train"]
+            ds["test"] = split["test"]
+
+    # Get label mapping
+    labels = ds["train"].features["label"].names
+    label2id = {label: i for i, label in enumerate(labels)}
+    id2label = {i: label for i, label in enumerate(labels)}
+
+    print(f"Classes: {labels}")
+    print(f"Number of classes: {len(labels)}")
+    print(f"Training samples: {len(ds['train'])}")
+    print(f"Test samples: {len(ds['test'])}")
+
+    # Load processor and model
+    print(f"Loading base model: {args.base_model}")
+    processor = AutoImageProcessor.from_pretrained(args.base_model)
+    model = AutoModelForImageClassification.from_pretrained(
+        args.base_model,
+        num_labels=len(labels),
+        id2label=id2label,
+        label2id=label2id,
+        ignore_mismatched_sizes=True,
+    )
+
+    # Optionally freeze backbone
+    if args.freeze_backbone:
+        print("Freezing backbone, only training head...")
+        trainable_heads = ("classifier", "score", "fc", "heads", "head")
+        for name, param in model.named_parameters():
+            if not any(head in name for head in trainable_heads):
+                param.requires_grad = False
+
+    # Build transforms
+    train_tfm, val_tfm = build_transforms(processor)
+
+    def transform_train(batch):
+        batch["pixel_values"] = [train_tfm(img.convert("RGB")) for img in batch["image"]]
+        return batch
+
+    def transform_val(batch):
+        batch["pixel_values"] = [val_tfm(img.convert("RGB")) for img in batch["image"]]
+        return batch
+
+    ds["train"].set_transform(transform_train)
+    ds["test"].set_transform(transform_val)
+
+    def collate_fn(batch):
+        import torch
+
+        return {
+            "pixel_values": torch.stack([example["pixel_values"] for example in batch]),
+            "labels": torch.tensor([example["label"] for example in batch]),
+        }
+
+    # Metrics
+    metric = evaluate.load("accuracy")
+
+    def compute_metrics(eval_pred):
+        logits, labels_ = eval_pred
+        predictions = np.argmax(logits, axis=1)
+        return metric.compute(predictions=predictions, references=labels_)
+
+    # Training arguments
+    training_args = TrainingArguments(
+        output_dir=args.output_dir,
+        remove_unused_columns=False,
+        eval_strategy="epoch",
+        save_strategy="epoch",
+        load_best_model_at_end=True,
+        logging_strategy="steps",
+        logging_steps=50,
+        learning_rate=args.learning_rate,
+        weight_decay=args.weight_decay,
+        warmup_ratio=args.warmup_ratio,
+        label_smoothing_factor=args.label_smoothing,
+        per_device_train_batch_size=args.batch_size,
+        per_device_eval_batch_size=args.batch_size,
+        num_train_epochs=args.epochs,
+        push_to_hub=args.push_to_hub,
+        hub_model_id=args.hub_model_id if args.hub_model_id else None,
+        report_to="none",
     )
 
-    matrix = confusion_matrix(y_test, predictions, labels=present_label_indices).tolist()
+    # Trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=ds["train"],
+        eval_dataset=ds["test"],
+        data_collator=collate_fn,
+        compute_metrics=compute_metrics,
+    )
 
-    print(f"Test accuracy: {accuracy:.4f}")
-    print("Per-class recall (test-set classes):")
-    for class_name in present_class_names:
-        print(f"- {class_name}: {report[class_name]['recall']:.4f}")
+    # Train
+    print("Starting training...")
+    trainer.train()
 
-    joblib.dump(model, VISION_MODEL_PATH, compress=3)
-    VISION_LABELS_PATH.write_text(json.dumps(classes, indent=2), encoding="utf-8")
+    # Evaluate
+    print("Evaluating...")
+    metrics = trainer.evaluate()
+    print(f"Test accuracy: {metrics.get('eval_accuracy', 0):.4f}")
 
+    # Save model and processor
+    print(f"Saving model to {args.output_dir}...")
+    trainer.save_model(args.output_dir)
+    processor.save_pretrained(args.output_dir)
+
+    # Generate detailed metrics
+    predictions = trainer.predict(ds["test"])
+    pred_labels = np.argmax(predictions.predictions, axis=1)
+    true_labels = predictions.label_ids
+
+    report = classification_report(
+        true_labels,
+        pred_labels,
+        target_names=labels,
+        output_dict=True,
+        zero_division=0,
+    )
+
+    # Save metadata
     metadata = {
         "created_at": datetime.now(timezone.utc).isoformat(),
-        "data_source": str(image_root),
-        "model_name": "ExtraTreesClassifier",
-        "train_epochs": epochs,
-        "classes": classes,
-        "test_accuracy": round(accuracy, 4),
-        "training_strategy": "balanced_sampling_target_200",
+        "base_model": args.base_model,
+        "number_of_classes": len(labels),
+        "class_names": labels,
+        "train_image_count": len(ds["train"]),
+        "test_image_count": len(ds["test"]),
+        "accuracy": float(metrics.get("eval_accuracy", 0)),
         "classification_report": report,
-        "confusion_matrix": matrix,
+        "note": "This model can only predict one of the trained vehicle brands/classes. It does not provide damage detection or technical condition assessment.",
     }
-    VISION_METADATA_PATH.write_text(json.dumps(metadata, indent=2), encoding="utf-8")
 
-    print(f"Saved model to: {VISION_MODEL_PATH}")
-    print(f"Saved labels to: {VISION_LABELS_PATH}")
-    print(f"Saved metadata to: {VISION_METADATA_PATH}")
+    metadata_path = Path(args.output_dir) / "vision_metadata.json"
+    metadata_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(metadata_path, "w") as f:
+        json.dump(metadata, f, indent=2)
+
+    print(f"Metadata saved to {metadata_path}")
+    print("\nTraining complete!")
+    print(f"Model saved to {args.output_dir}")
+    print(f"Test accuracy: {metrics.get('eval_accuracy', 0):.4f}")
+
+    if args.push_to_hub:
+        trainer.push_to_hub()
 
 
 if __name__ == "__main__":