|
|
|
import cv2 |
|
import os |
|
import keras |
|
from keras.applications.imagenet_utils import preprocess_input |
|
from keras.backend.tensorflow_backend import set_session |
|
from keras.models import Model |
|
from keras.preprocessing import image |
|
import matplotlib |
|
matplotlib.use('Agg') |
|
import matplotlib.pyplot as plt |
|
import numpy as np |
|
import pickle |
|
from random import shuffle |
|
from scipy.misc import imread |
|
from scipy.misc import imresize |
|
import tensorflow as tf |
|
import random |
|
from ssd_tools.ssd import SSD300 |
|
from ssd_tools.ssd_training import MultiboxLoss |
|
from ssd_tools.ssd_utils import BBoxUtility |
|
|
|
|
|
|
|
|
|
IMAGE_DIR=os.path.join('training','img') |
|
ANOTATION_FILE=os.path.join('training','page_layout.pkl') |
|
|
|
np.set_printoptions(suppress=True) |
|
random.seed(77) |
|
NUM_CLASSES = 2 |
|
input_shape = (300, 300, 3) |
|
|
|
priors = pickle.load(open(os.path.join('ssd_tools','prior_boxes_ssd300.pkl'), 'rb')) |
|
bbox_util = BBoxUtility(NUM_CLASSES, priors) |
|
|
|
gt = pickle.load(open(ANOTATION_FILE, 'rb')) |
|
keys = sorted(gt.keys()) |
|
random.shuffle(keys) |
|
num_train = int(round(0.9 * len(keys))) |
|
train_keys = keys[:num_train] |
|
val_keys = keys[num_train:] |
|
num_val = len(val_keys) |
|
|
|
class Generator(object): |
|
def __init__(self, gt, bbox_util, |
|
batch_size, path_prefix, |
|
train_keys, val_keys, image_size, |
|
saturation_var=0.5, |
|
brightness_var=0.5, |
|
contrast_var=0.5, |
|
lighting_std=0.5, |
|
hflip_prob=0.5, |
|
vflip_prob=0.5, |
|
do_crop=True, |
|
crop_area_range=[0.8, 1.0], |
|
aspect_ratio_range=[1.0,1.0]): |
|
self.gt = gt |
|
self.bbox_util = bbox_util |
|
self.batch_size = batch_size |
|
self.path_prefix = path_prefix |
|
self.train_keys = train_keys |
|
self.val_keys = val_keys |
|
self.train_batches = len(train_keys) |
|
self.val_batches = len(val_keys) |
|
self.image_size = image_size |
|
self.color_jitter = [] |
|
if saturation_var: |
|
self.saturation_var = saturation_var |
|
self.color_jitter.append(self.saturation) |
|
if brightness_var: |
|
self.brightness_var = brightness_var |
|
self.color_jitter.append(self.brightness) |
|
if contrast_var: |
|
self.contrast_var = contrast_var |
|
self.color_jitter.append(self.contrast) |
|
self.lighting_std = lighting_std |
|
self.hflip_prob = hflip_prob |
|
self.vflip_prob = vflip_prob |
|
self.do_crop = do_crop |
|
self.crop_area_range = crop_area_range |
|
self.aspect_ratio_range = aspect_ratio_range |
|
|
|
def grayscale(self, rgb): |
|
return rgb.dot([0.299, 0.587, 0.114]) |
|
|
|
|
|
def saturation(self, rgb): |
|
gs = self.grayscale(rgb) |
|
alpha = 2 * np.random.random() * self.saturation_var |
|
alpha += 1 - self.saturation_var |
|
rgb = rgb * alpha + (1 - alpha) * gs[:, :, None] |
|
return np.clip(rgb, 0, 255) |
|
|
|
def brightness(self, rgb): |
|
alpha = 2 * np.random.random() * self.brightness_var |
|
alpha += 1 - self.saturation_var |
|
rgb = rgb * alpha |
|
return np.clip(rgb, 0, 255) |
|
|
|
def contrast(self, rgb): |
|
gs = self.grayscale(rgb).mean() * np.ones_like(rgb) |
|
alpha = 2 * np.random.random() * self.contrast_var |
|
alpha += 1 - self.contrast_var |
|
rgb = rgb * alpha + (1 - alpha) * gs |
|
return np.clip(rgb, 0, 255) |
|
|
|
def lighting(self, img): |
|
cov = np.cov(img.reshape(-1, 3) / 255.0, rowvar=False) |
|
eigval, eigvec = np.linalg.eigh(cov) |
|
noise = np.random.randn(3) * self.lighting_std |
|
noise = eigvec.dot(eigval * noise) * 255 |
|
img += noise |
|
return np.clip(img, 0, 255) |
|
|
|
def horizontal_flip(self, img, y): |
|
if np.random.random() < self.hflip_prob: |
|
img = img[:, ::-1] |
|
y[:, [0, 2]] = 1 - y[:, [2, 0]] |
|
return img, y |
|
|
|
def vertical_flip(self, img, y): |
|
if np.random.random() < self.vflip_prob: |
|
img = img[::-1] |
|
y[:, [1, 3]] = 1 - y[:, [3, 1]] |
|
return img, y |
|
|
|
def random_sized_crop(self, img, targets): |
|
img_w = img.shape[1] |
|
img_h = img.shape[0] |
|
img_area = img_w * img_h |
|
random_scale = np.random.random() |
|
random_scale *= (self.crop_area_range[1] - |
|
self.crop_area_range[0]) |
|
random_scale += self.crop_area_range[0] |
|
target_area = random_scale * img_area |
|
random_ratio = np.random.random() |
|
random_ratio *= (self.aspect_ratio_range[1] - |
|
self.aspect_ratio_range[0]) |
|
random_ratio += self.aspect_ratio_range[0] |
|
w = np.round(np.sqrt(target_area * random_ratio)) |
|
h = np.round(np.sqrt(target_area / random_ratio)) |
|
if np.random.random() < 0.5: |
|
w, h = h, w |
|
w = min(w, img_w) |
|
w_rel = w / img_w |
|
w = int(w) |
|
h = min(h, img_h) |
|
h_rel = h / img_h |
|
h = int(h) |
|
x = np.random.random() * (img_w - w) |
|
x_rel = x / img_w |
|
x = int(x) |
|
y = np.random.random() * (img_h - h) |
|
y_rel = y / img_h |
|
y = int(y) |
|
img = img[y:y+h, x:x+w] |
|
new_targets = [] |
|
for box in targets: |
|
cx = 0.5 * (box[0] + box[2]) |
|
cy = 0.5 * (box[1] + box[3]) |
|
if (x_rel < cx < x_rel + w_rel and |
|
y_rel < cy < y_rel + h_rel): |
|
xmin = (box[0] - x_rel) / w_rel |
|
ymin = (box[1] - y_rel) / h_rel |
|
xmax = (box[2] - x_rel) / w_rel |
|
ymax = (box[3] - y_rel) / h_rel |
|
xmin = max(0, xmin) |
|
ymin = max(0, ymin) |
|
xmax = min(1, xmax) |
|
ymax = min(1, ymax) |
|
box[:4] = [xmin, ymin, xmax, ymax] |
|
new_targets.append(box) |
|
new_targets = np.asarray(new_targets).reshape(-1, targets.shape[1]) |
|
return img, new_targets |
|
|
|
def generate(self, train=True): |
|
while True: |
|
if train: |
|
shuffle(self.train_keys) |
|
keys = self.train_keys |
|
else: |
|
shuffle(self.val_keys) |
|
keys = self.val_keys |
|
inputs = [] |
|
targets = [] |
|
for key in keys: |
|
img_path = self.path_prefix + key |
|
img = imread(img_path,mode="RGB").astype('float32') |
|
y = self.gt[key].copy() |
|
if train and self.do_crop: |
|
img, y = self.random_sized_crop(img, y) |
|
img = imresize(img, self.image_size).astype('float32') |
|
|
|
|
|
if train: |
|
shuffle(self.color_jitter) |
|
for jitter in self.color_jitter: |
|
img = jitter(img) |
|
if self.lighting_std: |
|
img = self.lighting(img) |
|
if self.hflip_prob > 0: |
|
img, y = self.horizontal_flip(img, y) |
|
if self.vflip_prob > 0: |
|
img, y = self.vertical_flip(img, y) |
|
|
|
|
|
y = self.bbox_util.assign_boxes(y) |
|
|
|
inputs.append(img) |
|
targets.append(y) |
|
if len(targets) == self.batch_size: |
|
tmp_inp = np.array(inputs) |
|
|
|
tmp_targets = np.array(targets) |
|
inputs = [] |
|
targets = [] |
|
yield preprocess_input(tmp_inp), tmp_targets |
|
|
|
path_prefix = IMAGE_DIR+"/" |
|
gen = Generator(gt, bbox_util, 5, path_prefix, |
|
train_keys, val_keys, |
|
(input_shape[0], input_shape[1]), do_crop=True) |
|
|
|
model = SSD300(input_shape, num_classes=NUM_CLASSES) |
|
|
|
def schedule(epoch, decay=0.9): |
|
return base_lr * decay**(epoch) |
|
|
|
callbacks = [keras.callbacks.ModelCheckpoint('./checkpoints/weights.{epoch:02d}-{val_loss:.2f}.hdf5', |
|
verbose=1, |
|
save_weights_only=True), |
|
keras.callbacks.LearningRateScheduler(schedule)] |
|
|
|
base_lr = 3e-4 |
|
optim = keras.optimizers.Adam(lr=base_lr) |
|
model.compile(optimizer=optim, |
|
loss=MultiboxLoss(NUM_CLASSES, neg_pos_ratio=5.0).compute_loss) |
|
|
|
nb_epoch = 100 |
|
history = model.fit_generator(gen.generate(True), gen.train_batches, |
|
nb_epoch, verbose=1, |
|
callbacks=callbacks, |
|
validation_data=gen.generate(False), |
|
nb_val_samples=gen.val_batches, |
|
nb_worker=1) |
|
|
|
|
