create_plots.py

from pathlib import Path
from turtle import shape
from plotnine import (
    ggplot,
    aes,
    geom_line,
    geom_point,
    facet_grid,
    facet_wrap,
    scale_y_continuous,
    geom_hline,
    position_dodge,
    geom_errorbar,
    theme,
    element_text,
    ylab,
    xlab,
    scale_color_discrete,
    scale_fill_manual
)
from plotnine.data import economics
from pandas import Categorical, DataFrame
from plotnine.scales.limits import ylim
from plotnine.scales.scale_xy import scale_x_discrete
from plotnine.guides import guide_axis, guide_legend, guide
from glob import glob
import re
import fire

# files = glob("run/models/*/eval_*.txt")
# files = glob("run/models/tanet_*/eval_*.txt") + glob("run/models/vnn_tanet_*/eval_*.txt") + glob("run/models/vnnclass*/eval_*.txt")
# files = glob("run/models/*pointpillars*/eval_*.txt")
# files = glob("run/models/tanet_*/eval_*.txt") + glob("run/models/vnnclass*/eval_*.txt")
# files = glob("run/models/*lens*/eval_*.txt")


def read_data(file, frame):
    with open(file, "r") as f:
        lines = f.readlines()

        model_name = Path(file).parent.name

        if "vnnclass_" in model_name:
            model_name=model_name.replace("car", "tanet_car")

        words = model_name.split("_")
        agent = words[0] if len(words) > 2 else "classic"
        network = "_".join(model_name.split("_")[1:])
        method = words[1] if len(words) > 2 else words[0]

        if "vnnclass1" or "vnnclass3" in model_name:
            method=words[-3]

        if model_name in ["tanet_car_2", "tanet_car_3"]:
            agent="classic"
            network="tanet"
            method="tanet"
            

        last_samples = -1
        last_easy_ap = -1
        last_moderate_ap = -1
        last_hard_ap = -1
        last_mAP = -1

        for line in lines:

            if "samples" in line:
                last_samples = int(line.split(" = ")[-1])

            if "3d AP" in line:
                number_line = line.split("[")[-1].split("]")[0]
                numbers = number_line.split()
                values = [float(x.replace(",", "")) for x in numbers]
                if len(values) < 3 or values[0] < 0:
                    continue
                last_easy_ap, last_moderate_ap, last_hard_ap = values

                last_mAP = (last_easy_ap + last_moderate_ap + last_hard_ap) / 3

                frame["model"].append(model_name)
                frame["agent"].append(agent)
                frame["network"].append(network)
                frame["method"].append(method)
                frame["samples"].append(last_samples)
                frame["easy_ap"].append(last_easy_ap)
                frame["moderate_ap"].append(last_moderate_ap)
                frame["hard_ap"].append(last_hard_ap)
                frame["mAP"].append(last_mAP)

                print(file, "s", last_samples)

                last_samples = -1
                last_easy_ap = -1
                last_moderate_ap = -1
                last_hard_ap = -1
                last_mAP = -1



def plot_single_frame(frame, output_file_name, y="mAP", map_min=None, map_max=None):

    plot = (
        ggplot(frame)
        + aes(x="samples", y=y)
        + facet_wrap(["method", "model"], nrow=(len(frame["model"].unique()) // 6 ) + 1)
        + geom_point(
            aes(colour="agent", shape="method"),
            size=3,
            position=position_dodge(width=0.8),
            stroke=0.2,
        )
        # + scale_fill_manual(values=["purple", "yellow", "green", "cyan", "red"])
    )

    if map_max is not None:
        plot += ylim(map_min, map_max)

    plot = plot + theme(strip_text_x=element_text(size=4), axis_title=element_text(size=10), axis_text=element_text(size=5), figure_size=(14, 6))

    plot.save("run/plots/" + output_file_name + ".png", dpi=600)
    frame.to_csv("run/plots/" + output_file_name + ".csv")


def plot_all_single_frames(files, map_min=60, map_max=90):

    frame = {
        "model": [],
        "agent": [],
        "network": [],
        "method": [],
        "samples": [],
        "easy_ap": [],
        "moderate_ap": [],
        "hard_ap": [],
        "mAP": [],
    }

    for file in files:

        read_data(file, frame)

    data_frame = DataFrame(frame)
    # data_frame.sort_values(by=["method"], inplace=True)
    data_frame["method"] = Categorical(data_frame["method"], ["pointpillars", "tanet", "tapp"])

    for metric in ["mAP", "easy_ap", "moderate_ap", "hard_ap"]:
        plot_single_frame(
            data_frame,
            "plot_" + metric,
            y=metric,
            map_min=map_min,
            map_max=map_max,
        )
    
    data_frame.sort_values(by=["mAP"], inplace=True)
    i = 0
    s = 10
    with open("run/plots/plot_map.df.txt", "w") as f:
        while i < len(data_frame):
            print(data_frame[i:min(i + s, len(data_frame))])
            print(data_frame[i:min(i + s, len(data_frame))], file=f)
            i += s

def main(template="run/models/*/eval_*.txt", map_min=0, map_max=100):

    if not isinstance(template, list):
        template = [template]

    files = sum([glob(a) for a in template[1:]], glob(template[0]))

    plot_all_single_frames(files, map_min, map_max)

if __name__ == "__main__":
    fire.Fire(main)