get_visual_differences.py

# Copyright 2025 Adobe Research. All rights reserved.
# To view a copy of the license, visit LICENSE.md.

import logging
import os
from typing import Dict, List, Tuple

import click
import numpy as np
import pandas as pd
from omegaconf import OmegaConf
import random
from tqdm import tqdm

import weave
import wandb
from components.evaluator import GPTEvaluator, NullEvaluator
# from components.evolver import (
#     Evolver,
#     LLMEvolverAttribute,
#     LLMEvolverCLIP,
#     LLMEvolverCLIPAcc,
#     LLMEvolverCLIPAccExclusive,
# )
from components.proposer import (
    LLMProposer,
    LLMProposerDiffusion,
    VLMFeatureProposer,
    VLMProposer,
    VLMProposerPairwise,
    VLMProposerPairwiseModified,
    VLMProposerErrorAnalysis,
    TFIDFProposer,
)
from components.ranker import (
    CLIPPairwiseRanker,
    CLIPPairwiseRankerMulti,
    CLIPPairwiseRankerDedup,
    CLIPRanker,
    CLIPZSRanker,
    LLMRanker,
    NullRanker,
    VLMRanker,
)

# os.environ["WANDB_API_KEY"] = "local-5a5dbfd2a8a0e4e1ecf93025e64cf05bfec4e069"
# wandb.login(host="https://adobesensei.wandb.io/")

def load_config(config: str, overrides: list) -> Dict:
    base_cfg = OmegaConf.load("configs/base.yaml")
    print("Overrides ", overrides)
    cfg = OmegaConf.load(config)
    overrides_cfg = OmegaConf.from_dotlist(overrides)
    final_cfg = OmegaConf.merge(base_cfg, cfg, overrides_cfg)
    args = OmegaConf.to_container(final_cfg)
    args["config"] = config
    print(args)
    if args["wandb"]:
        wandb.init(
            project=args["project"],
            # entity="research-dmo",
            name=args["name"] if args["name"] else args["data"]["name"],
            group=f'{args["data"]["group1"]} - {args["data"]["group2"]} ({args["data"]["purity"]})',
            config=args,
        )
    else:
        wandb.init(mode="disabled")
    return args


def load_data(args: Dict, drop_duplicates: bool = True) -> Tuple[List[Dict], List[Dict], List[str]]:
    random.seed(args["seed"])
    data_args = args["data"]

    df = pd.read_csv(f"{data_args['root']}/{data_args['name']}.csv")
    print(len(df))

    # drop doplicates
    # if drop_duplicates:
    #     df = df.drop_duplicates(subset=["Prompt", "group_name"])

    if data_args["subset"]:
        old_len = len(df)
        df = df[df["subset"] == data_args["subset"]]
        print(
            f"Taking {data_args['subset']} subset (dataset size reduced from {old_len} to {len(df)})"
        )
    
    logging.warning(f"Purity is set to {data_args['purity']}.")
    assert len(df[df["group_name"] == data_args["group1"]]) == len(df[df["group_name"] == data_args["group2"]]), f"Groups must be of equal size: {len(df[df['group_name'] == data_args['group1']])} != {len(df[df['group_name'] == data_args['group2']])}"

    # Group by 'Prompt'
    grouped1 = df[df["group_name"] == data_args["group1"]].groupby("Prompt")
    grouped2 = df[df["group_name"] == data_args["group2"]].groupby("Prompt")

    # Convert groups to lists
    groups1 = [group.to_dict("records") for _, group in grouped1]
    groups2 = [group.to_dict("records") for _, group in grouped2]

    # Ensure equal number of groups
    assert len(groups1) == len(groups2), f"Groups must be of equal size: {len(groups1)} != {len(groups2)}"

    # print(groups1[0])
    # prompts_wo_filler = [d[0]["Prompt"] for d in groups1 if d[0]["visual_attribute"] != "filler"]
    # print(prompts_wo_filler)

    # if data_args["purity"] < 1:
    #     # Calculate number of groups to swap
    #     n_swap = int((1 - data_args["purity"]) * len(prompts_wo_filler))

    #     prompts_to_swap = random.sample(prompts_wo_filler, n_swap)

    #     new_groups1 = []
    #     new_groups2 = []
    #     for g1, g2 in zip(groups1, groups2):
    #         if g1[0]["Prompt"] in prompts_to_swap:
    #             new_groups1.append(g2)
    #             new_groups2.append(g1)
    #         else:
    #             new_groups1.append(g1)
    #             new_groups2.append(g2)

    #     groups1 = new_groups1
    #     groups2 = new_groups2

    # Flatten the lists back to datasets
    dataset1 = [item for group in groups1 for item in group]
    dataset2 = [item for group in groups2 for item in group]

    group_names = [data_args["group1"], data_args["group2"]]

    return dataset1, dataset2, group_names


def propose(args: Dict, dataset1: List[Dict], dataset2: List[Dict]) -> List[str]:
    proposer_args = args["proposer"]
    proposer_args["seed"] = args["seed"]
    proposer_args["captioner"] = args["captioner"]

    group_names = [args["data"]["group1"], args["data"]["group2"]]

    proposer = eval(proposer_args["method"])(proposer_args)

    if proposer_args["method"] == "VLMProposerPairwiseModified":
        (hypotheses1, hypotheses2), logs, images = proposer.propose(dataset1, dataset2)
        logs = pd.DataFrame(logs)
    elif proposer_args["method"] == "VLMProposerErrorAnalysis":
        (hypotheses1, hypotheses2), logs, images = proposer.propose(dataset1, dataset2)
        logs = pd.DataFrame(logs)
        logs["model"] = group_names[0]
    else:
        hypotheses1, logs, images = proposer.propose(dataset1, dataset2)
        logs = pd.DataFrame(logs)
        logs["model"] = group_names[0]

        hypotheses2, logs2, images2 = proposer.propose(dataset2, dataset1)
        logs2 = pd.DataFrame(logs2)
        logs2["model"] = group_names[1]

        logs = pd.concat([logs, logs2])

    if args["wandb"]:
        if "image" in logs.columns:
            logs["image"] = logs.apply(
                lambda row: wandb.Image(row["image"], caption=row["prompt"]), axis=1
            )
        wandb.log({"logs": wandb.Table(dataframe=logs)})

    return hypotheses1, hypotheses2


def rank(
    args: Dict,
    hypotheses: List[str],
    dataset1: List[Dict],
    dataset2: List[Dict],
    group_name: str,
) -> List[str]:
    ranker_args = args["ranker"]
    ranker_args["seed"] = args["seed"]

    ranker = eval(ranker_args["method"])(ranker_args)

    scored_hypotheses = ranker.rerank_hypotheses(hypotheses, dataset1, dataset2)
    print(scored_hypotheses)
    if args["wandb"]:
        table_hypotheses = wandb.Table(dataframe=pd.DataFrame(scored_hypotheses))
        print(pd.DataFrame(scored_hypotheses).head())
        wandb.log({f"scored hypotheses: {group_name}": table_hypotheses})
        for i in range(5):
            wandb.summary[f"top_{i + 1}_difference_{group_name}"] = scored_hypotheses[i][
                "hypothesis"
            ].replace('"', "")
            wandb.summary[f"top_{i + 1}_score_{group_name}"] = scored_hypotheses[i][
                args["ranker"]["metric"]
            ]

    # scored_groundtruth = ranker.rerank_hypotheses(
    #     group_names,
    #     dataset1,
    #     dataset2,
    # )
    # if args["wandb"]:
    #     table_groundtruth = wandb.Table(dataframe=pd.DataFrame(scored_groundtruth))
    #     wandb.log({"scored groundtruth": table_groundtruth})

    return [hypothesis["hypothesis"] for hypothesis in scored_hypotheses]


def evolve(
    args: Dict,
    hypotheses1: List[str],
    hypotheses2: List[str],
    dataset1: List[Dict],
    dataset2: List[Dict],
) -> List[str]:
    evolver_args = args["evolver"]
    evolver_args["seed"] = args["seed"]
    evolver_args["load"] = args["load"]
    evolver_args["load_checkpoint"] = args["load_checkpoint"]
    evolver_args["name"] = args["name"]
    ranker_args = args["ranker"]

    # evolver = LLMEvolverCLIP(evolver_args, ranker_args)
    evolver = eval(evolver_args["method"])(evolver_args, ranker_args)
    new_hypotheses = evolver.evolve(hypotheses1, hypotheses2, dataset1, dataset2)
    return new_hypotheses


def evaluate(args: Dict, ranked_hypotheses: List[str], group_names: List[str]) -> Dict:
    evaluator_args = args["evaluator"]

    evaluator = eval(evaluator_args["method"])(evaluator_args)

    metrics, evaluated_hypotheses = evaluator.evaluate(
        ranked_hypotheses,
        group_names[0],
        group_names[1],
    )

    if args["wandb"] and evaluator_args["method"] != "NullEvaluator":
        table_evaluated_hypotheses = wandb.Table(dataframe=pd.DataFrame(evaluated_hypotheses))
        wandb.log({"evaluated hypotheses": table_evaluated_hypotheses})
        wandb.log(metrics)
    return metrics


from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.cluster import KMeans
from sklearn.metrics.pairwise import cosine_similarity


def cluster_strings(strings, num_final_strings):
    # Step 1: Generate embeddings
    vectorizer = TfidfVectorizer()
    embeddings = vectorizer.fit_transform(strings)

    # Step 2: Perform clustering
    kmeans = KMeans(n_clusters=num_final_strings, random_state=42)
    cluster_labels = kmeans.fit_predict(embeddings)

    # Step 3: Select representative strings
    representative_strings = []
    for i in range(num_final_strings):
        cluster_strings = [s for s, label in zip(strings, cluster_labels) if label == i]
        cluster_embeddings = vectorizer.transform(cluster_strings)

        # Find the string closest to the cluster centroid
        centroid = kmeans.cluster_centers_[i]
        similarities = cosine_similarity(cluster_embeddings, centroid.reshape(1, -1)).flatten()
        representative_index = similarities.argmax()
        representative_strings.append(cluster_strings[representative_index])

    return representative_strings


@click.command()
@click.option("--config", help="Config file", type=str, default="configs/base.yaml")
@click.option("--overrides", multiple=True, help="List of overrides for the configuration")
def main(config, overrides):
    logging.info("Loading config...")
    args = load_config(config, list(overrides))
    # print(args)

    logging.info("Loading data...")
    dataset1, dataset2, group_names = load_data(args)
    print(len(dataset1), len(dataset2), group_names)

    logging.info("Proposing hypotheses...")
    hypotheses1, hypotheses2 = propose(args, dataset1, dataset2)
    if len(hypotheses1) < 100 or len(hypotheses2) < 100:
        reduced_hypotheses1 = hypotheses1
        reduced_hypotheses2 = hypotheses2
    else:
        reduced_hypotheses1 = cluster_strings(hypotheses1, 50)
        reduced_hypotheses2 = cluster_strings(hypotheses2, 50)
    reduced_hypotheses1 = [r.strip().replace("\"", "").replace("\'", "") for r in reduced_hypotheses1]
    reduced_hypotheses2 = [r.strip().replace("\"", "").replace("\'", "") for r in reduced_hypotheses2]
    data = [
        {
            "model": group_names[0],
            "hypotheses": "\n".join(hypotheses1),
            "reduced_hypotheses": "\n".join(reduced_hypotheses1),
        },
        {
            "model": group_names[1],
            "hypotheses": "\n".join(hypotheses2),
            "reduced_hypotheses": "\n".join(reduced_hypotheses1),
        },
    ]
    print(pd.DataFrame(data))
    wandb.log({"Reduced hypotheses": wandb.Table(dataframe=pd.DataFrame(data))})

    # try:
    logging.info("Ranking hypotheses...")
    ranked_hypotheses = rank(args, reduced_hypotheses1, dataset1, dataset2, group_names[0])
    print("Ranking group 1")
    ranked_hypotheses2 = rank(args, reduced_hypotheses2, dataset2, dataset1, group_names[1])
    print("First group")
    print(ranked_hypotheses)
    print("Second group")
    print(ranked_hypotheses2)
    # except Exception as e:
    #     print(f"Error in ranking: {e}")
    #     pass
    # logging.info("Evaluating hypotheses...")
    # metrics = evaluate(args, ranked_hypotheses, group_names)
    # print(metrics)


if __name__ == "__main__":
    main()