utils_analysis.py

"""Helpers for the Analysis / Explore tab."""

import pandas as pd

OPEN_LICENSE_MARKERS = (
    "apache", "mit", "bsd", "open", "cc0", "cc-by-4", "llama",
)

PROPRIETARY_LICENSE_MARKERS = (
    "nc", "non-commercial", "apple sample", "nvidia source",
    "proprietary", "all rights",
)

FAMILY_KEYWORDS = (
    ("convnext", "ConvNeXt"),
    ("swinv2", "SwinV2"),
    ("swin", "Swin"),
    ("deit3", "DeiT-III"),
    ("deit", "DeiT"),
    ("dinov2", "DINOv2"),
    ("internimage", "InternImage"),
    ("efficientnet", "EfficientNet"),
    ("efficientvit", "EfficientViT"),
    ("efficientformer", "EfficientFormer"),
    ("mobilenet", "MobileNet"),
    ("mobilevit", "MobileViT"),
    ("mobileone", "MobileOne"),
    ("fastvit", "FastViT"),
    ("repvit", "RepViT"),
    ("regnet", "RegNet"),
    ("beit", "BEiT"),
    ("hiera", "Hiera"),
    ("maxvit", "MaxViT"),
    ("eva02", "EVA-02"),
    ("vit", "ViT"),
    ("resnet", "ResNet"),
    ("resnext", "ResNeXt"),
    ("densenet", "DenseNet"),
    ("inception", "Inception"),
    ("poolformer", "PoolFormer"),
    ("gcvit", "GCViT"),
    ("levit", "LeViT"),
    ("nfnet", "NFNet"),
    ("coat", "CoAt"),
    ("mixer", "MLP-Mixer"),
    ("xcit", "XCiT"),
    ("volo", "VOLO"),
    ("mit", "MiT"),
)

X_AXIS_OPTIONS = {
    "Parameters (M)": "parameters_millions",
    "FLOPs (G)": "flops_giga",
    "Model Size (MB)": "model_size_mb",
    "Year": "year",
}

Y_AXIS_OPTIONS = {
    "Top-1 Accuracy (%)": "top1_accuracy",
    "Top-5 Accuracy (%)": "top5_accuracy",
}


def infer_architecture_family(model_path: str) -> str:
    text = model_path.lower()
    name = text.split("/")[-1]
    for keyword, label in FAMILY_KEYWORDS:
        if keyword in name or keyword in text:
            return label
    if "/" in model_path:
        return model_path.split("/")[0].title()
    return "Other"


def is_open_license(license_value) -> bool:
    if license_value is None or pd.isna(license_value):
        return True
    lic = str(license_value).lower()
    if any(marker in lic for marker in PROPRIETARY_LICENSE_MARKERS):
        return False
    return any(marker in lic for marker in OPEN_LICENSE_MARKERS)


def prepare_analysis_df(raw_df: pd.DataFrame) -> pd.DataFrame:
    df = raw_df.copy()
    df = df[df["top1_accuracy"].notna() & (df["top1_accuracy"] > 0)]
    df["year"] = pd.to_numeric(df.get("year"), errors="coerce")
    df["architecture_family"] = df["model"].map(infer_architecture_family)
    df["open_license"] = df.get("license", pd.Series(dtype=object)).map(is_open_license)
    return df


def filter_analysis_df(
    df: pd.DataFrame,
    open_only: bool = False,
    min_year: int | None = None,
    max_year: int | None = None,
) -> pd.DataFrame:
    out = df.copy()
    if open_only and "open_license" in out.columns:
        out = out[out["open_license"]]
    if min_year is not None:
        out = out[out["year"].isna() | (out["year"] >= min_year)]
    if max_year is not None:
        out = out[out["year"].isna() | (out["year"] <= max_year)]
    return out


def build_scatter_df(
    df: pd.DataFrame,
    x_label: str,
    y_label: str,
) -> pd.DataFrame:
    x_col = X_AXIS_OPTIONS[x_label]
    y_col = Y_AXIS_OPTIONS[y_label]
    cols = [x_col, y_col, "architecture_family", "model"]
    plot_df = df[cols].dropna(subset=[x_col, y_col])
    return plot_df.sort_values(y_col, ascending=False)


def build_year_trend_df(df: pd.DataFrame) -> pd.DataFrame:
    yearly = (
        df.dropna(subset=["year", "top1_accuracy"])
        .groupby("year", as_index=False)
        .agg(
            best_top1=("top1_accuracy", "max"),
            model_count=("model", "count"),
        )
        .sort_values("year")
    )
    return yearly