# AUTOGENERATED! DO NOT EDIT! File to edit: 01_loader.ipynb (unless otherwise specified).

__all__ = ['slugify', 'get_image_files', 'device', 'archive_loader', 'db_loader', 'treemap_loader']

# Cell
from pathlib import Path

# Cell
def slugify(filepath):
    return f'{filepath.stem}_{str(filepath.stat().st_mtime).split(".")[0]}'

def get_image_files(path):
    img_extensions = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp']
    return [(f, slugify(f)) for f in path.rglob('*') if f.suffix in img_extensions]

# Cell
import torch
import torchvision

# Cell
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device

# Cell
def archive_loader(filepaths, root, device):
    dbpath = root/'memery.pt'
#     dbpath_backup = root/'memery.pt'
    db = db_loader(dbpath, device)

    current_slugs = [slug for path, slug in filepaths]
    archive_db = {i:db[item[0]] for i, item in enumerate(db.items()) if item[1]['slug'] in current_slugs}
    archive_slugs = [v['slug'] for v in archive_db.values()]
    new_files = [(str(path), slug) for path, slug in filepaths if slug not in archive_slugs]

    return(archive_db, new_files)

# Cell
def db_loader(dbpath, device):

    # check for savefile or backup and extract
    if Path(dbpath).exists():
        db = torch.load(dbpath, device)
#     elif dbpath_backup.exists():
#         db = torch.load(dbpath_backup)
    else:
        db = {}
    return(db)

# Cell
from annoy import AnnoyIndex

# Cell
def treemap_loader(treepath):
    treemap = AnnoyIndex(512, 'angular')

    if treepath.exists():
        treemap.load(str(treepath))
    else:
        treemap = None
    return(treemap)