|
|
import math |
|
|
from random import shuffle |
|
|
|
|
|
import cv2 |
|
|
import tensorflow as tf |
|
|
import numpy as np |
|
|
from keras.utils import to_categorical |
|
|
from PIL import Image |
|
|
|
|
|
from .utils import cvtColor, preprocess_input |
|
|
|
|
|
|
|
|
class ClsDatasets(keras.utils.Sequence): |
|
|
def __init__(self, annotation_lines, input_shape, batch_size, num_classes, train, **kwargs): |
|
|
super().__init__() |
|
|
self.annotation_lines = annotation_lines |
|
|
self.length = len(self.annotation_lines) |
|
|
|
|
|
self.input_shape = input_shape |
|
|
self.batch_size = batch_size |
|
|
self.num_classes = num_classes |
|
|
self.train = train |
|
|
|
|
|
def __len__(self): |
|
|
return int(math.ceil(self.length / float(self.batch_size))) |
|
|
|
|
|
def __getitem__(self, index): |
|
|
X_train = [] |
|
|
Y_train = [] |
|
|
start = index * self.batch_size |
|
|
end = min((index + 1) * self.batch_size, self.length) |
|
|
for i in range(start, end): |
|
|
annotation_path = self.annotation_lines[i].split(';')[1].split()[0] |
|
|
image = Image.open(annotation_path) |
|
|
image = self.get_random_data(image, self.input_shape, random=self.train) |
|
|
image = preprocess_input(np.array(image).astype(np.float32)) |
|
|
|
|
|
X_train.append(image) |
|
|
Y_train.append(int(self.annotation_lines[i].split(';')[0])) |
|
|
|
|
|
X_train = np.array(X_train) |
|
|
Y_train = to_categorical(np.array(Y_train), num_classes = self.num_classes) |
|
|
return X_train, Y_train |
|
|
|
|
|
def on_epoch_end(self): |
|
|
if self.train: |
|
|
np.random.shuffle(self.annotation_lines) |
|
|
|
|
|
def rand(self, a=0, b=1): |
|
|
return np.random.rand()*(b-a) + a |
|
|
|
|
|
def get_random_data(self, image, input_shape, jitter=.3, hue=.1, sat=1.5, val=1.5, random=True): |
|
|
|
|
|
image = cvtColor(image) |
|
|
iw, ih = image.size |
|
|
h, w = input_shape |
|
|
|
|
|
if not random: |
|
|
scale = min(w/iw, h/ih) |
|
|
nw = int(iw*scale) |
|
|
nh = int(ih*scale) |
|
|
dx = (w-nw)//2 |
|
|
dy = (h-nh)//2 |
|
|
|
|
|
image = image.resize((nw,nh), Image.BICUBIC) |
|
|
new_image = Image.new('RGB', (w,h), (128,128,128)) |
|
|
new_image.paste(image, (dx, dy)) |
|
|
image_data = np.array(new_image, np.float32) |
|
|
|
|
|
return image_data |
|
|
|
|
|
new_ar = w/h * self.rand(1-jitter,1+jitter)/self.rand(1-jitter,1+jitter) |
|
|
scale = self.rand(.75, 1.25) |
|
|
if new_ar < 1: |
|
|
nh = int(scale*h) |
|
|
nw = int(nh*new_ar) |
|
|
else: |
|
|
nw = int(scale*w) |
|
|
nh = int(nw/new_ar) |
|
|
image = image.resize((nw,nh), Image.BICUBIC) |
|
|
|
|
|
dx = int(self.rand(0, w-nw)) |
|
|
dy = int(self.rand(0, h-nh)) |
|
|
new_image = Image.new('RGB', (w,h), (128,128,128)) |
|
|
new_image.paste(image, (dx, dy)) |
|
|
image = new_image |
|
|
|
|
|
flip = self.rand()<.5 |
|
|
if flip: image = image.transpose(Image.FLIP_LEFT_RIGHT) |
|
|
|
|
|
rotate = self.rand()<.5 |
|
|
if rotate: |
|
|
angle = np.random.randint(-15,15) |
|
|
a,b = w/2,h/2 |
|
|
M = cv2.getRotationMatrix2D((a,b),angle,1) |
|
|
image = cv2.warpAffine(np.array(image), M, (w,h), borderValue=[128,128,128]) |
|
|
|
|
|
hue = self.rand(-hue, hue) |
|
|
sat = self.rand(1, sat) if self.rand()<.5 else 1/self.rand(1, sat) |
|
|
val = self.rand(1, val) if self.rand()<.5 else 1/self.rand(1, val) |
|
|
x = cv2.cvtColor(np.array(image,np.float32)/255, cv2.COLOR_RGB2HSV) |
|
|
x[..., 1] *= sat |
|
|
x[..., 2] *= val |
|
|
x[x[:,:, 0]>360, 0] = 360 |
|
|
x[:, :, 1:][x[:, :, 1:]>1] = 1 |
|
|
x[x<0] = 0 |
|
|
image_data = cv2.cvtColor(x, cv2.COLOR_HSV2RGB)*255 |
|
|
return image_data |
|
|
|
