| import argparse |
| import json |
| import os |
| import re |
|
|
| import numpy as np |
| from data.data_loader import load_tvgbench_filter |
|
|
|
|
| def compute_IoU(pred, gt): |
| """Compute the IoU given predicted and ground truth windows.""" |
| assert isinstance(pred, list) and isinstance(gt, list) |
| pred_is_list = isinstance(pred[0], list) |
| gt_is_list = isinstance(gt[0], list) |
| if not pred_is_list: |
| pred = [pred] |
| if not gt_is_list: |
| gt = [gt] |
| pred, gt = np.array(pred), np.array(gt) |
| inter_left = np.maximum(pred[:, 0, None], gt[None, :, 0]) |
| inter_right = np.minimum(pred[:, 1, None], gt[None, :, 1]) |
| inter = np.maximum(0.0, inter_right - inter_left) |
| union_left = np.minimum(pred[:, 0, None], gt[None, :, 0]) |
| union_right = np.maximum(pred[:, 1, None], gt[None, :, 1]) |
| union = np.maximum(0.0, union_right - union_left) |
| overlap = 1.0 * inter / union |
| if not gt_is_list: |
| overlap = overlap[:, 0] |
| if not pred_is_list: |
| overlap = overlap[0] |
| return overlap |
|
|
|
|
| def calc_difficulty(pred, gt): |
| if None in pred: |
| return 0.0 |
| return compute_IoU(pred, gt) * 100.0 |
|
|
|
|
| def extract_answer_force(output_string): |
| number_pattern = r"\d+(?:\.\d+)?" |
| matches = re.findall(number_pattern, output_string) |
| output = [float(num) for num in matches[:2]] |
| if len(output) == 2: |
| return output |
| return [None, None] |
|
|
|
|
| def load_new_data(data_dir): |
| data = {} |
| for file in os.listdir(data_dir): |
| if "jsonl" not in file: |
| continue |
| file_path = os.path.join(data_dir, file) |
| for line in open(file_path): |
| tmp = json.loads(line) |
| if ( |
| None in tmp["pred"] |
| ): |
| tmp["pred"] = extract_answer_force(tmp["output_text"]) |
| data[tmp["qid"]] = { |
| "difficulty": calc_difficulty(tmp["pred"], tmp["target"]), |
| "pred": tmp["pred"], |
| } |
|
|
| return data |
|
|
|
|
| def calc_score(difficulty_data_dict): |
| data = list(difficulty_data_dict.values()) |
| for i in [30.0, 50.0, 70.0]: |
| cnt = len([itm for itm in data if itm["difficulty"] > i]) |
| score = round(cnt / len(difficulty_data_dict) * 100, 1) |
| print(score) |
|
|
|
|
| def main(input_dir=None, split="p03c_v2_top_2500", output_dir=None): |
| original_data = load_tvgbench_filter(split=split) |
| difficulty_data_dict = load_new_data(f"./{input_dir}") |
| |
|
|
| print(len(difficulty_data_dict)) |
| calc_score(difficulty_data_dict) |
|
|
| new_data = [] |
| for itm in original_data: |
| if itm["qid"] in difficulty_data_dict: |
| itm["difficulty"] = difficulty_data_dict[itm["qid"]]["difficulty"] |
| itm["pred"] = difficulty_data_dict[itm["qid"]]["pred"] |
| new_data.append(itm) |
|
|
| if len(new_data) != len(original_data): |
| print("Not All!! Attention!!") |
|
|
| if not os.path.exists(f"{output_dir}/{input_dir}"): |
| os.makedirs(f"{output_dir}/{input_dir}") |
| path_name = f"{output_dir}/{input_dir}/train_v4_cloud.json" |
|
|
| with open(path_name, "w") as f: |
| json.dump(new_data, f) |
|
|
| print(len(new_data)) |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description="") |
| parser.add_argument("--input", help="logging file") |
| parser.add_argument("--split", help="annotation loading") |
| parser.add_argument("--output_dir") |
| args = parser.parse_args() |
| main(input_dir=args.input, split=args.split, output_dir=args.output_dir) |
|
|
