Source code for tensorbay.opendataset.Elpv.loader

#!/usr/bin/env python3
#
# Copyright 2021 Graviti. Licensed under MIT License.
#
# pylint: disable=invalid-name

"""Dataloader of Elpv dataset."""

import os

from tensorbay.dataset import Data, Dataset
from tensorbay.label import Classification

DATASET_NAME = "Elpv"


[docs]def Elpv(path: str) -> Dataset: """`elpv <https://github.com/zae-bayern/elpv-dataset>`_ dataset. The file structure should be like:: <path> labels.csv images/ cell0001.png ... Arguments: path: The root directory of the dataset. Returns: Loaded :class:`~tensorbay.dataset.dataset.Dataset` instance. """ root_path = os.path.abspath(os.path.expanduser(path)) dataset = Dataset(DATASET_NAME) dataset.load_catalog(os.path.join(os.path.dirname(__file__), "catalog.json")) segment = dataset.create_segment() csv_path = os.path.join(root_path, "labels.csv") with open(csv_path, encoding="utf-8") as csv_file: for row in csv_file: image_name, attributes, category = row.strip().split() dirname, basename = image_name.split("/") image_path = os.path.join(root_path, dirname, basename) data = Data(image_path) data.label.classification = Classification( attributes={"defect probability": float(attributes)}, category=category ) segment.append(data) return dataset