# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/adapters.ipynb.

# %% auto 0
__all__ = ['np_2_b64', 'b64_2_np', 'b64_2_file', 'csvs_2_cvat', 'cvat_2_csvs']

# %% ../nbs/adapters.ipynb 2
import base64, cv2, numpy as np
from typing import Union
from .loader import *
from .paths import P, stems, stem, parent, makedir
from .markup import write_xml, read_xml, AttrDict

# %% ../nbs/adapters.ipynb 3
def np_2_b64(image: np.ndarray) -> str:
    """Convert a numpy image to base64 string"""
    img_str = cv2.imencode(".jpg", image)[1].tobytes()
    img_b64 = base64.b64encode(img_str).decode("utf-8")
    return img_b64


def b64_2_np(input: str) -> np.ndarray:
    input = bytes(input, "utf-8")
    input = base64.b64decode(input)
    img_nparr = np.frombuffer(input, np.uint8)
    img_cv2 = cv2.imdecode(img_nparr, cv2.IMREAD_COLOR)
    img_rgb = cv2.cvtColor(img_cv2, cv2.COLOR_BGR2RGB)
    return img_rgb


def b64_2_file(
    input: str, # base64 encoded string
    fpath: Union[str, P] # Place where you want to save the file
    ) -> None:
    """Save a file encoded as a base64 `input` at given `fpath`"""
    with open(fpath, "wb") as pdf_file:
        input = bytes(input, "utf-8")
        input = base64.b64decode(input)
        pdf_file.write(input)
        logger.info(f"wrote pdf file to {fpath}")

# %% ../nbs/adapters.ipynb 4
def _process(df: pd.DataFrame, label_column="readable_label", default_label="Others"):
    df["@xbr"] = df["X"]
    df["@xtl"] = df["x"]
    df["@ybr"] = df["Y"]
    df["@ytl"] = df["y"]
    df["@label"] = df[label_column] if label_column in df.columns else default_label
    df["@occluded"] = "0"
    df["@source"] = "manual"
    df["@z_order"] = "0"
    df.drop(
        [
            c
            for c in df.columns
            if c
            not in set(
                [
                    "@xbr",
                    "@xtl",
                    "@ybr",
                    "@ytl",
                    "@label",
                    "@occluded",
                    "@source",
                    "@z_order",
                ]
            )
        ],
        axis=1,
        inplace=True,
    )
    records = df.to_dict(orient="records")
    return records


def csvs_2_cvat(images_folder, csvs_folder, xml_output_file, items=None):
    images_folder, csvs_folder = [P(_) for _ in [images_folder, csvs_folder]]
    data = AttrDict({"annotations": {"image": []}})
    if items is None:
        items = common(stems(images_folder), stems(csvs_folder))

    items = sorted(items)

    for ix, item in enumerate(track(items)):
        _ia = _image_annotation = AttrDict({})
        image = images_folder / f"{item}.jpg"
        df = pd.read_csv(f"{csvs_folder}/{item}.csv")
        _ia["@height"], _ia["@width"] = read(image).shape[:2]
        _ia["@id"] = str(ix)
        _ia["@name"] = f"{item}.jpg"
        _ia["box"] = _process(df)
        data.annotations.image.append(_ia)
    write_xml(data, xml_output_file)

def _cvat_ann_2_csv(ann):
    df = pd.DataFrame([a.to_dict() for a in ann.box])
    df.rename({'@xtl': 'x', '@ytl': 'y', '@xbr': 'X', '@ybr': 'Y', '@label': 'readable_label'}, inplace=1, axis=1)
    for col in 'xyXY':
        df[col] = df[col].map(lambda x: int(float(x)))
    df.drop(['@z_order','@source','@occluded'], axis=1, inplace=True)
    height, width = int(ann['@height']), int(ann['@width'])
    df = to_relative(df, height, width)
    return df

def cvat_2_csvs(xmlfile, csvs_folder):
    data = read_xml(xmlfile)
    for item in data.annotations.image:
        try:
            df = _cvat_ann_2_csv(item)
            save_at = f'{csvs_folder}/{stem(item["@name"])}.csv'
            makedir(parent(save_at))
            df.to_csv(save_at)
        except Exception as e:
            Warn(f'{e} @ {item["@name"]}')

