intprogramming_per_image.py

import itertools
from pulp import LpProblem, LpMaximize, LpVariable, lpSum, value
from pulp import PULP_CBC_CMD
import time
from utils import *
import tqdm
import multiprocessing as mp
from functools import partial
import os


##########################
##########################
epsilon_values = [0.01, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1]
# epsilon_values = [0.5]
deltas = [0.01, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1]
# deltas = [0.9, 0.3, 0.01]
##########################
##########################
classes = [0, 1, 2, 3]
ic = set(list(itertools.permutations(classes, 2)))
models = [
    "normal-1000",
    "rain-1000",
    "fog-1000",
    "maple_leaf-1000",
    "snow-1000",
    "dust-1000",
]
##########################
##########################
condition_map = {
    "normal-1000": "pfca_normal-1000",
    "rain-1000": "pfca_rain-1000",
    "fog-1000": "pfca_fog-1000",
    "maple_leaf-1000": "pfca_maple_leaf-1000",
    "snow-1000": "pfca_snow-1000",
    "dust-1000": "pfca_dust-1000",
}


def optimize_for_threshold(objects, pred, delta):

    # Create the optimization problem
    prob = LpProblem(f"Optimization_Threshold_{delta:.2f}", LpMaximize)

    # Variables
    elim = LpVariable.dicts(
        "Elim", [(m, i) for m in models for i in classes], cat="Binary"
    )
    x = LpVariable.dicts(
        "X",
        [(omega, m, i) for omega in objects for m in models for i in classes],
        cat="Binary",
    )
    con = LpVariable.dicts(
        "Con", [(omega, (c1, c2)) for omega in objects for (c1, c2) in ic], cat="Binary"
    )
    a = LpVariable.dicts(
        "A", [(i, omega) for i in classes for omega in objects], cat="Binary"
    )

    # Objective Function: Maximize assignments while minimizing conflicts
    alpha = 0.8  # Weight for maximizing assignments
    beta = 0.2  # Weight for minimizing conflicts
    prob += alpha * lpSum(
        a[i, omega] for i in classes for omega in objects
    ) - beta * lpSum(con[omega, p] for omega in objects for p in ic)

    # Constraints
    for omega in objects:
        for m in models:
            for i in classes:
                prob += x[omega, m, i] <= 1 - elim[m, i]  # Elimination constraint
                prob += (
                    x[omega, m, i] * pred[m, i, omega] <= a[i, omega]
                )  # Consistency constraint

        for i in classes:
            prob += a[i, omega] <= lpSum(
                x[omega, m, i] * pred[m, i, omega] for m in models
            )  # Upper bound

        for i, j in ic:
            prob += con[omega, (i, j)] >= a[i, omega] + a[j, omega] - 1
            prob += con[omega, (i, j)] <= a[i, omega]
            prob += con[omega, (i, j)] <= a[j, omega]

        prob += (
            lpSum(a[i, omega] for i in classes) >= 1
        )  # Ensure at least one class per object

    # Global conflict threshold
    prob += lpSum(con[omega, p] for omega in objects for p in ic) <= delta * lpSum(
        a[i, omega] for i in classes for omega in objects
    )

    # Solve the problem
    solver = PULP_CBC_CMD(msg=False, timeLimit=3, timeMode="elapsed")
    start_time = time.time()
    prob.solve(solver)
    end_time = time.time()
    optimization_time = end_time - start_time
    total_conflicts = value(lpSum(con[omega, p] for omega in objects for p in ic))
    # Extract assignments for this threshold
    assign_prime = []
    for i in classes:
        for omega in objects:
            for m in models:
                if a[i, omega].varValue == 1 and x[omega, m, i].varValue == 1:
                    assign_prime.append((m, omega, i))

    return [assign_prime, optimization_time]


def get_assign_set(predictions):
    assigns = {
        "pfca_normal-1000": [],
        "pfca_rain-1000": [],
        "pfca_fog-1000": [],
        "pfca_maple_leaf-1000": [],
        "pfca_snow-1000": [],
        "pfca_dust-1000": [],
    }
    for model, obj, pred, conf in predictions:
        assigns[f"pfca_{model}"].append((obj, pred, conf))
    return assigns


def optimize_int_prog(predictions, gt):
    assigns_set = get_assign_set(predictions)
    objects = list(set([obj for _, obj, _, _ in predictions]))
    # Initialize Pred dictionary
    Pred = {}
    for m in models:
        for i in classes:
            for omega in objects:
                Pred[m, i, omega] = int(
                    any(
                        (obj == omega and clas == i)
                        for obj, clas, conf in assigns_set.get(condition_map[m], [])
                    )
                )

    number_of_objects = len(objects)

    results = {}
    for delta in deltas:
        opt_assigns, running_time = optimize_for_threshold(objects, Pred, delta)
        aux_format_preds = [(obj, pred) for _, obj, pred in opt_assigns]
        aux_format_preds = list(set(aux_format_preds))
        metrics = compute_performance_metrics(gt, aux_format_preds)
        inconsistencies = compute_incon(aux_format_preds)
        aux_format_conf_preds = [(model, obj, pred) for model, obj, pred in opt_assigns]
        preds_with_conf = [
            (obj, pred, conf)
            for (model, obj, pred, conf) in predictions
            if (model, obj, pred) in aux_format_conf_preds
        ]
        tb_assignments = resolve_inconsistencies_TB(preds_with_conf)
        tb_metrics = compute_performance_metrics(gt, tb_assignments)
        tb_inconsistencies = compute_incon(tb_assignments)

        results[f"delta-{delta:.2f}"] = {
            "assignments": aux_format_preds,
            "tb_assignments": tb_assignments,
            "gt": gt,
            "time": running_time,
            "number_of_objects": number_of_objects,
        }
    return results

    # Convert results to DataFrame and save to Excel
    # results_df = pd.DataFrame(results)
    # base_filename = os.path.basename(predictions_file_name).replace('.xlsx', '')
    # results_df.to_excel(f'{output_file_path}/optimization_results_{base_filename}.xlsx', index=False)


def resolve_inconsistencies_TB(assignments):
    # Diccionario para almacenar la clase con mayor conf para cada obj
    best_class_for_obj = defaultdict(lambda: (None, -1))

    # Iterar sobre las triplas y seleccionar la clase con mayor conf para cada obj
    for obj, cls, conf in assignments:
        if conf > best_class_for_obj[obj][1]:
            best_class_for_obj[obj] = (cls, conf)

    # Crear una lista de tuplas (obj, class) con las clases seleccionadas
    resolved_tuples = [(obj, cls) for obj, (cls, conf) in best_class_for_obj.items()]

    return resolved_tuples


def process_single_image(image_path, predictions, ground_truth):
    """Process a single image and return optimization results."""
    image_predictions = [
        pred for pred in predictions if pred[1].startswith(image_path)
    ]
    image_ground_truth = set(
        [gt for gt in ground_truth if gt[0].startswith(image_path)]
    )

    return optimize_int_prog(image_predictions, image_ground_truth)


def run_optimization_for_epsilons(test_set_path, output_folder_path):
    os.makedirs(output_folder_path, exist_ok=True)

    df = pd.read_excel(
        f"{test_set_path}/epsilon-0.5.xlsx",
        usecols=["object", "ground_truth_category_id"],
    )
    ground_truth = get_ground_truth(df)
    epsilon_predictions = {}

    print("#" * 80)
    print(f"Running optimization for {test_set_path}...")
    print("#" * 80)
    for epsilon in epsilon_values:
        print(f"Running optimization for epsilon {epsilon}...")
        if (
            f"epsilon-{epsilon}-no-assignments.jsonl"
            in os.listdir(output_folder_path)
        ):
            print(f"Skipping epsilon {epsilon}...")
            continue

        print("Reading predictions...")
        df = pd.read_excel(
            f"{test_set_path}/epsilon-{epsilon}.xlsx",
            usecols=["model", "object", "after_edcr", "prediction_confidence"],
        )
        predictions = [
            (model, obj, pred, conf)
            for model, obj, pred, conf in zip(
                df["model"], df["object"], df["after_edcr"], df["prediction_confidence"]
            )
            if pd.notnull(pred)
        ]
        epsilon_predictions[epsilon] = predictions

        total_results = []
        predictions = epsilon_predictions[epsilon]

        unique_image_paths = set([obj.split("/")[0] for _, obj, _, _ in predictions])

        # Create a partial function with the constant arguments
        process_func = partial(process_single_image,
                               predictions=predictions,
                               ground_truth=ground_truth)

        # Create a multiprocessing pool
        with mp.Pool(processes=mp.cpu_count()) as pool:
            # Use imap to process images in parallel with a progress bar
            total_results = list(tqdm.tqdm(
                pool.imap(process_func, unique_image_paths),
                total=len(unique_image_paths),
                desc="Processing images"
            ))

        combined_results = optimize_int_prog(predictions, ground_truth)

        metrics_results = []
        for delta in deltas:
            total_assignments = [
                result[f"delta-{delta:.2f}"]["assignments"] for result in total_results
            ]
            total_tb_assignments = [
                result[f"delta-{delta:.2f}"]["tb_assignments"]
                for result in total_results
            ]

            # Flatten the list of lists
            total_assignments = [
                item for sublist in total_assignments for item in sublist
            ]
            total_tb_assignments = [
                item for sublist in total_tb_assignments for item in sublist
            ]
            metrics = calculate_metrics(total_assignments, ground_truth)
            tb_metrics = calculate_metrics(total_tb_assignments, ground_truth)
            metrics_results.append(
                {
                    "epsilon": epsilon,
                    "delta": delta,
                    "metrics": metrics,
                    "tb_metrics": tb_metrics,
                }
            )

        metric_results_df = pd.DataFrame(metrics_results)
        metric_results_df.to_json(
            f"{output_folder_path}/epsilon-{epsilon}-metrics.jsonl",
            orient="records",
            lines=True,
        )

        total_results_df = pd.DataFrame(total_results)
        total_results_df.to_json(
            f"{output_folder_path}/epsilon-{epsilon}.jsonl",
            orient="records",
            lines=True,
        )

        def remove_assignments(x):
            del x["assignments"]
            del x["tb_assignments"]
            del x["gt"]

            return x

        total_results_df = total_results_df.applymap(remove_assignments)
        total_results_df.to_json(
            f"{output_folder_path}/epsilon-{epsilon}-no-assignments.jsonl",
            orient="records",
            lines=True,
        )


def calculate_metrics(assignments, gt):
    metrics = compute_performance_metrics(gt, assignments)
    inconsistencies = compute_incon(assignments)
    return {
        "precision": metrics["precision"],
        "recall": metrics["recall"],
        "f1": metrics["f1"],
        "accuracy": metrics["accuracy"],
        "inconsistencies": inconsistencies,
    }



def main():
    CLOUD_FOLDER = r''
    DATA_FOLDER = os.path.expanduser(rf'{CLOUD_FOLDER}/data')

    run_optimization_for_epsilons(test_set_path=fr"{DATA_FOLDER}/test-1",
                                  output_folder_path="output/test-1")
    # run_optimization_for_epsilons("data/test-2", "output/test-2")
    # run_optimization_for_epsilons("data/test-3", "output/test-3")


if __name__ == "__main__":
    mp.freeze_support()
    main()