wko_anime-face-similarity/load_test_data.py

43 lines
1.0 KiB
Python

import numpy as np
import os
from skimage.io import imread
import cv2 as cv
from pathlib import Path
def load_source(filename: str) -> np.ndarray:
return cv.imread(filename)[..., ::-1]
def load_data(input_dir):
image_path = Path(input_dir)
file_names = os.listdir(image_path)
categories_name = []
categories_count = []
count = 0
n = file_names[0]
for name in file_names:
if name != n:
categories_count.append(count)
n = name
count = 1
else:
count += 1
if not name in categories_name:
categories_name.append(name)
categories_count.append(count)
test_img = []
labels = []
for n in file_names:
p = image_path / n
img = load_source(str(p)) # zwraca ndarry postaci xSize x ySize x colorDepth
test_img.append(img)
labels.append(n)
X = {}
X["values"] = np.array(test_img)
X["name"] = categories_name
X["names_count"] = categories_count
X["labels"] = labels
return X