torchvisionのDatasetFolderでSVGを読み込む

ImageFolderのようにベクター画像のSVGを読み込みたい。

💡アプローチ

1. ImageFolderの親クラスDatasetFolderはloaderメソッドでファイルの読み込みロジックを定義しているので、SVG用のloaderメソッドsvg_loaderを作る。
2. svg_loaderでは、読み込んだSVGファイルをcairosvgのsvg2pngでラスタ画像化。ラスタ画像化したデータをBytesIOとPillowを組み合わせてImageとしてロード。
3. svg_loaderをDatasetFolderのインスタンス化の際にloaderに指定する。

📄コード

ecample.py

from PIL import Image
from io import BytesIO
from cairosvg import svg2png
from torchvision.datasets.folder import DatasetFolder

def svg_loader(path: Union[str, Path]) -> Image.Image:
  # svgからpngへ変換
  # (background_color="white"とか与えると背景色を指定できる。)
  png = svg2png(bytestring=svg_data)
  # Imageにして返す。
  return Image.open(BytesIO(png)).convert('RGB')
  
dataset = DatasetFolder(root=dataset_svg, loader=svg_loader, extensions=("svg"))

👍