File size: 3,767 Bytes
9825f94 04e599a 9825f94 302d684 9825f94 302d684 9825f94 e916c8e 9825f94 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
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):
# Read and convert images to RGB
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
|