|
| 1 | +import glob |
1 | 2 | import json |
2 | 3 | import logging |
3 | 4 | from pathlib import Path |
4 | | -import glob |
5 | 5 |
|
6 | 6 | import pandas as pd |
| 7 | + |
7 | 8 | from ..exceptions import SABaseException |
| 9 | + |
8 | 10 | logger = logging.getLogger("superannotate-python-sdk") |
9 | 11 |
|
10 | 12 |
|
@@ -32,7 +34,12 @@ def df_to_annotations(df, output_dir): |
32 | 34 | image_height = None |
33 | 35 | image_width = None |
34 | 36 | image_df = df[df["imageName"] == image] |
35 | | - image_annotation = [] |
| 37 | + image_annotation = { |
| 38 | + "instances": [], |
| 39 | + "metadata": {}, |
| 40 | + "tags": [], |
| 41 | + "comments": [] |
| 42 | + } |
36 | 43 | instances = image_df["instanceId"].dropna().unique() |
37 | 44 | for instance in instances: |
38 | 45 | instance_df = image_df[image_df["instanceId"] == instance] |
@@ -67,27 +74,31 @@ def df_to_annotations(df, output_dir): |
67 | 74 | "name": row["attributeName"] |
68 | 75 | } |
69 | 76 | ) |
70 | | - image_annotation.append(instance_annotation) |
| 77 | + image_annotation["instances"].append(instance_annotation) |
71 | 78 | image_width = image_width or instance_df.iloc[0]["imageWidth"] |
72 | 79 | image_height = image_height or instance_df.iloc[0]["imageHeight"] |
73 | 80 | image_pinned = image_pinned or instance_df.iloc[0]["imagePinned"] |
74 | 81 | image_status = image_status or instance_df.iloc[0]["imageStatus"] |
75 | 82 |
|
76 | 83 | comments = image_df[image_df["type"] == "comment"] |
77 | 84 | for _, comment in comments.iterrows(): |
78 | | - comment_json = {"type": "comment"} |
| 85 | + comment_json = {} |
79 | 86 | comment_json.update(comment["meta"]) |
| 87 | + comment_json["correspondence"] = comment_json["comments"] |
| 88 | + del comment_json["comments"] |
80 | 89 | comment_json["resolved"] = comment["commentResolved"] |
81 | | - image_annotation.append(comment_json) |
| 90 | + image_annotation["comments"].append(comment_json) |
| 91 | + |
| 92 | + tags = image_df[image_df["type"] == "tag"] |
| 93 | + for _, tag in tags.iterrows(): |
| 94 | + image_annotation["tags"].append(tag["tag"]) |
82 | 95 |
|
83 | | - meta = { |
84 | | - "type": "meta", |
| 96 | + image_annotation["metadata"] = { |
85 | 97 | "width": int(image_width), |
86 | 98 | "height": int(image_height), |
87 | 99 | "status": image_status, |
88 | 100 | "pinned": bool(image_pinned) |
89 | 101 | } |
90 | | - image_annotation.append(meta) |
91 | 102 | json.dump( |
92 | 103 | image_annotation, |
93 | 104 | open(output_dir / f"{image}___{project_suffix}", "w"), |
@@ -236,13 +247,11 @@ def __append_annotation(annotation_dict): |
236 | 247 |
|
237 | 248 | def __get_image_metadata(image_name, annotations): |
238 | 249 | image_metadata = {"imageName": image_name} |
239 | | - for annotation in annotations: |
240 | | - if "type" in annotation and annotation["type"] == "meta": |
241 | | - image_metadata["imageHeight"] = annotation.get("height") |
242 | | - image_metadata["imageWidth"] = annotation.get("width") |
243 | | - image_metadata["imageStatus"] = annotation.get("status") |
244 | | - image_metadata["imagePinned"] = annotation.get("pinned") |
245 | | - break |
| 250 | + |
| 251 | + image_metadata["imageHeight"] = annotations["metadata"].get("height") |
| 252 | + image_metadata["imageWidth"] = annotations["metadata"].get("width") |
| 253 | + image_metadata["imageStatus"] = annotations["metadata"].get("status") |
| 254 | + image_metadata["imagePinned"] = annotations["metadata"].get("pinned") |
246 | 255 | return image_metadata |
247 | 256 |
|
248 | 257 | def __get_user_metadata(annotation): |
@@ -289,38 +298,30 @@ def __get_user_metadata(annotation): |
289 | 298 | image_name = annotation_path.name.split(type_postfix)[0] |
290 | 299 | image_metadata = __get_image_metadata(image_name, annotation_json) |
291 | 300 | annotation_instance_id = 0 |
292 | | - for annotation in annotation_json: |
| 301 | + if include_comments: |
| 302 | + for annotation in annotation_json["comments"]: |
| 303 | + comment_resolved = annotation["resolved"] |
| 304 | + comment_meta = { |
| 305 | + "x": annotation["x"], |
| 306 | + "y": annotation["y"], |
| 307 | + "comments": annotation["correspondence"] |
| 308 | + } |
| 309 | + annotation_dict = { |
| 310 | + "type": "comment", |
| 311 | + "meta": comment_meta, |
| 312 | + "commentResolved": comment_resolved, |
| 313 | + } |
| 314 | + user_metadata = __get_user_metadata(annotation) |
| 315 | + annotation_dict.update(user_metadata) |
| 316 | + annotation_dict.update(image_metadata) |
| 317 | + __append_annotation(annotation_dict) |
| 318 | + if include_tags: |
| 319 | + for annotation in annotation_json["tags"]: |
| 320 | + annotation_dict = {"type": "tag", "tag": annotation} |
| 321 | + annotation_dict.update(image_metadata) |
| 322 | + __append_annotation(annotation_dict) |
| 323 | + for annotation in annotation_json["instances"]: |
293 | 324 | annotation_type = annotation.get("type", "mask") |
294 | | - if annotation_type == "meta": |
295 | | - continue |
296 | | - if annotation_type == "comment": |
297 | | - if include_comments: |
298 | | - comment_resolved = annotation["resolved"] |
299 | | - comment_meta = { |
300 | | - "x": annotation["x"], |
301 | | - "y": annotation["y"], |
302 | | - "comments": annotation["comments"] |
303 | | - } |
304 | | - annotation_dict = { |
305 | | - "type": annotation_type, |
306 | | - "meta": comment_meta, |
307 | | - "commentResolved": comment_resolved, |
308 | | - } |
309 | | - user_metadata = __get_user_metadata(annotation) |
310 | | - annotation_dict.update(user_metadata) |
311 | | - annotation_dict.update(image_metadata) |
312 | | - __append_annotation(annotation_dict) |
313 | | - continue |
314 | | - if annotation_type == "tag": |
315 | | - if include_tags: |
316 | | - annotation_tag = annotation["name"] |
317 | | - annotation_dict = { |
318 | | - "type": annotation_type, |
319 | | - "tag": annotation_tag |
320 | | - } |
321 | | - annotation_dict.update(image_metadata) |
322 | | - __append_annotation(annotation_dict) |
323 | | - continue |
324 | 325 | annotation_class_name = annotation.get("className") |
325 | 326 | if annotation_class_name is None or annotation_class_name not in class_name_to_color: |
326 | 327 | logger.warning( |
|
0 commit comments