52_train_classifier.py

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models, transforms
from torch.utils.data import DataLoader, Dataset
import shutil
import logging
from tqdm import tqdm
import pandas as pd
from PIL import Image
import re
import argparse
import json
import sys

# --- Default Configuration ---
DEFAULT_RUNS_SEGMENT_DIR = "runs/segment"
DEFAULT_ALL_SPECIES_RUN_DIR = os.path.join(DEFAULT_RUNS_SEGMENT_DIR, "all_species")
DEFAULT_PER_SPECIES_RUN_DIR = os.path.join(DEFAULT_RUNS_SEGMENT_DIR, "species")
DEFAULT_LABELED_CUTOUTS_DIR = "inference_out/cutouts_masked"
DEFAULT_ANNOTATION_FILE = "inference_out/cutouts_masked/annotations.csv"
DEFAULT_RESULTS_CSV_PATH = "inference_out/numeric/results.csv"
DEFAULT_LOG_FILE = "master_log.txt"

# Sample data paths for validation
SAMPLE_DATA_BASE_DIR = "models/pretrained_model_data/classification/training_input_data"
SAMPLE_LABELED_CUTOUTS_DIR = os.path.join(SAMPLE_DATA_BASE_DIR, "images")
SAMPLE_ANNOTATION_FILE =  os.path.join(SAMPLE_DATA_BASE_DIR, "annotations.csv")

# --- Hyperparameters ---
DEFAULT_NUM_EPOCHS = 300
DEFAULT_BATCH_SIZE = 1024
DEFAULT_LEARNING_RATE = 0.001

def setup_logging(log_file):
    """Configures logging to output to both console and file."""
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)

    if logger.hasHandlers():
        logger.handlers.clear()

    # File handler for master_log.txt
    file_handler = logging.FileHandler(log_file, mode='a', encoding="utf-8")
    file_handler.setLevel(logging.INFO)
    file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    file_handler.setFormatter(file_formatter)
    logger.addHandler(file_handler)

    # Console handler for INFO and above
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setLevel(logging.INFO)
    console_formatter = logging.Formatter('%(levelname)s: %(message)s')
    console_handler.setFormatter(console_formatter)
    logger.addHandler(console_handler)

def build_label_groups(annotations_df, target_columns=None):
    """
    Dynamically build LABEL_GROUPS from the annotation DataFrame.
    Rules:
      - Exclude rows with 'skipped' in that column when collecting classes.
      - If a label starting with 'unknown_' exists OR label == 'not_visible', pick one as ignore label (key '0').
      - Remaining labels get keys '1','2',... (stable alphabetical order).
      - If no ignore label existes, only numbered keys ('1','2',...) are created and no '0' key.
    """
    if target_columns is None:
        # All columns except filename (and any added columns like species)
        target_columns = [c for c in annotations_df.columns if c not in ('filename', 'species')]
    groups = {}
    for col in target_columns:
        if col == 'filename':
            continue
        if col not in annotations_df.columns:
            continue
        vals = annotations_df[col].dropna()
        vals = [v for v in vals if v != 'skipped']
        if not vals:
            continue
        unique = sorted(set(vals))
        # Determine ignore candidate
        ignore_candidates = [v for v in unique if v.startswith('unknown_')]
        if not ignore_candidates and 'not_visible' in unique:
            ignore_candidates = ['not_visible']
        mapping = {}
        next_id = 1
        if ignore_candidates:
            ignore_label = ignore_candidates[0]
            mapping['0'] = ignore_label
            # Remove chosen ignore from list for numbering
            unique_no_ignore = [u for u in unique if u != ignore_label]
        else:
            unique_no_ignore = unique
        for u in unique_no_ignore:
            mapping[str(next_id)] = u
            next_id += 1
        groups[col] = mapping
        logging.info(f"Built label group '{col}': {mapping}")
    return groups

# --- Custom Dataset for Multi-Label Data ---
class MultiLabelDataset(Dataset):
    def __init__(self, annotations_df, img_dir, label_groups, transform=None):
        self.annotations_df = annotations_df.reset_index(drop=True)
        self.img_dir = img_dir
        self.transform = transform
        self.label_groups = label_groups
        # Build label -> index maps per group in the order they appear in values()
        self.label_maps = {
            group: {label: i for i, label in enumerate(classes.values())}
            for group, classes in self.label_groups.items()
        }

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

    def __getitem__(self, idx):
        img_name = os.path.join(self.img_dir, self.annotations_df.iloc[idx]['filename'])
        image = Image.open(img_name).convert('RGB')
        labels = {}
        for group, label_map in self.label_maps.items():
            label_name = self.annotations_df.iloc[idx][group]
            labels[group] = torch.tensor(label_map[label_name], dtype=torch.long)
        if self.transform:
            image = self.transform(image)
        return image, labels

# --- Multi-Head Model ---
class MultiHeadResNet(nn.Module):
    def __init__(self, base_model, label_groups):
        super().__init__()
        # Remove the final layer of the base model
        self.base = nn.Sequential(*list(base_model.children())[:-1])
        in_features = base_model.fc.in_features
        
        # Create a separate output layer (head) for each label group
        self.heads = nn.ModuleDict({
            g: nn.Linear(in_features, len(classes))
            for g, classes in label_groups.items()
        })

    def forward(self, x):
        x = self.base(x)
        x = torch.flatten(x, 1)
        # Get outputs from each head
        return {g: head(x) for g, head in self.heads.items()}

def train_single_model(tag, annotations_df, args, label_groups):
    """
    Trains a single model (all species or per-species).
    tag: 'all_species' or species name.
    """
    # Always store models in runs/segment
    save_dir = os.path.join(DEFAULT_RUNS_SEGMENT_DIR, tag if tag != "all_species" else "all_species")
    model_save_path = os.path.join(save_dir, "best_multihead_classifier.pt")
    results_save_path = os.path.join(save_dir, "results.csv")

    df = annotations_df.copy()
    if df.empty:
        logging.warning(f"[{tag}] No samples – skipping.")
        return
    logging.info(f"[{tag}] Samples: {len(df)}")

    # Split (80/10/10)
    train_df = df.sample(frac=0.8, random_state=42)
    temp_df = df.drop(train_df.index)
    val_df = temp_df.sample(frac=0.5, random_state=42)
    test_df = temp_df.drop(val_df.index)

    data_transforms = {
        'train': transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.RandomHorizontalFlip(),
            transforms.RandomRotation(15),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406],[0.229,0.224,0.225])
        ]),
        'val': transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406],[0.229,0.224,0.225])
        ]),
        'test': transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406],[0.229,0.224,0.225])
        ]),
    }

    image_datasets = {
        'train': MultiLabelDataset(train_df, args.labeled_cutouts_dir, label_groups, data_transforms['train']),
        'val': MultiLabelDataset(val_df, args.labeled_cutouts_dir, label_groups, data_transforms['val']),
        'test': MultiLabelDataset(test_df, args.labeled_cutouts_dir, label_groups, data_transforms['test'])
    }
    dataloaders = {
        'train': DataLoader(image_datasets['train'], batch_size=args.batch_size, shuffle=True, num_workers=4),
        'val': DataLoader(image_datasets['val'], batch_size=args.batch_size, shuffle=False, num_workers=4),
        'test': DataLoader(image_datasets['test'], batch_size=args.batch_size, shuffle=False, num_workers=4)
    }

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

    base_model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
    model = MultiHeadResNet(base_model, label_groups).to(device)

    # A separate loss function for each head
    criterions = {g: nn.CrossEntropyLoss() for g in label_groups.keys()}
    optimizer = optim.Adam(model.parameters(), lr=args.lr)

    best_acc = 0.0
    history = []

    for epoch in range(args.epochs):
        logging.info(f"Epoch {epoch+1}/{args.epochs}")
        for phase in ['train', 'val']:
            model.train() if phase == 'train' else model.eval()

            running_loss = 0.0
            running_corrects = {g: 0 for g in label_groups.keys()}
            running_totals = {g: 0 for g in label_groups.keys()}
            
            for inputs, labels in tqdm(dataloaders[phase], desc=f"{phase.capitalize()} Phase"):
                inputs = inputs.to(device)
                labels = {g: l.to(device) for g,l in labels.items()}
                optimizer.zero_grad()
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    total_loss = 0
                    for g, logits in outputs.items():
                        group_labels = labels[g]
                        unknown_label_name = label_groups[g].get('0', None)
                        if unknown_label_name:
                            unknown_idx = image_datasets[phase].label_maps[g][unknown_label_name]
                            valid_indices = (group_labels != unknown_idx)
                        else:
                            valid_indices = torch.ones_like(group_labels, dtype=torch.bool)
                        if valid_indices.any():
                            valid_logits = logits[valid_indices]
                            valid_labels = group_labels[valid_indices]
                            loss = criterions[g](valid_logits, valid_labels)
                            total_loss += loss

                            # Calculate accuracy only on valid samples
                            _, preds = torch.max(valid_logits, 1)
                            running_corrects[g] += torch.sum(preds == valid_labels.data)
                            running_totals[g] += len(valid_labels)

                    if phase == 'train':
                        if isinstance(total_loss, torch.Tensor):
                            total_loss.backward()
                            optimizer.step()

                if isinstance(total_loss, torch.Tensor):
                    running_loss += total_loss.item() * inputs.size(0)

            epoch_loss = running_loss / len(image_datasets[phase])
            # Calculate accuracy based on valid samples only and convert to float
            epoch_accs = {
                g: (running_corrects[g].double() / running_totals[g]).item() if running_totals[g] > 0 else 0.0
                for g in label_groups.keys()
            }
            # Average only over groups with at least one valid sample
            valid_groups = [g for g in label_groups.keys() if running_totals[g] > 0]
            avg_acc = (sum(epoch_accs[g] for g in valid_groups) / len(valid_groups)) if valid_groups else 0.0

            # Log only validation metrics to keep the output clean
            if phase == 'val':
                logging.info(f'Validation Loss: {epoch_loss:.4f}')
                for g,a in epoch_accs.items():
                    logging.info(f'  - {g} Acc: {a:.4f}')
                logging.info(f'  - Average Acc: {avg_acc:.4f}')

            # Store results for both train and val phases
            row = {'epoch': epoch+1, 'phase': phase, 'loss': epoch_loss, 'avg_acc': avg_acc}
            row.update({f'{g}_acc': a for g,a in epoch_accs.items()})
            history.append(row)

            # Save model based on average validation accuracy
            if phase == 'val' and avg_acc > best_acc:
                best_acc = avg_acc
                os.makedirs(save_dir, exist_ok=True)
                torch.save(model.state_dict(), model_save_path)
                logging.info(f"[{tag}] New best model saved to {model_save_path} (avg_acc={best_acc:.4f})")

        # Save training history to CSV after each epoch
        os.makedirs(save_dir, exist_ok=True)
        pd.DataFrame(history).to_csv(results_save_path, index=False)

    logging.info(f"[{tag}] Training complete.")
    logging.info(f"[{tag}] Evaluating on test set.")
    model.load_state_dict(torch.load(model_save_path))
    model.eval()

    test_corrects = {g: 0 for g in label_groups.keys()}
    test_totals = {g: 0 for g in label_groups.keys()}
    test_loss = 0.0
    with torch.no_grad():
        for inputs, labels in tqdm(dataloaders['test'], desc="Test Phase"):
            inputs = inputs.to(device)
            labels = {g:l.to(device) for g,l in labels.items()}
            outputs = model(inputs)
            total_loss = 0
            for g, logits in outputs.items():
                group_labels = labels[g]
                unknown_label_name = label_groups[g].get('0', None)
                if unknown_label_name:
                    unknown_idx = image_datasets['test'].label_maps[g][unknown_label_name]
                    valid_indices = (group_labels != unknown_idx)
                else:
                    valid_indices = torch.ones_like(group_labels, dtype=torch.bool)
                if valid_indices.any():
                    valid_logits = logits[valid_indices]
                    valid_labels = group_labels[valid_indices]
                    loss = criterions[g](valid_logits, valid_labels)
                    total_loss += loss
                    _, preds = torch.max(valid_logits,1)
                    test_corrects[g] += torch.sum(preds == valid_labels.data)
                    test_totals[g] += len(valid_labels)
            test_loss += total_loss.item() * inputs.size(0)

    final_loss = test_loss / len(image_datasets['test'])
    final_accs = {
        g: (test_corrects[g].double() / test_totals[g]).item() if test_totals[g] > 0 else 0.0
        for g in label_groups.keys()
    }
    valid_test_groups = [g for g in label_groups.keys() if test_totals[g] > 0]
    avg_final = (sum(final_accs[g] for g in valid_test_groups)/len(valid_test_groups)) if valid_test_groups else 0.0

    logging.info(f"[{tag}] Test Loss: {final_loss:.4f}")
    for g,a in final_accs.items():
        logging.info(f"[{tag}]  - {g} Test Acc: {a:.4f}")
    logging.info(f"[{tag}]  - Average Test Acc: {avg_final:.4f}")

def attach_class_to_annotations(annotations_df, results_csv_path):
    """
    Adds 'class' by joining with the specified results CSV.
    """
    if not os.path.exists(results_csv_path):
        logging.warning(f"No detection results file at {results_csv_path}; skipping class mapping.")
        annotations_df['class'] = ""
        return annotations_df
    try:
        det_df = pd.read_csv(results_csv_path)
    except Exception as e:
        logging.error(f"Failed reading {results_csv_path}: {e}")
        annotations_df['class'] = ""
        return annotations_df

    # Always use 'class' as the grouping column
    if 'class' not in det_df.columns:
        logging.warning(f"{results_csv_path} missing required column 'class'; cannot assign class labels.")
        annotations_df['class'] = ""
        return annotations_df

    # Updated regex to match filenames like ..._20_0_leaf_black_spot.tiff
    pattern = re.compile(r"^(?P<core>.+)_(?P<idx>\d+)_.*\.tiff$")
    cores, insts = [], []
    for fn in annotations_df['filename']:
        m = pattern.match(fn)
        if m:
            cores.append(m.group('core') + ".tiff")
            insts.append(int(m.group('idx')) + 1)
        else:
            cores.append(None)
            insts.append(None)
    annotations_df['image_name_join'] = cores
    annotations_df['instance_number_join'] = insts

    # Ensure join keys have the same data type to prevent merge errors
    annotations_df['instance_number_join'] = annotations_df['instance_number_join'].astype('float64')
    det_df['instance_number'] = det_df['instance_number'].astype('float64')

    merged = annotations_df.merge(
        det_df[['image_name', 'instance_number', 'class']],
        left_on=['image_name_join', 'instance_number_join'],
        right_on=['image_name', 'instance_number'],
        how='left'
    )
    merged['class'] = merged['class'].fillna("")

    return merged.drop(columns=['image_name_join', 'instance_number_join', 'image_name', 'instance_number'])

def train_classifier(args):
    logging.info("--- Dynamic Multi-Head Classifier Training ---")
    if not os.path.exists(args.annotation_file):
        logging.error(f"Annotation file not found: {args.annotation_file}")
        return
    annotations = pd.read_csv(args.annotation_file)
    # Build dynamic groups BEFORE filtering 'skipped' for columns discovery
    label_groups = build_label_groups(annotations, target_columns=[c for c in annotations.columns if c != 'filename'])
    # Remove skipped rows
    for g in label_groups.keys():
        annotations = annotations[annotations[g] != 'skipped']
    logging.info(f"Filtered annotations count: {len(annotations)}")
    # Attach class information from the main results file
    annotations = attach_class_to_annotations(annotations, args.results_csv_path)
    
    # Log detected classes for user verification
    class_vals = [c for c in annotations['class'].unique() if c]
    if class_vals:
        logging.info(f"Detected classes for training: {sorted(class_vals)}")
    elif not os.path.exists(args.results_csv_path):
        logging.warning(f"Results file not found at '{args.results_csv_path}'. Training a single global model.")
    else:
        logging.info("No class information was found after processing the results file. Training a single global model.")

    # Always train global model
    train_single_model("all_species", annotations, args, label_groups)
    # Train per-class if more than one class present
    if len(class_vals) > 1:
        logging.info(f"Per-class models will be trained for: {sorted(class_vals)}")
        for c in class_vals:
            train_single_model(c, annotations[annotations['class']==c], args, label_groups)
    else:
        logging.info("No multiple classes detected; skipping per-class models.")

def show_label_groups(args):
    if not os.path.exists(args.annotation_file):
        print(f"Annotation file not found: {args.annotation_file}")
        return
    df = pd.read_csv(args.annotation_file)
    groups = build_label_groups(df, target_columns=[c for c in df.columns if c != 'filename'])
    print(json.dumps(groups, indent=2))

def main():
    parser = argparse.ArgumentParser(description="Train a multi-head classifier for leaf traits.",
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    # Path arguments
    parser.add_argument("--labeled-cutouts-dir", type=str, default=DEFAULT_LABELED_CUTOUTS_DIR, help="Directory containing labeled cutout images.")
    parser.add_argument("--annotation-file", type=str, default=DEFAULT_ANNOTATION_FILE, help="Path to the CSV file with annotations.")
    parser.add_argument("--results-csv-path", type=str, default=DEFAULT_RESULTS_CSV_PATH, help="Path to the main results CSV for class mapping.")
    # Always use runs/segment for model outputs
    parser.add_argument("--all-species-run-dir", type=str, default=DEFAULT_ALL_SPECIES_RUN_DIR, help="Directory to save the global model.")
    parser.add_argument("--per-species-run-dir", type=str, default=DEFAULT_PER_SPECIES_RUN_DIR, help="Base directory to save per-class models.")
    parser.add_argument("--log-file", type=str, default=DEFAULT_LOG_FILE, help="Path to the master log file.")
    
    # Hyperparameter arguments
    parser.add_argument("--epochs", type=int, default=DEFAULT_NUM_EPOCHS, help="Number of training epochs.")
    parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE, help="Batch size for training and evaluation.")
    parser.add_argument("--lr", type=float, default=DEFAULT_LEARNING_RATE, help="Learning rate for the optimizer.")

    # Action arguments
    parser.add_argument("--show-groups", action="store_true", help="Build and display dynamic LABEL_GROUPS then exit.")
    parser.add_argument("--use-sample-data", action="store_true", help="Use sample data and annotations for validation.")

    args = parser.parse_args()

    # Store the original path before any overrides
    original_results_csv_path = args.results_csv_path

    # Override paths if using sample data
    if args.use_sample_data:
        logging.info("--- Using sample data mode ---")
        args.labeled_cutouts_dir = SAMPLE_LABELED_CUTOUTS_DIR
        args.annotation_file = SAMPLE_ANNOTATION_FILE
        # IMPORTANT: When using sample data, we still want to map classes
        # from the *actual* inference results, not sample results.
        args.results_csv_path = DEFAULT_RESULTS_CSV_PATH


    setup_logging(args.log_file)
    logging.info("--- Starting script: 52_train_classifier.py ---")

    if args.show_groups:
        show_label_groups(args)
        logging.info("--- Finished script: 52_train_classifier.py (show-groups) ---")
        return

    train_classifier(args)
    logging.info("--- Finished script: 52_train_classifier.py ---")

if __name__ == "__main__":
    main()