Blablablab/audio-classification
0
1"""2Export CLI3 4Command-line interface for exporting Potato annotations to various formats.5 6Usage:7 python -m potato.export --config config.yaml --format coco --output ./out/8 python -m potato.export --config config.yaml --format conll_2003 --output ./out/9 python -m potato.export --list-formats10"""11 12import argparse13import json14import os15import sys16import logging17import glob18 19import yaml20 21from .base import ExportContext22from .registry import export_registry23 24logger = logging.getLogger(__name__)25 26 27def load_annotations_from_output_dir(output_dir: str, schemas: list) -> list:28 """29 Load user annotations from the Potato output directory.30 31 Reads user_state.json files from each user subdirectory32 and flattens annotations into a list of records.33 34 Args:35 output_dir: Path to the annotation output directory36 schemas: List of annotation scheme configs37 38 Returns:39 List of annotation dicts40 """41 annotations = []42 43 if not os.path.isdir(output_dir):44 logger.warning(f"Output directory not found: {output_dir}")45 return annotations46 47 for user_dir in sorted(os.listdir(output_dir)):48 user_path = os.path.join(output_dir, user_dir)49 if not os.path.isdir(user_path):50 continue51 52 state_file = os.path.join(user_path, "user_state.json")53 if not os.path.exists(state_file):54 continue55 56 with open(state_file, "r") as f:57 user_state = json.load(f)58 59 user_id = user_state.get("user_id", user_dir)60 61 # Extract label annotations62 label_data = user_state.get("instance_id_to_label_to_value", {})63 span_data = user_state.get("instance_id_to_span_to_value", {})64 65 # Collect all instance IDs66 all_instances = set(label_data.keys()) | set(span_data.keys())67 68 for instance_id in all_instances:69 # Labels may be stored as a list of [[{schema, name}, value], ...]70 # or as a dict of {schema_name: {label_name: value}}.71 # Normalize to dict format.72 raw_labels = label_data.get(instance_id, {})73 if isinstance(raw_labels, list):74 labels_dict = {}75 for entry in raw_labels:76 if isinstance(entry, (list, tuple)) and len(entry) == 2:77 label_obj, value = entry78 if isinstance(label_obj, dict):79 schema = label_obj.get("schema", "")80 name = label_obj.get("name", "")81 else:82 schema, name = str(label_obj), ""83 labels_dict.setdefault(schema, {})[name] = value84 raw_labels = labels_dict85 86 record = {87 "instance_id": instance_id,88 "user_id": user_id,89 "labels": raw_labels,90 "spans": {},91 "links": {},92 "image_annotations": {},93 }94 95 # Process span data96 instance_spans = span_data.get(instance_id, {})97 for schema_name, span_list in instance_spans.items():98 if isinstance(span_list, list):99 record["spans"][schema_name] = span_list100 elif isinstance(span_list, dict):101 # Span data might be stored as a dict of span_id -> span_obj102 record["spans"][schema_name] = list(span_list.values())103 104 # Extract image annotations from labels105 # Image annotations are stored as JSON strings in label values106 for schema_name, label_dict in record["labels"].items():107 schema_config = _find_schema(schemas, schema_name)108 if schema_config and schema_config.get("annotation_type") == "image_annotation":109 # Image annotation data is stored in the label value110 for label_key, value in label_dict.items():111 if isinstance(value, str):112 try:113 parsed = json.loads(value)114 if isinstance(parsed, list):115 record["image_annotations"][schema_name] = parsed116 except (json.JSONDecodeError, TypeError):117 pass118 elif isinstance(value, list):119 record["image_annotations"][schema_name] = value120 121 annotations.append(record)122 123 return annotations124 125 126def load_phase_responses_from_output_dir(output_dir: str) -> list:127 """128 Load phase/surveyflow responses from the Potato output directory.129 130 Reads phase_to_page_to_label_to_value from each user's user_state.json131 and flattens into a list of records.132 133 Returns:134 List of dicts with keys: user_id, phase, page, schema, label_name, value135 """136 responses = []137 138 if not os.path.isdir(output_dir):139 return responses140 141 for user_dir in sorted(os.listdir(output_dir)):142 user_path = os.path.join(output_dir, user_dir)143 if not os.path.isdir(user_path):144 continue145 146 state_file = os.path.join(user_path, "user_state.json")147 if not os.path.exists(state_file):148 continue149 150 with open(state_file, "r") as f:151 user_state = json.load(f)152 153 user_id = user_state.get("user_id", user_dir)154 phase_data = user_state.get("phase_to_page_to_label_to_value", {})155 156 for phase, pages in phase_data.items():157 for page, label_values in pages.items():158 # label_values is a list of [[{schema, name}, value], ...]159 if isinstance(label_values, list):160 for entry in label_values:161 if isinstance(entry, (list, tuple)) and len(entry) == 2:162 label_obj, value = entry163 if isinstance(label_obj, dict):164 schema = label_obj.get("schema", "")165 label_name = label_obj.get("name", "")166 else:167 schema, label_name = str(label_obj), ""168 responses.append({169 "user_id": user_id,170 "phase": phase,171 "page": page,172 "schema": schema,173 "label_name": label_name,174 "value": value,175 })176 elif isinstance(label_values, dict):177 for label_obj, value in label_values.items():178 responses.append({179 "user_id": user_id,180 "phase": phase,181 "page": page,182 "schema": str(label_obj),183 "label_name": "",184 "value": value,185 })186 187 return responses188 189 190def load_items_from_data_files(config: dict, config_dir: str) -> dict:191 """192 Load item data from the data files specified in config.193 194 Args:195 config: Full Potato configuration dict196 config_dir: Directory containing the config file197 198 Returns:199 Dict mapping instance_id -> item data200 """201 items = {}202 item_props = config.get("item_properties", {})203 id_key = item_props.get("id_key", "id")204 205 data_files = config.get("data_files", [])206 if isinstance(data_files, str):207 data_files = [data_files]208 209 task_dir = config.get("task_dir", ".")210 base_dir = os.path.normpath(os.path.join(config_dir, task_dir))211 212 for data_file_entry in data_files:213 if isinstance(data_file_entry, dict):214 path = data_file_entry.get("path", "")215 else:216 path = str(data_file_entry)217 218 if not os.path.isabs(path):219 path = os.path.join(base_dir, path)220 221 if not os.path.exists(path):222 logger.warning(f"Data file not found: {path}")223 continue224 225 with open(path, "r") as f:226 for line_num, line in enumerate(f, 1):227 line = line.strip()228 if not line:229 continue230 try:231 item = json.loads(line)232 item_id = str(item.get(id_key, f"item_{line_num}"))233 items[item_id] = item234 except json.JSONDecodeError:235 # Try CSV/TSV236 logger.debug(f"Line {line_num} in {path} is not JSON, skipping")237 238 return items239 240 241def _find_schema(schemas: list, name: str) -> dict:242 """Find a schema config by name."""243 for s in schemas:244 if s.get("name") == name:245 return s246 return {}247 248 249def build_export_context(config_path: str) -> ExportContext:250 """251 Build an ExportContext from a Potato config file.252 253 Args:254 config_path: Path to YAML config file255 256 Returns:257 ExportContext ready for export258 """259 config_path = os.path.abspath(config_path)260 config_dir = os.path.dirname(config_path)261 262 with open(config_path, "r") as f:263 config = yaml.safe_load(f)264 265 schemas = config.get("annotation_schemes", [])266 267 # Determine output directory268 task_dir = config.get("task_dir", ".")269 base_dir = os.path.normpath(os.path.join(config_dir, task_dir))270 output_annotation_dir = config.get(271 "output_annotation_dir",272 os.path.join(base_dir, "annotation_output")273 )274 if not os.path.isabs(output_annotation_dir):275 output_annotation_dir = os.path.join(base_dir, output_annotation_dir)276 277 items = load_items_from_data_files(config, config_dir)278 annotations = load_annotations_from_output_dir(output_annotation_dir, schemas)279 phase_responses = load_phase_responses_from_output_dir(output_annotation_dir)280 281 return ExportContext(282 config=config,283 annotations=annotations,284 items=items,285 schemas=schemas,286 output_dir=output_annotation_dir,287 phase_responses=phase_responses,288 )289 290 291def main():292 parser = argparse.ArgumentParser(293 description="Export Potato annotations to standard formats"294 )295 parser.add_argument(296 "--config", "-c",297 help="Path to Potato YAML config file",298 )299 parser.add_argument(300 "--format", "-f",301 help="Export format (e.g., coco, yolo, pascal_voc, conll_2003, conll_u)",302 )303 parser.add_argument(304 "--output", "-o",305 help="Output directory",306 default="./export_output",307 )308 parser.add_argument(309 "--list-formats",310 action="store_true",311 help="List available export formats and exit",312 )313 parser.add_argument(314 "--option",315 action="append",316 default=[],317 help="Format-specific option as key=value (can be repeated)",318 )319 parser.add_argument(320 "--verbose", "-v",321 action="store_true",322 help="Enable verbose logging",323 )324 325 args = parser.parse_args()326 327 logging.basicConfig(328 level=logging.DEBUG if args.verbose else logging.INFO,329 format="%(levelname)s: %(message)s",330 )331 332 if args.list_formats:333 formats = export_registry.list_exporters()334 if not formats:335 print("No export formats registered.")336 else:337 print("Available export formats:\n")338 for fmt in formats:339 exts = ", ".join(fmt["file_extensions"])340 print(f" {fmt['format_name']:15s} {fmt['description']}")341 print(f" {'':15s} Extensions: {exts}")342 print()343 return344 345 if not args.config:346 parser.error("--config is required (unless using --list-formats)")347 if not args.format:348 parser.error("--format is required (unless using --list-formats)")349 350 if not os.path.exists(args.config):351 print(f"Error: Config file not found: {args.config}", file=sys.stderr)352 sys.exit(1)353 354 # Parse options355 options = {}356 for opt in args.option:357 if "=" in opt:358 k, v = opt.split("=", 1)359 options[k.strip()] = v.strip()360 361 # Build context362 print(f"Loading config from: {args.config}")363 context = build_export_context(args.config)364 print(f"Loaded {len(context.items)} items, {len(context.annotations)} annotations")365 366 # Export367 print(f"Exporting to {args.format} format...")368 result = export_registry.export(args.format, context, args.output, options)369 370 if result.success:371 print(f"\nExport successful!")372 print(f"Files written:")373 for f in result.files_written:374 print(f" {f}")375 if result.stats:376 print(f"\nStatistics:")377 for k, v in result.stats.items():378 print(f" {k}: {v}")379 else:380 print(f"\nExport failed!", file=sys.stderr)381 for err in result.errors:382 print(f" ERROR: {err}", file=sys.stderr)383 384 if result.warnings:385 print(f"\nWarnings:")386 for w in result.warnings:387 print(f" WARNING: {w}")388 389 sys.exit(0 if result.success else 1)390 391 392if __name__ == "__main__":393 main()394