|
| 1 | +""" |
| 2 | +COCO import run by the DataTorch server: the JSON is uploaded to the |
| 3 | +project's default storage and the server's import job (createImport) maps |
| 4 | +its images to dataset files by name, creates missing labels (with their |
| 5 | +keypoint names and skeleton) and writes boxes, polygons and keypoints. |
| 6 | +""" |
| 7 | + |
| 8 | +import os |
| 9 | +import re |
| 10 | +import time |
| 11 | +import uuid |
| 12 | +import logging |
| 13 | +from typing import List, Optional |
| 14 | + |
| 15 | +import requests |
| 16 | + |
| 17 | +from .. import ApiClient, Dataset, Project, Where |
| 18 | + |
| 19 | +_LOGGER = logging.getLogger(__name__) |
| 20 | + |
| 21 | +_CREATE_IMPORT = """ |
| 22 | + mutation CreateImport($input: CreateImportInput!) { |
| 23 | + createImport(input: $input) { id } |
| 24 | + } |
| 25 | +""" |
| 26 | + |
| 27 | +_GET_IMPORT = """ |
| 28 | + query GetImport($id: ID!) { |
| 29 | + import(id: $id) { |
| 30 | + annotationsCreated |
| 31 | + labelsCreated |
| 32 | + filesProcessed |
| 33 | + job { state logs { level message } } |
| 34 | + } |
| 35 | + } |
| 36 | +""" |
| 37 | + |
| 38 | +_DONE_STATES = ("SUCCESS", "FAILED", "CANCELED") |
| 39 | + |
| 40 | + |
| 41 | +class CocoImportError(Exception): |
| 42 | + pass |
| 43 | + |
| 44 | + |
| 45 | +def _plain(message: str) -> str: |
| 46 | + """Job log messages carry <b> tags for the web UI.""" |
| 47 | + return re.sub(r"<[^>]+>", "", message) |
| 48 | + |
| 49 | + |
| 50 | +def _pick_dataset(project: Project, dataset_name: Optional[str]) -> Dataset: |
| 51 | + if dataset_name: |
| 52 | + return project.dataset(dataset_name) |
| 53 | + datasets = project.datasets() |
| 54 | + if len(datasets) == 1: |
| 55 | + return datasets[0] |
| 56 | + names = [d.name for d in datasets] |
| 57 | + raise CocoImportError( |
| 58 | + f"Project has {len(datasets)} datasets, pass one with --dataset: {names}" |
| 59 | + ) |
| 60 | + |
| 61 | + |
| 62 | +def _upload(api: ApiClient, project: Project, file_path: str) -> str: |
| 63 | + """ |
| 64 | + Uploads the COCO file as a project file and returns its id. The upload |
| 65 | + route returns no id and renames on a path clash, so the file goes into a |
| 66 | + fresh folder and is found again by its exact path. |
| 67 | + """ |
| 68 | + # The server reads the file through whichever link holds it, so any link |
| 69 | + # works: DataTorch Storage when the project has it, else the first one. |
| 70 | + links = project.storage_links() |
| 71 | + if not links: |
| 72 | + raise CocoImportError("Project has no storage to upload the COCO file to") |
| 73 | + link = next((l for l in links if l.name == "DataTorch Storage"), links[0]) |
| 74 | + folder = f"imports/{uuid.uuid4().hex}" |
| 75 | + name = os.path.basename(file_path) |
| 76 | + with open(file_path, "rb") as f: |
| 77 | + r = requests.post( |
| 78 | + f"{api.api_url}/file/v1/upload/{link.id}", |
| 79 | + params={"path": folder, "import": "true"}, |
| 80 | + files={"file": (name, f, "application/json")}, |
| 81 | + headers={api.token_header: api._api_token}, |
| 82 | + ) |
| 83 | + if not r.ok: |
| 84 | + raise CocoImportError(f"Upload failed ({r.status_code}): {r.text[:300]}") |
| 85 | + |
| 86 | + files = project.files(Where(path=f"{folder}/{name}"), limit=2) |
| 87 | + if len(files) != 1: |
| 88 | + raise CocoImportError(f"Uploaded {name} but could not find it in the project") |
| 89 | + return files[0].id |
| 90 | + |
| 91 | + |
| 92 | +def import_coco_server( |
| 93 | + file_path: str, |
| 94 | + project_string: str, |
| 95 | + dataset_name: Optional[str] = None, |
| 96 | + api: ApiClient = None, |
| 97 | + timeout: float = 3600, |
| 98 | + poll_seconds: float = 2, |
| 99 | +) -> dict: |
| 100 | + """ |
| 101 | + Imports a COCO file into a dataset on the server and waits for the job. |
| 102 | + Returns the import's counts and its warning messages; raises |
| 103 | + CocoImportError with the job's error messages if it fails. |
| 104 | + """ |
| 105 | + if not os.path.isfile(file_path): |
| 106 | + raise ValueError(f"Provided path '{file_path}' is not a file.") |
| 107 | + |
| 108 | + if api is None: |
| 109 | + api = ApiClient() |
| 110 | + |
| 111 | + project: Project = api.project(project_string) |
| 112 | + dataset = _pick_dataset(project, dataset_name) |
| 113 | + |
| 114 | + _LOGGER.info("Uploading %s...", os.path.basename(file_path)) |
| 115 | + file_id = _upload(api, project, file_path) |
| 116 | + |
| 117 | + import_id = api.execute( |
| 118 | + _CREATE_IMPORT, |
| 119 | + params={ |
| 120 | + "input": {"datasetId": dataset.id, "fileId": file_id, "format": "COCO"} |
| 121 | + }, |
| 122 | + )["createImport"]["id"] |
| 123 | + _LOGGER.info("Import %s queued for dataset '%s'.", import_id, dataset.name) |
| 124 | + |
| 125 | + deadline = time.time() + timeout |
| 126 | + while True: |
| 127 | + imp = api.execute(_GET_IMPORT, params={"id": import_id})["import"] |
| 128 | + state = imp["job"]["state"] |
| 129 | + if state in _DONE_STATES: |
| 130 | + break |
| 131 | + if time.time() > deadline: |
| 132 | + raise CocoImportError( |
| 133 | + f"Timed out waiting for import {import_id} (last state {state})" |
| 134 | + ) |
| 135 | + time.sleep(poll_seconds) |
| 136 | + |
| 137 | + logs = imp["job"]["logs"] or [] |
| 138 | + |
| 139 | + def messages(level: str) -> List[str]: |
| 140 | + return [_plain(log["message"]) for log in logs if log["level"] == level] |
| 141 | + |
| 142 | + if state != "SUCCESS": |
| 143 | + # The job logs one failure at several layers; the last is the cause. |
| 144 | + errors = messages("error") or [f"Import job ended {state}"] |
| 145 | + raise CocoImportError(errors[-1]) |
| 146 | + |
| 147 | + return { |
| 148 | + "import_id": import_id, |
| 149 | + "dataset": dataset.name, |
| 150 | + "files_processed": imp["filesProcessed"], |
| 151 | + "labels_created": imp["labelsCreated"], |
| 152 | + "annotations_created": imp["annotationsCreated"], |
| 153 | + "warnings": messages("warn"), |
| 154 | + } |
0 commit comments