-
Notifications
You must be signed in to change notification settings - Fork 1.2k
/
image.py
58 lines (50 loc) 路 1.58 KB
/
image.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
import base64
import os
from typing import TYPE_CHECKING, Dict, List, Tuple
from dvc.render import FILENAME_FIELD, REVISION_FIELD, SRC_FIELD
from . import Converter
if TYPE_CHECKING:
from dvc.types import StrPath
class ImageConverter(Converter):
@staticmethod
def _write_image(
path: "StrPath",
revision: str,
filename: str,
image_data: bytes,
) -> "StrPath":
img_path = os.path.join(
path,
f"{revision}_{filename}".replace(os.sep, "_").replace("/", "_"),
)
with open(img_path, "wb") as fd:
fd.write(image_data)
return img_path
@staticmethod
def _encode_image(
image_data: bytes,
) -> str:
base64_str = base64.b64encode(image_data).decode()
return f"data:image;base64,{base64_str}"
def convert(
self, data, revision: str, filename: str, **kwargs
) -> Tuple[List[Dict], Dict]:
"""
Convert the DVC Plots content to DVC Render datapoints.
Return both generated datapoints and updated properties.
"""
path = self.plot_properties.get("out")
if path:
if not os.path.isdir(path):
os.makedirs(path, exist_ok=True)
src = self._write_image(
os.path.abspath(path), revision, filename, data
)
else:
src = self._encode_image(data)
datapoint = {
REVISION_FIELD: revision,
FILENAME_FIELD: filename,
SRC_FIELD: src,
}
return [datapoint], self.plot_properties