Skip to content

Commit

Permalink
plots: add bboxes to dvc plots diff so VSCode can used it to displa…
Browse files Browse the repository at this point in the history
…y bbox on images

Fixes #10198
  • Loading branch information
AlexandreKempf authored and skshetry committed Mar 5, 2024
1 parent b0384f8 commit b1c547a
Show file tree
Hide file tree
Showing 4 changed files with 46 additions and 6 deletions.
1 change: 1 addition & 0 deletions dvc/render/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,3 +6,4 @@
ANCHOR_DEFINITIONS = "anchor_definitions"
TYPE_KEY = "type"
SRC = "src"
ANNOTATIONS = "annotations"
3 changes: 2 additions & 1 deletion dvc/render/convert.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from typing import Union

from dvc.render import REVISION, REVISIONS, SRC, TYPE_KEY
from dvc.render import REVISION, REVISIONS, SRC, TYPE_KEY, ANNOTATIONS
from dvc.render.converter.image import ImageConverter
from dvc.render.converter.vega import VegaConverter

Expand Down Expand Up @@ -42,6 +42,7 @@ def to_json(renderer, split: bool = False) -> list[dict]:
{
TYPE_KEY: renderer.TYPE,
REVISIONS: [datapoint.get(REVISION)],
**datapoint.get(ANNOTATIONS),
"url": datapoint.get(SRC),
}
for datapoint in renderer.datapoints
Expand Down
21 changes: 16 additions & 5 deletions dvc/render/converter/image.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
import base64
import os
import json
from typing import TYPE_CHECKING, Any

from dvc.render import FILENAME, REVISION, SRC
from dvc.render import FILENAME, REVISION, SRC, ANNOTATIONS

from . import Converter

Expand Down Expand Up @@ -46,15 +47,25 @@ def flat_datapoints(self, revision: str) -> tuple[list[dict], dict]:
path = self.properties.get("out")
datapoints = []
datas, properties = self.convert()
for filename, _, image_data in datas:

if "annotations" in properties:
with open(properties["annotations"], encoding="utf-8") as annotations_path:
annotations = json.load(annotations_path)

for filename, _, image_content in datas:
if path:
if not os.path.isdir(path):
os.makedirs(path, exist_ok=True)
src = self._write_image(
os.path.abspath(path), revision, filename, image_data
os.path.abspath(path), revision, filename, image_content
)
else:
src = self._encode_image(image_data)
datapoint = {REVISION: revision, FILENAME: filename, SRC: src}
src = self._encode_image(image_content)
datapoint = {
REVISION: revision,
FILENAME: filename,
SRC: src,
ANNOTATIONS: annotations,
}
datapoints.append(datapoint)
return datapoints, properties
27 changes: 27 additions & 0 deletions dvc/repo/plots/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from copy import deepcopy
from functools import partial
from multiprocessing import cpu_count
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Optional, Union

import dpath
Expand Down Expand Up @@ -357,6 +358,10 @@ def infer_data_sources(plot_id, config=None):
if isinstance(x, dict):
sources.append(first(x.keys()))

# annotation = config.get("annotations", None)
# if isinstance(annotation, str):
# sources.append(annotation)

return ldistinct(source for source in sources)


Expand Down Expand Up @@ -467,6 +472,7 @@ def _resolve_definitions(
for k, v in unpacked["data"].items()
if _closest_parent(fs, k, plot_ids_parents) == data_path
}
unpacked = _add_annotations_to_image_definition(unpacked)
dpath.merge(result, unpacked)
elif _matches(targets, config_path, plot_id):
adjusted_props = _adjust_sources(fs, plot_props, config_dir)
Expand All @@ -483,6 +489,27 @@ def _closest_parent(fs, path, parents):
best_result = common_path
return best_result

def _add_annotations_to_image_definition(target):
if "data" not in target:
return target

path_to_remove = []
for path in target["data"]:
annotation_file = Path(path).with_suffix(".json").as_posix()
if ImageRenderer.matches(path) and annotation_file in target["data"]:
annotations = {"annotations": annotation_file}
# empty dict all share the same reference, so override them
if target["data"][path]:
target["data"][path].update({"annotations": annotation_file})
else:
target["data"][path] = annotations
path_to_remove.append(annotation_file)
# remove the annotation files from the data once they are added to the config
target["data"] = {
k: v for k, v in target["data"].items() if k not in path_to_remove
}
return target


def _collect_pipeline_files(repo, targets: list[str], props, onerror=None):
result: dict[str, dict] = {}
Expand Down

0 comments on commit b1c547a

Please sign in to comment.