1.py

import os

import numpy as np
import pandas as pd
from beat.backend.python.database import View
from PIL import Image


class Test(View):
    def index(self, root_folder, parameters):
        """Creates the data for the database indexation"""

        csv_path = os.path.join(root_folder, "test.csv")
        df = pd.read_csv(csv_path)

        df["filename"] = df["filename"].apply(lambda x: os.path.join(root_folder, x))
        df = df.rename(columns={"filename": "image"})

        # ------------- v1 random labels -----
        # remove this part for v2 release
        num_files = len(df)
        np.random.seed(0)
        df["label"] = np.random.randint(0, 2, size=num_files, dtype=bool)
        df["category"] = np.random.randint(1, 7, size=num_files, dtype=int)
        # ------------------------------------

        return list(df.itertuples(index=False))

    def get(self, output, index):
        """Returns the data for the output based on the index content"""

        obj = self.objs[index]

        if output == "image":
            img = np.asarray(Image.open(obj.image))
            return {"value": img}
        elif output == "label":
            return {"value": bool(obj.label)}
        elif output == "category":
            return {"value": np.cast["int32"](obj.category)}