resnet-classifier/image_classifier/data.py

# todo: clean this mess up

import copy
# import h5py
import json
import logging
import numpy as np
import os
import random
import torch
import math

from .config import cfg

from collections import defaultdict
from functools import cache, cached_property
from itertools import groupby, zip_longest
from pathlib import Path
from typing import Any

from torch import Tensor
from torch.utils.data import DataLoader, Dataset as _Dataset
import torchvision.transforms as transforms
from tqdm.auto import tqdm

from PIL import Image

# torch.multiprocessing.set_sharing_strategy("file_system")

_logger = logging.getLogger(__name__)

@cache
def get_symmap():
	return { " ": 0, "<s>": 1, "</s>": 2, "0": 3, "2": 4, "4": 5, "8": 6, "A": 7, "D": 8, "G": 9, "H": 10, "J": 11, "K": 12, "M": 13, "N": 14, "P": 15, "R": 16, "S": 17, "T": 18, "V": 19, "W": 20, "X": 21, "Y": 22 }

@cache
def _get_symbols( content ):
	content = content.replace("O", "0")
	return [f"<s>"] + [ p for p in content ] + [f"</s>"]

class Dataset(_Dataset):
	def __init__(
		self,
		paths,
		width=300,
		height=80,

		symmap=get_symmap(),
		training=False,
	):
		super().__init__()

		self._head = None

		self.paths = paths
		self.width = width
		self.height = height
		
		self.symmap = symmap
		self.training = training

		self.transform = transforms.Compose([
			#transforms.Resize((self.height, self.width)), # for some reason, running the validation dataset breaks when this is set. all images *should* be normalized anyhow
			transforms.ToTensor(),
			transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
		])

	@cached_property
	def symbols(self):
		return sorted(set().union(*[_get_symbols(path.stem) for path in self.paths]))


	def __getitem__(self, index):
		path = self.paths[index]

		# stupid try/except when the original VALL-E training framework was able to insert foreign symbols into the symmap, but that functionality isn't really necessary here
		try:
			text = torch.tensor([*map(self.symmap.get, _get_symbols(path.stem))]).to(torch.uint8)
		except Exception as e:
			print("Invalid symbol:", _get_symbols(path.stem), [*map(self.symmap.get, _get_symbols(path.stem))], path.stem)
			raise e

		image = self.transform(Image.open(path).convert('RGB')).to(cfg.trainer.dtype) # resnet has to be RGB

		return dict(
			index=index,
			path=path,
			image=image,
			text=text,
		)

	def head_(self, n):
		self._head = n

	def training_(self, value):
		self.training = value

	def __len__(self):
		return min(len(self.paths), self._head or len(self.paths))

	def pin_memory(self):
		self.text = self.text.pin_memory()
		self.image = self.image.pin_memory()
		return self


def collate_fn(samples: list[dict]):
	batch: dict[str, Any] = {k: [s[k] for s in samples] for k in samples[0]}
	return batch


def _seed_worker(worker_id):
	worker_seed = torch.initial_seed() % 2**32
	np.random.seed(worker_seed)
	random.seed(worker_seed)


def _create_dataloader(dataset, training):
	return DataLoader(
		dataset=dataset,
		batch_size=cfg.hyperparameters.batch_size if training else cfg.evaluation.batch_size,
		shuffle=True, # training
		drop_last=training,
		num_workers=cfg.dataset.workers,
		collate_fn=collate_fn,
		persistent_workers=cfg.dataset.workers > 0,
		pin_memory=False, # True,
		worker_init_fn=_seed_worker,
	)

def _load_train_val_paths( val_ratio=0.1 ):
	paths = []
	train_paths = []
	val_paths = []

	print(cfg.dataset.training)
	for data_dir in cfg.dataset.training:
		paths.extend(data_dir.rglob("*.png"))

	if len(paths) > 0:
		random.seed(0)
		random.shuffle(paths)
		train_paths.extend(paths)

	if len(cfg.dataset.validation) == 0:
		val_len = math.floor(len(train_paths) * val_ratio)
		train_len = math.floor(len(train_paths) * (1 - val_ratio))

		print(val_len, train_len)

		val_paths = train_paths[:-val_len]
		train_paths = train_paths[:train_len]
	else:
		for data_dir in cfg.dataset.validation:
			paths.extend(data_dir.rglob("*.png"))

		if len(paths) > 0:
			random.seed(0)
			random.shuffle(paths)
			val_paths.extend(paths)

	train_paths, val_paths = map(sorted, [train_paths, val_paths])

	if len(train_paths) == 0:
		raise RuntimeError(f"Failed to find any .png file in {cfg.dataset.training}.")
	# to get it to shut up
	if len(val_paths) == 0:
		val_paths = [ train_paths[0] ]

	return train_paths, val_paths

@cfg.diskcache()
def create_datasets():
	train_paths, val_paths = _load_train_val_paths()

	train_dataset = Dataset(
		train_paths,
		training=True,
	)

	val_dataset = Dataset(
		val_paths,
		train_dataset.symmap,
	)

	val_dataset.head_(cfg.evaluation.size)

	return train_dataset, val_dataset


def create_train_val_dataloader():
	train_dataset, val_dataset = create_datasets()

	subtrain_dataset = copy.deepcopy(train_dataset)
	subtrain_dataset.head_(cfg.evaluation.size)
	subtrain_dataset.training_(False)

	train_dl = _create_dataloader(train_dataset, training=True)
	val_dl = _create_dataloader(val_dataset, training=False)
	subtrain_dl = _create_dataloader(subtrain_dataset, training=False)

	_logger.info(str(train_dataset.symmap))
	

	_logger.info(f"#samples (train): {len(train_dataset)}.")
	_logger.info(f"#samples (val): {len(val_dataset)}.")
	_logger.info(f"#samples (subtrain): {len(subtrain_dataset)}.")

	assert isinstance(subtrain_dl.dataset, Dataset)

	return train_dl, subtrain_dl, val_dl

"""
if __name__ == "__main__":
	create_dataset_hdf5()

	train_dl, subtrain_dl, val_dl = create_train_val_dataloader()
	sample = train_dl.dataset[0]
	print(sample)
"""