diff --git a/.dockerignore b/.dockerignore index 0e47b55..cdf1181 100644 --- a/.dockerignore +++ b/.dockerignore @@ -24,6 +24,9 @@ models/* !models/captcha_svm_model.pkl !models/label_encoder.pkl !models/scaler.pkl +!models/efficientnet_v2_best.pth +!models/v2_classes.txt +!models/cnn_best_model.pth !app/ !scripts/ diff --git a/.gitignore b/.gitignore index fdfbd3d..78b5562 100644 --- a/.gitignore +++ b/.gitignore @@ -44,6 +44,9 @@ models/* !models/captcha_svm_model.pkl !models/label_encoder.pkl !models/scaler.pkl +!models/efficientnet_v2_best.pth +!models/v2_classes.txt +!models/cnn_best_model.pth # Project Specific app/uploads/* diff --git a/app/routers/captcha_v2.py b/app/routers/captcha_v2.py index bd2830a..3da63c0 100644 --- a/app/routers/captcha_v2.py +++ b/app/routers/captcha_v2.py @@ -3,6 +3,7 @@ from fastapi import APIRouter, Request from fastapi.responses import JSONResponse from ..services.dataset_service import dataset_cache +from ..services.captcha_v2_service import captcha_v2_service router = APIRouter(prefix="/api/v2") @@ -32,18 +33,42 @@ async def get_challenge(): "grid": grid_with_images } +# ============================ +# Model Integration +# ============================ +from ..services.captcha_v2_service import captcha_v2_service + +@router.on_event("startup") +async def load_v2_model(): + captcha_v2_service.load_model() + @router.post("/solve") async def solve_v2(request: Request): - """Simulate AI solving the 3x3 grid.""" + """Real AI solving the 3x3 grid using EfficientNet.""" data = await request.json() - grid_ids = data.get("grid_ids", []) # IDs of images currently in the grid + grid_ids = data.get("grid_ids", []) target_name = data.get("target") - if target_name not in dataset_cache.v2_data: - return {"correct_indices": []} - - target_files = set(dataset_cache.v2_data[target_name]) - correct_indices = [idx for idx, img_id in enumerate(grid_ids) if img_id in target_files] + correct_indices = [] + + # Reverse lookup map (ID -> Path) + id_to_path = {} + for cat, ids in dataset_cache.v2_data.items(): + for img_id in ids: + id_to_path[img_id] = dataset_cache.v2_root / cat / img_id + + for idx, img_id in enumerate(grid_ids): + if img_id not in id_to_path: + continue + + img_path = id_to_path[img_id] + predicted_class = captcha_v2_service.predict(img_path) + + # Determine match + # Note: predicted_class is singular (e.g. 'bus'), target_name is 'bus'. + # Exact match. + if predicted_class and predicted_class.lower() == target_name.lower(): + correct_indices.append(idx) return {"correct_indices": correct_indices} diff --git a/app/services/__init__.py b/app/services/__init__.py index e69de29..a0dfeef 100644 --- a/app/services/__init__.py +++ b/app/services/__init__.py @@ -0,0 +1,2 @@ +from .dataset_service import dataset_cache +from .captcha_v2_service import captcha_v2_service diff --git a/app/services/captcha_v2_service.py b/app/services/captcha_v2_service.py new file mode 100644 index 0000000..03ea2b2 --- /dev/null +++ b/app/services/captcha_v2_service.py @@ -0,0 +1,115 @@ +import torch +import torch.nn as nn +from torchvision import models, transforms +from PIL import Image +from pathlib import Path +import json + +# Constants +DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +if torch.backends.mps.is_available(): + DEVICE = torch.device('mps') + +# Paths are relative to project root usually, but let's make them robust +BASE_DIR = Path(__file__).parent.parent.parent +MODEL_PATH = BASE_DIR / "models" / "cnn_best_model.pth" +METADATA_PATH = BASE_DIR / "models" / "model_metadata_v2.json" +CLASSES_PATH = BASE_DIR / "models" / "v2_classes.txt" + +class CaptchaV2Service: + _instance = None + + def __new__(cls): + if cls._instance is None: + cls._instance = super(CaptchaV2Service, cls).__new__(cls) + cls._instance._initialized = False + return cls._instance + + def __init__(self): + if self._initialized: return + + self.model = None + self.classes = [] + self.transform = transforms.Compose([ + transforms.Resize((224, 224)), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + + self.load_model() + self._initialized = True + + def load_model(self): + try: + # Load Classes from metadata + if METADATA_PATH.exists(): + with open(METADATA_PATH, "r") as f: + meta = json.load(f) + self.classes = meta.get("classes", []) + print(f"✅ [V2 Service] Loaded {len(self.classes)} classes from metadata") + + if not self.classes: + # Fallback classes (11 classes as detected in model weights) + self.classes = [ + "Bicycle", "Bridge", "Bus", "Car", "Chimney", + "Crosswalk", "Hydrant", "Motorcycle", "Other", "Palm", "Stair" + ] + print(f"⚠️ [V2 Service] Using fallback 11 classes") + + # Load Model: cnn_best_model.pth is EfficientNet-B1 + self.model = models.efficientnet_b1(weights=None) + + # Match the nested classifier structure found in the state_dict (classifier.1.1) + num_ftrs = self.model.classifier[1].in_features + self.model.classifier[1] = nn.Sequential( + nn.Identity(), + nn.Linear(num_ftrs, len(self.classes)) + ) + + if MODEL_PATH.exists(): + state_dict = torch.load(MODEL_PATH, map_location=DEVICE) + # Handle 'model.' prefix if present + new_state_dict = {} + for k, v in state_dict.items(): + if k.startswith("model."): + new_state_dict[k[6:]] = v + else: + new_state_dict[k] = v + + self.model.load_state_dict(new_state_dict) + self.model.to(DEVICE) + self.model.eval() + print(f"✅ [V2 Service] EfficientNet-B1 Model loaded from {MODEL_PATH}") + else: + print(f"⚠️ [V2 Service] Model file not found at {MODEL_PATH}") + self.model = None + + except Exception as e: + print(f"❌ [V2 Service] Failed to load model: {e}") + self.model = None + + def predict(self, img_path: Path): + """ + Predicts the class of the image at img_path. + Returns the class name (str) or None if prediction fails. + """ + if not self.model: + # Try reloading if missing (e.g. usage before startup) + self.load_model() + if not self.model: return None + + try: + img = Image.open(img_path).convert('RGB') + img_tensor = self.transform(img).unsqueeze(0).to(DEVICE) + + with torch.no_grad(): + outputs = self.model(img_tensor) + _, pred_idx = torch.max(outputs, 1) + + return self.classes[pred_idx.item()] + except Exception as e: + print(f"❌ [V2 Service] Prediction error for {img_path}: {e}") + return None + +# Singleton export +captcha_v2_service = CaptchaV2Service() diff --git a/app/services/metadata_service.py b/app/services/metadata_service.py index d9d2148..c917e70 100644 --- a/app/services/metadata_service.py +++ b/app/services/metadata_service.py @@ -26,7 +26,29 @@ def load_metadata(): with open(v2_path, "r") as f: data = json.load(f) STATE["v2_metadata"].update(data) - print(f"✅ V2 Metadata loaded from {v2_path.name}") + + # Parse Stats for Dashboard + if "fold_results" in data: + best_acc = 0.0 + avg_loss = 0.0 + total_folds = len(data["fold_results"]) + + for fold in data["fold_results"]: + if fold["best_val_acc"] > best_acc: + best_acc = fold["best_val_acc"] + + # Take last validation loss as proxy + if "history" in fold and "val_loss" in fold["history"]: + avg_loss += fold["history"]["val_loss"][-1] + + if total_folds > 0: + avg_loss /= total_folds + + STATE["v2_metadata"]["accuracy"] = f"{best_acc*100:.2f}%" + STATE["v2_metadata"]["loss"] = f"{avg_loss:.4f}" + STATE["v2_metadata"]["type"] = "EfficientNet-B0 (K-Fold)" + + print(f"✅ V2 Metadata loaded from {v2_path.name} (Acc: {STATE['v2_metadata']['accuracy']})") except Exception as e: print(f"❌ Failed to load V2 metadata: {e}") diff --git a/models/cnn_best_model.pth b/models/cnn_best_model.pth new file mode 100644 index 0000000..3f075c9 Binary files /dev/null and b/models/cnn_best_model.pth differ diff --git a/models/efficientnet_v2_best.pth b/models/efficientnet_v2_best.pth new file mode 100644 index 0000000..4706247 Binary files /dev/null and b/models/efficientnet_v2_best.pth differ diff --git a/models/model_metadata_v2.json b/models/model_metadata_v2.json index 00c05f1..f84eba4 100644 --- a/models/model_metadata_v2.json +++ b/models/model_metadata_v2.json @@ -1,12 +1,25 @@ { - "accuracy": "12.34%", - "char_accuracy": "12.34%", - "precision": "0.5", - "recall": "0.5", - "f1_score": "0.5", - "loss_value": "0.1", - "type": "CNN", - "loss": "CrossEntropy", - "features": "Image Classification", - "preprocessing": "Resize (28x28), Normalization" -} + "accuracy": "72.98%", + "char_accuracy": "72.98%", + "loss": "1.1722", + "loss_value": "1.1722", + "precision": "0.67", + "recall": "0.73", + "f1_score": "0.70", + "type": "EfficientNet-B0 (Re-eval)", + "timestamp": "2026-01-28T01:00:10", + "classes": [ + "Bicycle", + "Bridge", + "Bus", + "Car", + "Chimney", + "Crosswalk", + "Hydrant", + "Motorcycle", + "Other", + "Palm", + "Stair" + ], + "dataset_path": "/Users/parkyoungdu/Downloads/samples_v2/images" +} \ No newline at end of file diff --git a/models/v2_classes.txt b/models/v2_classes.txt new file mode 100644 index 0000000..d45d85d --- /dev/null +++ b/models/v2_classes.txt @@ -0,0 +1,12 @@ +Bicycle +Bridge +Bus +Car +Chimney +Crosswalk +Hydrant +Motorcycle +Other +Palm +Stair +Traffic Light \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index a461e26..c4bbd20 100644 --- a/requirements.txt +++ b/requirements.txt @@ -17,3 +17,7 @@ websockets # ML Production (CPU optimized) torch --index-url https://download.pytorch.org/whl/cpu torchinfo +torchvision +mlflow +opencv-python +accelerate diff --git a/scripts/generate_v2_metadata.py b/scripts/generate_v2_metadata.py new file mode 100644 index 0000000..58f34f0 --- /dev/null +++ b/scripts/generate_v2_metadata.py @@ -0,0 +1,173 @@ +import os +import json +import argparse +import time +from pathlib import Path +import torch +import torch.nn as nn +from torchvision import models, transforms, datasets +from torch.utils.data import DataLoader +from sklearn.metrics import precision_recall_fscore_support +from tqdm import tqdm + +# Constants +DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +if torch.backends.mps.is_available(): + DEVICE = torch.device('mps') + +# Paths +BASE_DIR = Path(__file__).resolve().parent.parent +MODELS_DIR = BASE_DIR / "models" +MODEL_PATH = MODELS_DIR / "cnn_best_model.pth" +CLASSES_PATH = MODELS_DIR / "v2_classes.txt" +OUTPUT_PATH = MODELS_DIR / "model_metadata_v2.json" + +# Import app config to get correct dataset path +import sys +sys.path.append(str(BASE_DIR)) +from app.core.config import get_v2_dataset_path + +DATA_DIR = Path(get_v2_dataset_path()) + +def load_classes(): + if not CLASSES_PATH.exists(): + raise FileNotFoundError(f"Classes file not found at {CLASSES_PATH}") + with open(CLASSES_PATH, "r") as f: + return [line.strip() for line in f.readlines()] + +def load_model(num_classes): + if not MODEL_PATH.exists(): + raise FileNotFoundError(f"Model file not found at {MODEL_PATH}") + + print(f"Loading model from {MODEL_PATH}...") + # It's EfficientNet-B1 based on block structure + model = models.efficientnet_b1(weights=None) + + # Match the nested classifier structure found in the state_dict (classifier.1.1) + num_ftrs = model.classifier[1].in_features + model.classifier[1] = nn.Sequential( + nn.Identity(), # 1.0 + nn.Linear(num_ftrs, num_classes) # 1.1 + ) + + state_dict = torch.load(MODEL_PATH, map_location=DEVICE) + + # Handle 'model.' prefix if present + new_state_dict = {} + for k, v in state_dict.items(): + if k.startswith("model."): + new_state_dict[k[6:]] = v + else: + new_state_dict[k] = v + + print(f"Sample keys in new_state_dict: {list(new_state_dict.keys())[:5]}") + model.load_state_dict(new_state_dict) + model.to(DEVICE) + model.eval() + return model + +def evaluate(model, classes): + if not DATA_DIR.exists(): + raise FileNotFoundError(f"Dataset directory not found at {DATA_DIR}") + + print(f"Evaluating on dataset at {DATA_DIR}...") + transform = transforms.Compose([ + transforms.Resize((224, 224)), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + + dataset = datasets.ImageFolder(str(DATA_DIR), transform=transform) + dataloader = DataLoader(dataset, batch_size=32, shuffle=False, num_workers=0) + + criterion = nn.CrossEntropyLoss() + + running_loss = 0.0 + all_preds = [] + all_labels = [] + + with torch.no_grad(): + for inputs, labels in tqdm(dataloader, desc="Evaluating"): + inputs = inputs.to(DEVICE) + labels = labels.to(DEVICE) + + outputs = model(inputs) + loss = criterion(outputs, labels) + + running_loss += loss.item() * inputs.size(0) + _, preds = torch.max(outputs, 1) + + all_preds.extend(preds.cpu().numpy()) + all_labels.extend(labels.cpu().numpy()) + + total_loss = running_loss / len(dataset) + + # Calculate Metrics + correct = sum(p == l for p, l in zip(all_preds, all_labels)) + accuracy = correct / len(dataset) + + precision, recall, f1, _ = precision_recall_fscore_support(all_labels, all_preds, average='weighted', zero_division=0) + + return { + "loss": total_loss, + "accuracy": accuracy, + "precision": precision, + "recall": recall, + "f1_score": f1 + } + +def main(): + try: + print(f"Using device: {DEVICE}") + + # Determine actual num_classes from model state_dict first + state_dict = torch.load(MODEL_PATH, map_location=DEVICE) + actual_num_classes = 0 + for k, v in state_dict.items(): + if 'classifier.1.1.weight' in k or 'classifier.1.weight' in k: + actual_num_classes = v.size(0) + break + + if actual_num_classes == 0: + # Fallback to loading from classes file + classes = load_classes() + actual_num_classes = len(classes) + else: + print(f"Detected {actual_num_classes} classes from model weights.") + # We still need class names for metadata, but we'll handle the mismatch if any + classes = load_classes() + if len(classes) != actual_num_classes: + print(f"⚠️ Warning: Classes file has {len(classes)} classes, but model has {actual_num_classes}.") + classes = classes[:actual_num_classes] # Simple truncation for now + + model = load_model(actual_num_classes) + + metrics = evaluate(model, classes) + + # Format Metadata + metadata = { + "accuracy": f"{metrics['accuracy']*100:.2f}%", + "char_accuracy": f"{metrics['accuracy']*100:.2f}%", # Reuse for consistent UI if needed, or omit + "loss": f"{metrics['loss']:.4f}", # Display string + "loss_value": f"{metrics['loss']:.4f}", # Raw value string + "precision": f"{metrics['precision']:.2f}", + "recall": f"{metrics['recall']:.2f}", + "f1_score": f"{metrics['f1_score']:.2f}", + "type": "EfficientNet-B0 (Re-eval)", + "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"), + "classes": classes, + "dataset_path": str(DATA_DIR) + } + + # Save + with open(OUTPUT_PATH, "w") as f: + json.dump(metadata, f, indent=4) + + print(f"✅ Metadata saved to {OUTPUT_PATH}") + print(json.dumps(metadata, indent=4)) + + except Exception as e: + print(f"❌ Error: {e}") + +if __name__ == "__main__": + main() diff --git a/scripts/train_v2_classifier.py b/scripts/train_v2_classifier.py new file mode 100644 index 0000000..313403c --- /dev/null +++ b/scripts/train_v2_classifier.py @@ -0,0 +1,250 @@ +import os +import argparse +import shutil +import time +from pathlib import Path +import numpy as np +import matplotlib.pyplot as plt +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader, WeightedRandomSampler +from torchvision import datasets, models, transforms +from sklearn.model_selection import train_test_split +from tqdm import tqdm +import mlflow +import mlflow.pytorch +import kagglehub + +# ============================ +# 1. Configuration & Constants +# ============================ +DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +# Force CPU on Mac IF MPS is unstable, but usually MPS is fine for CNNs. +# Safe bet is CPU or MPS. efficientnet might differ. +if torch.backends.mps.is_available(): + DEVICE = torch.device('mps') + print("Using MPS (Metal Performance Shaders) acceleration") +else: + print(f"Using device: {DEVICE}") + +IMG_SIZE = 224 +BATCH_SIZE = 32 +LEARNING_RATE = 0.001 + +# ============================ +# 2. Data Preparation +# ============================ +def prepare_data(data_root="data"): + """ + Ensures data exists at data_root/samples_v2/images. + Downloads from Kaggle if missing. + """ + v2_dest = Path(data_root) / "samples_v2" / "images" + + if v2_dest.exists() and any(v2_dest.iterdir()): + print(f"✅ Dataset found at {v2_dest}") + return v2_dest + + print(f"Dataset not found at {v2_dest}. Downloading...") + try: + # Download from Kaggle + path = kagglehub.dataset_download("mikhailma/test-dataset") + print(f"Downloaded to cache: {path}") + + # Copy to local project data dir for structure + v2_dest.parent.mkdir(parents=True, exist_ok=True) + if v2_dest.exists(): + shutil.rmtree(v2_dest) + + # The dataset structure from kaggle might be specific, let's copy the content + # Check if 'path' contains folders directly or an 'images' folder + # For 'mikhailma/test-dataset', based on download_datasets.py logic, it seems direct. + shutil.copytree(path, v2_dest) + print(f"✅ Copied to {v2_dest}") + return v2_dest + except Exception as e: + print(f"❌ Failed to download dataset: {e}") + # Fallback to local if running in Docker/Expected Env + return v2_dest + +def perform_eda(dataset_path): + """ + Prints class distribution and checks for corrupt images. + """ + print("\n--- EDA: Dataset Statistics ---") + classes = sorted([d.name for d in dataset_path.iterdir() if d.is_dir()]) + print(f"Classes ({len(classes)}): {classes}") + + class_counts = {} + total_images = 0 + + for cls in classes: + cls_dir = dataset_path / cls + count = len(list(cls_dir.glob("*"))) + class_counts[cls] = count + total_images += count + + print(f"Total Images: {total_images}") + print("Class Distribution:") + for cls, count in class_counts.items(): + ratio = count / total_images + print(f" - {cls}: {count} ({ratio:.1%})") + + return classes, class_counts + +# ============================ +# 3. Training Logic +# ============================ +def train_model(args): + dataset_path = prepare_data() + classes, class_counts = perform_eda(dataset_path) + + # Transforms + # Helper to clean/check images? ImageFolder deals with basic loading. + + train_transform = transforms.Compose([ + transforms.Resize((IMG_SIZE, IMG_SIZE)), + transforms.RandomHorizontalFlip(), + transforms.RandomRotation(10), + transforms.ColorJitter(brightness=0.2, contrast=0.2), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + + val_transform = transforms.Compose([ + transforms.Resize((IMG_SIZE, IMG_SIZE)), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + + # Dataset + full_dataset = datasets.ImageFolder(root=str(dataset_path)) + + # Train/Val Split (indices) + train_idx, val_idx = train_test_split( + list(range(len(full_dataset))), + test_size=0.2, + stratify=full_dataset.targets, + random_state=42 + ) + + # Subsets with different transforms + train_dataset = torch.utils.data.Subset(full_dataset, train_idx) + # Hack to apply different transform to subset? + # Proper way is creating two Datasets or Wrapper. + # For simplicity, we'll apply transform in the Dataset but ImageFolder applies ONE transform. + # Pattern: Re-init ImageFolder with transform, or custom wrapper. + # Let's use two ImageFolders for simplicity (Valid since data is same file structure) + + train_ds = datasets.ImageFolder(str(dataset_path), transform=train_transform) + val_ds = datasets.ImageFolder(str(dataset_path), transform=val_transform) + + train_subset = torch.utils.data.Subset(train_ds, train_idx) + val_subset = torch.utils.data.Subset(val_ds, val_idx) + + train_loader = DataLoader(train_subset, batch_size=args.batch_size, shuffle=True, num_workers=0) + val_loader = DataLoader(val_subset, batch_size=args.batch_size, shuffle=False, num_workers=0) + + # Model Setup + print(f"\nInitializing EfficientNet-B0 for {len(classes)} classes...") + model = models.efficientnet_b0(weights='IMAGENET1K_V1') + + # Replace Head + num_ftrs = model.classifier[1].in_features + model.classifier[1] = nn.Linear(num_ftrs, len(classes)) + + model = model.to(DEVICE) + + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=args.cv_lr) + + # MLflow + mlflow.set_experiment("Captcha_V2_EfficientNet") + + with mlflow.start_run(): + mlflow.log_params({ + "model": "EfficientNet-B0", + "epochs": args.epochs, + "batch_size": args.batch_size, + "learning_rate": args.cv_lr + }) + + best_acc = 0.0 + + for epoch in range(args.epochs): + print(f"Epoch {epoch+1}/{args.epochs}") + + # Train + model.train() + running_loss = 0.0 + correct = 0 + total = 0 + + for inputs, labels in tqdm(train_loader, desc="Training"): + inputs, labels = inputs.to(DEVICE), labels.to(DEVICE) + + optimizer.zero_grad() + outputs = model(inputs) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() * inputs.size(0) + _, predicted = torch.max(outputs, 1) + total += labels.size(0) + correct += (predicted == labels).sum().item() + + epoch_loss = running_loss / len(train_subset) + epoch_acc = correct / total + + # Val + model.eval() + val_loss = 0.0 + val_correct = 0 + val_total = 0 + + with torch.no_grad(): + for inputs, labels in tqdm(val_loader, desc="Validation"): + inputs, labels = inputs.to(DEVICE), labels.to(DEVICE) + outputs = model(inputs) + loss = criterion(outputs, labels) + + val_loss += loss.item() * inputs.size(0) + _, predicted = torch.max(outputs, 1) + val_total += labels.size(0) + val_correct += (predicted == labels).sum().item() + + val_epoch_loss = val_loss / len(val_subset) + val_epoch_acc = val_correct / val_total + + print(f" Train Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}") + print(f" Val Loss: {val_epoch_loss:.4f} Acc: {val_epoch_acc:.4f}") + + mlflow.log_metric("train_loss", epoch_loss, step=epoch) + mlflow.log_metric("train_acc", epoch_acc, step=epoch) + mlflow.log_metric("val_loss", val_epoch_loss, step=epoch) + mlflow.log_metric("val_acc", val_epoch_acc, step=epoch) + + # Save Best Model + if val_epoch_acc > best_acc: + best_acc = val_epoch_acc + os.makedirs("models", exist_ok=True) + torch.save(model.state_dict(), "models/efficientnet_v2_best.pth") + print(f" 🌟 New Best Model Saved! ({best_acc:.4f})") + + # Save Final Metadata + print("Training Complete.") + # Log classes to artifacts for inference mapping + with open("models/v2_classes.txt", "w") as f: + f.write("\n".join(classes)) + mlflow.log_artifact("models/v2_classes.txt") + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--epochs", type=int, default=10) + parser.add_argument("--batch_size", type=int, default=32) + parser.add_argument("--cv_lr", type=float, default=0.001) + args = parser.parse_args() + + train_model(args) diff --git a/scripts/train_v2_kfold.py b/scripts/train_v2_kfold.py new file mode 100644 index 0000000..6990a57 --- /dev/null +++ b/scripts/train_v2_kfold.py @@ -0,0 +1,288 @@ +import os +import argparse +import json +import time +import sys +from pathlib import Path +import numpy as np +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader, Subset +from torchvision import datasets, models, transforms +from sklearn.model_selection import KFold +from tqdm import tqdm +import mlflow +import mlflow.pytorch +import matplotlib.pyplot as plt + +# Add project root to path to import config +sys.path.append(str(Path(__file__).resolve().parent.parent)) +try: + from app.core.config import get_v2_dataset_path +except ImportError: + # Fallback if running outside expected structure + def get_v2_dataset_path(): + return "data/samples_v2" + +# ============================ +# 1. Configuration & Constants +# ============================ +DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +if torch.backends.mps.is_available(): + DEVICE = torch.device('mps') + +IMG_SIZE = 224 + +class EarlyStopping: + def __init__(self, patience=10, delta=0): + self.patience = patience + self.delta = delta + self.counter = 0 + self.best_loss = None + self.early_stop = False + + def __call__(self, val_loss): + if self.best_loss is None: + self.best_loss = val_loss + elif val_loss > self.best_loss - self.delta: + self.counter += 1 + if self.counter >= self.patience: + self.early_stop = True + self.best_loss = val_loss + self.counter = 0 + +def plot_history(metadata): + if not metadata or "fold_results" not in metadata: + return + + folds = metadata["fold_results"] + plt.figure(figsize=(15, 6)) + + # Plot Loss + plt.subplot(1, 2, 1) + for f in folds: + plt.plot(f['history']['train_loss'], alpha=0.3, label=f"Fold {f['fold']} Train") + plt.plot(f['history']['val_loss'], label=f"Fold {f['fold']} Val") + plt.title("Loss per Fold") + plt.xlabel("Epoch") + plt.ylabel("Loss") + plt.legend() + + # Plot Accuracy + plt.subplot(1, 2, 2) + for f in folds: + plt.plot(f['history']['train_acc'], alpha=0.3, label=f"Fold {f['fold']} Train") + plt.plot(f['history']['val_acc'], label=f"Fold {f['fold']} Val") + plt.title("Accuracy per Fold") + plt.xlabel("Epoch") + plt.ylabel("Accuracy") + plt.legend() + + os.makedirs("docs/visualizations", exist_ok=True) + save_path = "docs/visualizations/v2_training_history.png" + plt.savefig(save_path) + print(f"📊 Training history saved to {save_path}") + +def train_kfold(args): + print(f"Using device: {DEVICE}") + + # 1. Prepare Data + data_path_str = get_v2_dataset_path() + if not data_path_str or not os.path.exists(data_path_str): + print(f"❌ Dataset not found. Please ensure images are available.") + # Try to use download logic if needed or just fail + return + + print(f"✅ Loading dataset from: {data_path_str}") + dataset_path = Path(data_path_str) + + # Basic Transforms + data_transforms = { + 'train': transforms.Compose([ + transforms.Resize((IMG_SIZE, IMG_SIZE)), + transforms.RandomHorizontalFlip(), + transforms.RandomRotation(10), + transforms.ColorJitter(brightness=0.2, contrast=0.2), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]), + 'val': transforms.Compose([ + transforms.Resize((IMG_SIZE, IMG_SIZE)), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]), + } + + # Load Full Dataset (We will apply transforms via custom wrapper or subset logic) + # Since ImageFolder applies one transform, we cheat slightly by reloading or using a wrapper. + # Simple approach: Load ImageFolder with 'val' transform for underlying data, + # but that's wrong for training. + # Proper way for K-Fold: Indicies. + # We will instantiate two ImageFolders, one for train (augmented) and one for val (clean), + # pointing to the SAME data directory. Then use Subsets. + + full_train_dataset = datasets.ImageFolder(data_path_str, transform=data_transforms['train']) + full_val_dataset = datasets.ImageFolder(data_path_str, transform=data_transforms['val']) + + classes = full_train_dataset.classes + print(f"Classes: {classes}") + + # Export Metadata early + metadata = { + "classes": classes, + "num_classes": len(classes), + "config": vars(args), + "dataset_path": data_path_str, + "fold_results": [] + } + + # 2. K-Fold Setup + kfold = KFold(n_splits=args.folds, shuffle=True, random_state=42) + + # MLflow + mlflow.set_experiment("Captcha_V2_KFold") + + best_overall_acc = 0.0 + + # Indices for the whole dataset + indices = np.arange(len(full_train_dataset)) + + for fold, (train_idx, val_idx) in enumerate(kfold.split(indices)): + print(f"\n--- Fold {fold+1}/{args.folds} ---") + + # Subsets + train_subset = Subset(full_train_dataset, train_idx) + val_subset = Subset(full_val_dataset, val_idx) + + train_loader = DataLoader(train_subset, batch_size=args.batch_size, shuffle=True, num_workers=0) + val_loader = DataLoader(val_subset, batch_size=args.batch_size, shuffle=False, num_workers=0) + + # Init Model + model = models.efficientnet_b0(weights='IMAGENET1K_V1') + num_ftrs = model.classifier[1].in_features + model.classifier[1] = nn.Linear(num_ftrs, len(classes)) + model = model.to(DEVICE) + + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=args.learning_rate) + + early_stopping = EarlyStopping(patience=10, delta=0.001) + + fold_best_acc = 0.0 + history = {'train_loss': [], 'val_loss': [], 'train_acc': [], 'val_acc': []} + + with mlflow.start_run(run_name=f"Fold_{fold+1}"): + mlflow.log_params(vars(args)) + mlflow.log_param("fold", fold+1) + + for epoch in range(args.epochs): + # Train + model.train() + running_loss = 0.0 + correct = 0 + total = 0 + + # Progress bar for Training + pbar_train = tqdm(train_loader, desc=f"Epoch {epoch+1}/{args.epochs} [Train]", leave=False) + for inputs, labels in pbar_train: + inputs, labels = inputs.to(DEVICE), labels.to(DEVICE) + optimizer.zero_grad() + outputs = model(inputs) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() * inputs.size(0) + _, predicted = torch.max(outputs, 1) + total += labels.size(0) + correct += (predicted == labels).sum().item() + + # Update pbar description with current loss + pbar_train.set_postfix({'loss': loss.item()}) + + + train_loss = running_loss / len(train_subset) + train_acc = correct / total + + # Val + model.eval() + val_loss_accum = 0.0 + val_correct = 0 + val_total = 0 + + # Progress bar for Validation + with torch.no_grad(): + for inputs, labels in tqdm(val_loader, desc=f"Epoch {epoch+1}/{args.epochs} [Val]", leave=False): + inputs, labels = inputs.to(DEVICE), labels.to(DEVICE) + outputs = model(inputs) + loss = criterion(outputs, labels) + + val_loss_accum += loss.item() * inputs.size(0) + _, predicted = torch.max(outputs, 1) + val_total += labels.size(0) + val_correct += (predicted == labels).sum().item() + + val_loss = val_loss_accum / len(val_subset) + val_acc = val_correct / val_total + + # History + history['train_loss'].append(train_loss) + history['val_loss'].append(val_loss) + history['train_acc'].append(train_acc) + history['val_acc'].append(val_acc) + + print(f"Epoch {epoch+1}/{args.epochs} | Loss: {train_loss:.4f}/{val_loss:.4f} | Acc: {train_acc:.4f}/{val_acc:.4f}") + + mlflow.log_metric("train_loss", train_loss, step=epoch) + mlflow.log_metric("val_loss", val_loss, step=epoch) + mlflow.log_metric("val_acc", val_acc, step=epoch) + + # Updated Best Model Logic (Per Fold and Overall) + if val_acc > fold_best_acc: + fold_best_acc = val_acc + + if val_acc > best_overall_acc: + best_overall_acc = val_acc + os.makedirs("models", exist_ok=True) + torch.save(model.state_dict(), "models/efficientnet_v2_best.pth") + print(f" 🌟 New Overall Best Model! ({val_acc:.4f})") + + # Early Stopping + early_stopping(val_loss) + if early_stopping.early_stop: + print(f" ⏹️ Early stopping at epoch {epoch+1}") + break + + # Log Fold Result + metadata["fold_results"].append({ + "fold": fold + 1, + "best_val_acc": fold_best_acc, + "epochs_run": len(history['train_loss']), + "history": history + }) + + # Save Metadata + print("\nTraining Complete.") + print(f"Best Overall Accuracy: {best_overall_acc:.4f}") + + with open("models/model_metadata_v2.json", "w") as f: + json.dump(metadata, f, indent=4) + print("MetaData saved to models/model_metadata_v2.json") + + # Save Class List separately for simpler loading + with open("models/v2_classes.txt", "w") as f: + f.write("\n".join(classes)) + + # Plot History + plot_history(metadata) + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--epochs", type=int, default=50) # Recommendation: 50 with EarlyStop + parser.add_argument("--batch_size", type=int, default=32) + parser.add_argument("--learning_rate", type=float, default=0.001) + parser.add_argument("--folds", type=int, default=5) + args = parser.parse_args() + + train_kfold(args) diff --git a/scripts/visualize_v2_results.py b/scripts/visualize_v2_results.py new file mode 100644 index 0000000..280c825 --- /dev/null +++ b/scripts/visualize_v2_results.py @@ -0,0 +1,200 @@ +import os +import json +import argparse +import numpy as np +import matplotlib.pyplot as plt +import cv2 +import torch +import torch.nn as nn +from torchvision import models, transforms, datasets +from PIL import Image + +# ============================ +# 1. Configuration +# ============================ +DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +if torch.backends.mps.is_available(): + DEVICE = torch.device('mps') + +# ============================ +# 2. Grad-CAM Utils +# ============================ +class GradCAM: + def __init__(self, model, target_layer): + self.model = model + self.target_layer = target_layer + self.gradients = None + + # Hooks + target_layer.register_forward_hook(self.save_activation) + target_layer.register_full_backward_hook(self.save_gradient) + + def save_activation(self, module, input, output): + self.activation = output + + def save_gradient(self, module, grad_input, grad_output): + self.gradients = grad_output[0] + + def __call__(self, x, class_idx=None): + # Forward pass + output = self.model(x) + if class_idx is None: + class_idx = output.argmax(dim=1).item() + + # Backward pass + self.model.zero_grad() + score = output[0, class_idx] + score.backward() + + # Generate CAM + gradients = self.gradients[0].cpu().data.numpy() + activations = self.activation[0].cpu().data.numpy() + + weights = np.mean(gradients, axis=(1, 2)) + cam = np.zeros(activations.shape[1:], dtype=np.float32) + + for i, w in enumerate(weights): + cam += w * activations[i] + + cam = np.maximum(cam, 0) + cam = cv2.resize(cam, (224, 224)) + cam = cam - np.min(cam) + cam = cam / np.max(cam) + return cam, class_idx, torch.softmax(output, dim=1).max().item() + +# ============================ +# 3. Main Logic +# ============================ +def load_metadata(path="models/model_metadata_v2.json"): + if not os.path.exists(path): + print("Metadata not found. Plotting skipped.") + return None + with open(path, "r") as f: + return json.load(f) + +def plot_history(metadata): + if not metadata or "fold_results" not in metadata: + return + + folds = metadata["fold_results"] + plt.figure(figsize=(15, 6)) + + # Plot Loss + plt.subplot(1, 2, 1) + for f in folds: + plt.plot(f['history']['train_loss'], alpha=0.3, label=f"Fold {f['fold']} Train") + plt.plot(f['history']['val_loss'], label=f"Fold {f['fold']} Val") + plt.title("Loss per Fold") + plt.xlabel("Epoch") + plt.ylabel("Loss") + plt.legend() + + # Plot Accuracy + plt.subplot(1, 2, 2) + for f in folds: + plt.plot(f['history']['train_acc'], alpha=0.3, label=f"Fold {f['fold']} Train") + plt.plot(f['history']['val_acc'], label=f"Fold {f['fold']} Val") + plt.title("Accuracy per Fold") + plt.xlabel("Epoch") + plt.ylabel("Accuracy") + plt.legend() + + os.makedirs("docs/visualizations", exist_ok=True) + plt.savefig("docs/visualizations/v2_training_history.png") + print("Saved training history to docs/visualizations/v2_training_history.png") + +def visualize_heatmaps(model, classes, data_root, num_samples=3): + img_paths = [] + # Collect random images + for root, _, files in os.walk(data_root): + for f in files: + if f.lower().endswith(('.png', '.jpg')): + img_paths.append(os.path.join(root, f)) + + if not img_paths: + return + + selected = np.random.choice(img_paths, num_samples, replace=False) + + # Target Layer: Last Conv layer of EfficientNet features + # EfficientNet parts: features (Sequential) -> avgpool -> classifier + target_layer = model.features[-1] + grad_cam = GradCAM(model, target_layer) + + transform = transforms.Compose([ + transforms.Resize((224, 224)), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + + plt.figure(figsize=(15, 5 * num_samples)) + + for i, img_path in enumerate(selected): + img_pil = Image.open(img_path).convert('RGB') + img_tensor = transform(img_pil).unsqueeze(0).to(DEVICE) + + # Run Grad-CAM + cam, pred_idx, conf = grad_cam(img_tensor) + + # Prepare visualization + img_np = np.array(img_pil.resize((224, 224))) + heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_JET) + heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB) + + overlay = cv2.addWeighted(img_np, 0.6, heatmap, 0.4, 0) + + true_label = Path(img_path).parent.name + pred_label = classes[pred_idx] + + # Plot + plt.subplot(num_samples, 3, i*3 + 1) + plt.imshow(img_np) + plt.title(f"Original: {true_label}") + plt.axis('off') + + plt.subplot(num_samples, 3, i*3 + 2) + plt.imshow(heatmap) + plt.title("Grad-CAM Heatmap") + plt.axis('off') + + plt.subplot(num_samples, 3, i*3 + 3) + plt.imshow(overlay) + plt.title(f"Pred: {pred_label} ({conf:.2f})") + plt.axis('off') + + plt.tight_layout() + plt.savefig("docs/visualizations/v2_gradcam.png") + print("Saved Grad-CAM heatmaps to docs/visualizations/v2_gradcam.png") + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--model_path", default="models/efficientnet_v2_best.pth") + # Try to load path from metadata, if not default + parser.add_argument("--data_root", default=None) + args = parser.parse_args() + + metadata = load_metadata() + if metadata: + plot_history(metadata) + data_root = args.data_root if args.data_root else metadata.get("dataset_path", "data/samples_v2") + else: + data_root = args.data_root or "data/samples_v2" + + if not os.path.exists(args.model_path): + print("Model not found. Train first.") + exit() + + # Load Model + classes = load_metadata()["classes"] if metadata else [] + if not classes and os.path.exists("models/v2_classes.txt"): + with open("models/v2_classes.txt") as f: + classes = [l.strip() for l in f] + + print(f"Loading model with {len(classes)} classes...") + model = models.efficientnet_b0(weights=None) + model.classifier[1] = nn.Linear(model.classifier[1].in_features, len(classes)) + model.load_state_dict(torch.load(args.model_path, map_location=DEVICE)) + model.to(DEVICE) + model.eval() + + visualize_heatmaps(model, classes, data_root)