# -*- coding: utf-8 -*- """Train/evaluate an image-only density estimator. Labels such as brix_raw_label / soup density are used only as targets. They are explicitly excluded from the input feature matrix. """ from __future__ import annotations import argparse import csv import json from pathlib import Path from typing import Any import numpy as np from sklearn.ensemble import RandomForestClassifier from sklearn.metrics import accuracy_score, confusion_matrix, classification_report, f1_score from sklearn.model_selection import StratifiedKFold, cross_val_predict from sklearn.neighbors import KNeighborsClassifier from sklearn.pipeline import make_pipeline from sklearn.preprocessing import StandardScaler from sklearn.svm import SVC EXCLUDED_COLUMNS = { "id", "video_path", "csv_row_index", "store", "slot", "brix_raw_label", "human_soup_score", "human_total_score", "human_comment", "frame_count", "frame_ratios", "frames_json", "error", } def main() -> int: parser = argparse.ArgumentParser(description="Evaluate image-only density features.") parser.add_argument("--features", required=True, help="CSV created by build_density_features.py.") parser.add_argument("--label-column", default="brix_raw_label") parser.add_argument("--folds", type=int, default=5) parser.add_argument("--output", default="") parser.add_argument("--merge3", action="store_true", help="Map 10/11=thin, 12=standard, 13/14=thick.") args = parser.parse_args() rows = read_rows(Path(args.features)) clean_rows = [row for row in rows if not row.get("error") and row.get(args.label_column)] x, y, feature_names, ids = build_matrix(clean_rows, args.label_column) if args.merge3: y = np.asarray([merge3_label(value) for value in y]) if len(set(y)) < 2: raise SystemExit("Need at least two label classes to evaluate.") min_class_count = min(np.bincount(np.unique(y, return_inverse=True)[1])) folds = max(2, min(args.folds, min_class_count, len(y))) if folds < 2: raise SystemExit("Not enough samples per class for cross validation.") cv = StratifiedKFold(n_splits=folds, shuffle=True, random_state=42) labels = [str(value) for value in sorted(set(y))] model_results = {} predictions_by_model = {} for model_name, model in build_models().items(): pred = cross_val_predict(model, x, y, cv=cv) model_results[model_name] = evaluate_predictions(y, pred, labels) predictions_by_model[model_name] = [ {"id": str(item_id), "actual": str(actual), "predicted": str(predicted)} for item_id, actual, predicted in zip(ids, y, pred) ] best_model = max(model_results, key=lambda name: model_results[name]["accuracy"]) report = { "rows": len(clean_rows), "features": feature_names, "folds": folds, "merge3": args.merge3, "labels": labels, "best_model": best_model, "accuracy": model_results[best_model]["accuracy"], "within1_accuracy": model_results[best_model]["within1_accuracy"], "macro_f1": model_results[best_model]["macro_f1"], "model_results": model_results, "predictions": predictions_by_model[best_model], } report = to_jsonable(report) print(json.dumps(report, ensure_ascii=False, indent=2)) if args.output: output_path = Path(args.output) output_path.parent.mkdir(parents=True, exist_ok=True) output_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8") return 0 def build_models() -> dict[str, Any]: return { "svm_rbf": make_pipeline( StandardScaler(), SVC(kernel="rbf", C=4.0, gamma="scale", class_weight="balanced"), ), "knn": make_pipeline( StandardScaler(), KNeighborsClassifier(n_neighbors=7, weights="distance"), ), "rf_depth8": make_pipeline( StandardScaler(), RandomForestClassifier( n_estimators=500, max_depth=8, min_samples_leaf=2, random_state=42, ), ), "rf_depth4": make_pipeline( StandardScaler(), RandomForestClassifier( n_estimators=300, max_depth=4, min_samples_leaf=2, random_state=42, ), ), } def evaluate_predictions(y_true: np.ndarray, y_pred: np.ndarray, labels: list[str]) -> dict[str, Any]: return { "accuracy": float(accuracy_score(y_true, y_pred)), "within1_accuracy": within1_accuracy(y_true, y_pred), "macro_f1": float(f1_score(y_true, y_pred, labels=labels, average="macro", zero_division=0)), "confusion_matrix": confusion_matrix(y_true, y_pred, labels=labels).tolist(), "classification_report": classification_report(y_true, y_pred, labels=labels, zero_division=0, output_dict=True), } def within1_accuracy(y_true: np.ndarray, y_pred: np.ndarray) -> float: try: true_values = np.asarray([int(value) for value in y_true]) pred_values = np.asarray([int(value) for value in y_pred]) except ValueError: return float(accuracy_score(y_true, y_pred)) return float((np.abs(true_values - pred_values) <= 1).mean()) def merge3_label(value: Any) -> str: number = int(float(str(value))) if number <= 11: return "thin" if number == 12: return "standard" return "thick" def read_rows(path: Path) -> list[dict[str, str]]: with path.open("r", encoding="utf-8-sig", newline="") as fh: return list(csv.DictReader(fh)) def build_matrix(rows: list[dict[str, str]], label_column: str) -> tuple[np.ndarray, np.ndarray, list[str], list[str]]: feature_names = numeric_feature_names(rows, label_column) values: list[list[float]] = [] labels: list[str] = [] ids: list[str] = [] for row in rows: vector: list[float] = [] for name in feature_names: vector.append(float(row[name])) values.append(vector) labels.append(str(row[label_column])) ids.append(str(row.get("id", ""))) return np.asarray(values, dtype=np.float64), np.asarray(labels), feature_names, ids def numeric_feature_names(rows: list[dict[str, str]], label_column: str) -> list[str]: excluded = set(EXCLUDED_COLUMNS) excluded.add(label_column) names: list[str] = [] for key in rows[0].keys(): if key in excluded: continue try: for row in rows: if row.get(key, "") != "": float(row[key]) names.append(key) except ValueError: continue return names def to_jsonable(value: Any) -> Any: if isinstance(value, dict): return {str(key): to_jsonable(item) for key, item in value.items()} if isinstance(value, list): return [to_jsonable(item) for item in value] if isinstance(value, tuple): return [to_jsonable(item) for item in value] if isinstance(value, np.integer): return int(value) if isinstance(value, np.floating): return float(value) if isinstance(value, np.ndarray): return value.tolist() return value if __name__ == "__main__": raise SystemExit(main())