fernandoperlar/preprocessing_image

1import os2import pandas as pd3import numpy as np4import cv2 as cv5import matplotlib.pyplot as plt6from sklearn import model_selection7from keras import preprocessing8from .Misc import *9 10class Data:11	def __init__(self, path):12		self.images = self.__extract_images(path)13		self.images.category, self.labels = self.images.category.factorize()14		self.images.category = self.images.category.astype(str)15		self.training, self.test = None, None16 17	def train_test_split(self, test_size=0.15, shuffle=True, stratify=False):18		return model_selection.train_test_split(19			self.images,20			test_size=test_size,21			random_state=42,22			shuffle=shuffle,23			stratify=self.images.category if stratify else None24		)25 26	def count_labels(self, data, name):27		amount = data.category.value_counts().values28		29		print(f"{name}: {amount} {np.round(amount/len(data), 2)}")30 31	def image_generator(self, shuffle=True):32		train_datagen = preprocessing.image.ImageDataGenerator(rescale=1./255, validation_split=0.2)33		test_datagen = preprocessing.image.ImageDataGenerator(rescale=1./255)34 35		generator_properties = {36			"x_col": "image",37			"y_col": "category",38			"target_size": (215, 538),39			"color_mode": "rgb",40			"class_mode": "categorical"41		}42 43		train_generator = train_datagen.flow_from_dataframe(44			**generator_properties,45 46			dataframe=self.training,47			batch_size=10,48			shuffle=shuffle,49			subset="training"50		)51 52		validation_generator = train_datagen.flow_from_dataframe(53			**generator_properties,54 55			dataframe=self.training,56			batch_size=10,57			shuffle=shuffle,58			subset="validation"59		)60 61		test_generator = test_datagen.flow_from_dataframe(62			**generator_properties,63 64			dataframe=self.test,65			batch_size=1,66			shuffle=False67		)68 69		return train_generator, validation_generator, test_generator70 71	def detectColor(self, image, lower, upper):72		if tf.is_tensor(image):73			temp_image = image.numpy().copy()74		else:75			temp_image = image.copy()76 77		hsv_image = temp_image.copy()78		hsv_image = cv.cvtColor(hsv_image, cv.COLOR_RGB2HSV)79		mask = cv.inRange(hsv_image, lower, upper)80 81		result = temp_image.copy()82		result[np.where(mask == 0)] = 083		84		return result85 86	def getImageTensor(self, images, lower, upper):87		results = []88 89		for img in images:90			results.append(np.expand_dims(self.detectColor(img, lower, upper), axis=0))91 92		return np.concatenate(results, axis=0)93 94	def show_images(self, generator, filters, name):95		generator.reset()96 97		img, label = generator.next()98 99		fig, axs = plt.subplots(nrows=3, ncols=1, constrained_layout=True)100		fig.suptitle(name)101 102		for ax in axs:103			ax.remove()104 105		gridspec = axs[0].get_subplotspec().get_gridspec()106		subfigs = [fig.add_subfigure(gs) for gs in gridspec]107 108		for row, subfig in enumerate(subfigs):109			subfig.suptitle(str(self.labels[np.argmax(label[row], axis=-1)]).title())110 111			axs = subfig.subplots(nrows=1, ncols=4)112 113			for col, ax in enumerate(axs):114				ax.imshow(list(filters.values())[col](img)[row])115				ax.set_title(list(filters)[col].title())116				ax.axis("off")117				118				ax.plot()119 120	def __extract_images(self, path):121		images = []122		123		for category in os.listdir(path):124			for filename in os.listdir(path + category):125				images.append([path + category + "/" + filename, category])126 127		return pd.DataFrame(images, columns=["image", "category"])128