Spaces:
Runtime error
Runtime error
| import json | |
| from argparse import ArgumentParser | |
| from pathlib import Path | |
| def convert_thermal_mot_to_tao(dataset_dir, output_json): | |
| """ | |
| Converts a thermal MOT Challenge dataset (COCO-like per video) to a single TAO JSON file. | |
| Args: | |
| dataset_dir (str or Path): Path to the main dataset directory (contains video subdirs). | |
| output_json (str or Path): Path to save the resulting TAO JSON. | |
| """ | |
| dataset_dir = Path(dataset_dir) | |
| tao_dict = { | |
| "videos": [], | |
| "images": [], | |
| "annotations": [], | |
| "categories": [{"id": 0, "name": "person"}], | |
| "tracks": [], | |
| } | |
| global_image_id = 1 | |
| global_ann_id = 1 | |
| global_track_id = 1 | |
| track_mapping = {} # map (video_name, original_track_id) -> global_track_id | |
| video_mapping = {} | |
| for video_idx, video_dir in enumerate(sorted(dataset_dir.iterdir())): | |
| if not video_dir.is_dir(): | |
| continue | |
| annotation_path = video_dir / "thermal" / "COCO" / "annotations.json" | |
| if not annotation_path.exists(): | |
| print(f"Warning: {annotation_path} does not exist, skipping.") | |
| continue | |
| with open(annotation_path) as f: | |
| data = json.load(f) | |
| # Verify category assumption | |
| input_categories = {cat['id']: cat['name'] for cat in data.get("categories", [])} | |
| if len(input_categories) != 1 or 1 not in input_categories or input_categories[1] != "person": | |
| raise ValueError(f"Video {video_dir.name} has unexpected categories: {input_categories}") | |
| # Add video entry | |
| video_id = video_idx + 1 | |
| tao_dict["videos"].append({ | |
| "id": video_id, | |
| "name": video_dir.name, | |
| "neg_category_ids": list(range(1, 80)), | |
| "not_exhaustive_category_ids": [], | |
| }) | |
| video_mapping[video_dir.name] = video_id | |
| # Map original image_id to global_image_id | |
| img_id_mapping = {} | |
| for img in sorted(data["images"], key=lambda x: x["id"]): | |
| old_img_id = img["id"] | |
| img_id_mapping[old_img_id] = global_image_id | |
| tao_dict["images"].append({ | |
| "id": global_image_id, | |
| "video_id": video_id, | |
| "frame_index": old_img_id, # keep original frame index | |
| "file_name": img["file_name"] | |
| }) | |
| global_image_id += 1 | |
| # Convert annotations | |
| for ann in data["annotations"]: | |
| old_track_id = ann["track_id"] | |
| key = (video_dir.name, old_track_id) | |
| if key not in track_mapping: | |
| track_mapping[key] = global_track_id | |
| tao_dict["tracks"].append({ | |
| "id": global_track_id, | |
| "category_id": 0, | |
| "video_id": video_mapping[video_dir.name] | |
| }) | |
| global_track_id += 1 | |
| tao_dict["annotations"].append({ | |
| "id": global_ann_id, | |
| "image_id": img_id_mapping[ann["image_id"]], | |
| "category_id": 0, | |
| "bbox": ann["bbox"], | |
| "track_id": track_mapping[key], | |
| "video_id": video_id, | |
| }) | |
| global_ann_id += 1 | |
| # Save final TAO JSON | |
| with open(output_json, "w") as f: | |
| json.dump(tao_dict, f, indent=2) | |
| print(f"TAO JSON saved to {output_json}.") | |
| print( | |
| f"Total videos: {len(tao_dict['videos'])}, images: {len(tao_dict['images'])}, annotations: {len(tao_dict['annotations'])}") | |
| if __name__ == '__main__': | |
| parser = ArgumentParser() | |
| parser.add_argument("tmot_train_ann_dir") | |
| parser.add_argument("output_json") | |
| args = parser.parse_args() | |
| convert_thermal_mot_to_tao(args.tmot_train_ann_dir, args.output_json) | |