File size: 5,871 Bytes
1da285f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | import argparse
import json
from pathlib import Path
TASKS = ("semantics", "interactable", "interaction")
def extract_image_id(image_name):
stem = Path(image_name).stem
parts = stem.split("_")
if len(parts) != 2:
raise ValueError(f"Cannot infer image_id from image name: {image_name}")
app_id, frame_id = parts
return int(f"{app_id}{int(frame_id):03d}")
def load_questions(path):
if path is None:
return {}
mapping = {}
with Path(path).open() as file:
for line_number, line in enumerate(file, start=1):
line = line.strip()
if not line:
continue
question = json.loads(line)
try:
question_id = str(question["question_id"])
image_id = question.get("image_id")
if image_id is None:
image_id = extract_image_id(question["image"])
except KeyError as exc:
raise ValueError(
f"Question file {path} line {line_number} is missing {exc.args[0]!r}"
) from exc
mapping[question_id] = image_id
return mapping
def normalize_bbox(item):
bbox = item.get("bbox", item.get("bbox_pixels"))
if bbox is None:
raise ValueError(f"Prediction is missing bbox/bbox_pixels: {item}")
if len(bbox) != 4:
raise ValueError(f"Prediction bbox must have four values: {item}")
return bbox
def normalize_score(item):
return item.get("score", item.get("probability", 1.0))
def normalize_category(category, item, task):
if task == "interactable":
return 1
category_id = item.get("category_id", category)
if category_id is None:
raise ValueError(f"Prediction is missing category/category_id for {task}: {item}")
return category_id
def normalize_image_id(question_id, content, item, questions):
for source in (item, content):
if isinstance(source, dict) and "image_id" in source:
return source["image_id"]
if isinstance(source, dict) and "image" in source:
return extract_image_id(source["image"])
if question_id is not None and str(question_id) in questions:
return questions[str(question_id)]
raise ValueError(
"Prediction is missing image_id/image. Provide --questions for old "
f"question-keyed prediction files. question_id={question_id!r}"
)
def iter_old_format(data):
for question_id, content in data.items():
if not isinstance(content, dict):
raise ValueError(f"Prediction for question {question_id!r} must be an object")
results = content.get("oovd_result")
if results is None:
continue
if not isinstance(results, dict):
raise ValueError(f"oovd_result for question {question_id!r} must be an object")
for category, objects in results.items():
if not isinstance(objects, list):
raise ValueError(
f"oovd_result[{category!r}] for question {question_id!r} must be a list"
)
for item in objects:
if not isinstance(item, dict):
raise ValueError(f"Prediction item must be an object: {item!r}")
yield question_id, content, category, item
def iter_prediction_items(data):
if isinstance(data, dict):
yield from iter_old_format(data)
return
if not isinstance(data, list):
raise ValueError("Prediction input must be a list or a question-keyed object")
for item in data:
if not isinstance(item, dict):
raise ValueError(f"Prediction item must be an object: {item!r}")
yield None, {}, item.get("category_id"), item
def convert_predictions(input_path, task, questions_path=None):
with Path(input_path).open() as file:
data = json.load(file)
questions = load_questions(questions_path)
output = []
for question_id, content, category, item in iter_prediction_items(data):
output.append(
{
"image_id": normalize_image_id(question_id, content, item, questions),
"category_id": normalize_category(category, item, task),
"bbox": normalize_bbox(item),
"score": normalize_score(item),
}
)
return output
def output_path_for_task(output_path, task):
path = Path(output_path)
if path.suffix:
return path.with_name(f"{path.stem}_{task}{path.suffix}")
return path / f"{task}.json"
def write_json(path, data):
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("w") as file:
json.dump(data, file, indent=2)
file.write("\n")
def build_parser():
parser = argparse.ArgumentParser(
description="Convert Orienter prediction files to COCO-style result JSON."
)
parser.add_argument(
"--task",
choices=TASKS + ("all",),
required=True,
help="Evaluation task to convert for.",
)
parser.add_argument("--input", required=True, help="Prediction JSON path.")
parser.add_argument(
"--questions",
help="Question JSONL path. Required only for old question-keyed inputs without image_id.",
)
parser.add_argument("--output", required=True, help="Output JSON path or directory.")
return parser
def main(argv=None):
args = build_parser().parse_args(argv)
tasks = TASKS if args.task == "all" else (args.task,)
for task in tasks:
converted = convert_predictions(args.input, task, args.questions)
output_path = (
output_path_for_task(args.output, task) if args.task == "all" else Path(args.output)
)
write_json(output_path, converted)
if __name__ == "__main__":
main()
|