wko_anime-face-similarity/load_test_data.py

40 lines
1.0 KiB
Python
Raw Normal View History

2023-01-29 22:18:19 +01:00
import numpy as np
import os
from skimage.io import imread
import cv2 as cv
from pathlib import Path
2023-01-29 23:18:50 +01:00
2023-01-29 22:18:19 +01:00
def load_data(input_dir, newSize=(64,64)):
image_path = Path(input_dir)
file_names = os.listdir(image_path)
categories_name = []
2023-01-29 23:18:50 +01:00
categories_count = []
2023-01-29 22:18:19 +01:00
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 = []
2023-01-29 23:18:50 +01:00
labels = []
2023-01-29 22:18:19 +01:00
for n in file_names:
p = image_path / n
img = imread(p) # zwraca ndarry postaci xSize x ySize x colorDepth
img = cv.resize(img, newSize, interpolation=cv.INTER_AREA) # zwraca ndarray
test_img.append(img)
labels.append(n)
2023-01-29 23:18:50 +01:00
X = {}
2023-01-29 22:18:19 +01:00
X["values"] = np.array(test_img)
2023-01-29 23:18:50 +01:00
X["name"] = categories_name
X["names_count"] = categories_count
X["labels"] = labels
return X