mirror of
https://github.com/jantic/DeOldify.git
synced 2026-09-01 15:36:32 +08:00
Jason Antic
127 lines
4.4 KiB
Python
127 lines
4.4 KiB
Python
import math, os, json, sys, re, numpy as np, pickle, PIL, scipy
|
|
from PIL import Image
|
|
from glob import glob
|
|
from matplotlib import pyplot as plt
|
|
from operator import itemgetter, attrgetter, methodcaller
|
|
from collections import OrderedDict
|
|
import itertools
|
|
from itertools import chain
|
|
|
|
import pandas as pd
|
|
from numpy.random import random, permutation, randn, normal, uniform, choice
|
|
from numpy import newaxis
|
|
from scipy import misc, ndimage
|
|
from scipy.ndimage.interpolation import zoom
|
|
from scipy.ndimage import imread
|
|
from sklearn.metrics import confusion_matrix
|
|
from sklearn.preprocessing import OneHotEncoder
|
|
from sklearn.manifold import TSNE
|
|
import bcolz
|
|
|
|
from IPython.lib.display import FileLink
|
|
|
|
import keras
|
|
from keras import backend as K
|
|
from keras.utils.data_utils import get_file
|
|
from keras.utils import np_utils
|
|
from keras.utils.np_utils import to_categorical
|
|
from keras.models import Sequential, Model
|
|
from keras.layers import Input, Embedding, Reshape, merge, LSTM, Bidirectional
|
|
from keras.layers import TimeDistributed, Activation, SimpleRNN, GRU
|
|
from keras.layers import Flatten, Dense, Dropout, Lambda
|
|
from keras.regularizers import l2, l1
|
|
from keras.layers.normalization import BatchNormalization
|
|
from keras.optimizers import SGD, RMSprop, Adam
|
|
from keras.layers import deserialize as layer_from_config
|
|
from keras.metrics import categorical_crossentropy, categorical_accuracy
|
|
from keras.layers.convolutional import *
|
|
from keras.preprocessing import image, sequence
|
|
from keras.preprocessing.text import Tokenizer
|
|
from vgg16 import Vgg16
|
|
np.set_printoptions(precision=4, linewidth=100)
|
|
|
|
|
|
to_bw = np.array([0.299, 0.587, 0.114])
|
|
|
|
def gray(img): return np.rollaxis(img, 0, 1).dot(to_bw)
|
|
def to_plot(img): return np.rollaxis(img, 0, 1).astype(np.uint8)
|
|
def plot(img): plt.imshow(to_plot(img))
|
|
|
|
def floor(x): return int(math.floor(x))
|
|
def ceil(x): return int(math.ceil(x))
|
|
|
|
def plots(ims, figsize=(12,6), rows=1, interp=False, titles=None):
|
|
if type(ims[0]) is np.ndarray:
|
|
ims = np.array(ims).astype(np.uint8)
|
|
if (ims.shape[-1] != 3): ims = ims.transpose((0,2,3,1))
|
|
f = plt.figure(figsize=figsize)
|
|
|
|
for i in range(len(ims)):
|
|
sp = f.add_subplot(rows, len(ims)//rows, i+1)
|
|
sp.axis('Off')
|
|
if titles is not None: sp.set_title(titles[i], fontsize=16)
|
|
plt.imshow(ims[i], interpolation=None if interp else 'none')
|
|
|
|
|
|
def do_clip(arr, mx):
|
|
clipped = np.clip(arr, (1-mx)/1, mx)
|
|
return clipped/clipped.sum(axis=1)[:, np.newaxis]
|
|
|
|
|
|
def wrap_config(layer):
|
|
return {'class_name': layer.__class__.__name__, 'config': layer.get_config()}
|
|
|
|
def copy_layer(layer): return layer_from_config(wrap_config(layer))
|
|
|
|
def copy_layers(layers): return [copy_layer(layer) for layer in layers]
|
|
|
|
def copy_weights(from_layers, to_layers):
|
|
for from_layer,to_layer in zip(from_layers, to_layers):
|
|
to_layer.set_weights(from_layer.get_weights())
|
|
|
|
def save_array(fname, arr):
|
|
c=bcolz.carray(arr, rootdir=fname, mode='w')
|
|
c.flush()
|
|
|
|
def load_array(fname): return bcolz.open(fname)[:]
|
|
|
|
def get_classes(path):
|
|
batches = get_batches(path+'train', shuffle=False, batch_size=1)
|
|
val_batches = get_batches(path+'valid', shuffle=False, batch_size=1)
|
|
test_batches = get_batches(path+'test', shuffle=False, batch_size=1)
|
|
return (val_batches.classes, batches.classes, onehot(val_batches.classes), onehot(batches.classes),
|
|
val_batches.filenames, batches.filenames, test_batches.filenames)
|
|
|
|
def limit_mem():
|
|
K.get_session().close()
|
|
cfg = K.tf.ConfigProto()
|
|
cfg.gpu_options.allow_growth = True
|
|
K.set_session(K.tf.Session(config=cfg))
|
|
|
|
class MixIterator(object):
|
|
def __init__(self, iters):
|
|
self.iters = iters
|
|
self.multi = type(iters) is list
|
|
if self.multi:
|
|
self.N = sum([it[0].N for it in self.iters])
|
|
else:
|
|
self.N = sum([it.N for it in self.iters])
|
|
|
|
def reset(self):
|
|
for it in self.iters: it.reset()
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def next(self, *args, **kwargs):
|
|
if self.multi:
|
|
nexts = [[next(it) for it in o] for o in self.iters]
|
|
n0 = np.concatenate([n[0] for n in nexts])
|
|
n1 = np.concatenate([n[1] for n in nexts])
|
|
return (n0, n1)
|
|
else:
|
|
nexts = [next(it) for it in self.iters]
|
|
n0 = np.concatenate([n[0] for n in nexts])
|
|
n1 = np.concatenate([n[1] for n in nexts])
|
|
return (n0, n1)
|