Skip to content

Instantly share code, notes, and snippets.

@devforfu
Created April 8, 2018 09:29
Show Gist options
  • Save devforfu/917bde2ab9386153a617b19e05dc91e0 to your computer and use it in GitHub Desktop.
Save devforfu/917bde2ab9386153a617b19e05dc91e0 to your computer and use it in GitHub Desktop.
For Medium post on Dog Breed Identification
class FeaturesExtractor:
"""Runs pretrained model without top layers on dataset and saves generated
bottleneck features onto disk.
"""
def __init__(self, build_fn, preprocess_fn, source,
target_size=(299, 299, 3), batch_size=128):
self.build_fn = build_fn
self.preprocess_fn = preprocess_fn
self.source = source
self.target_size = target_size
self.batch_size = batch_size
def __call__(self, folder, filename, pool='avg'):
model = self.build_fn(weights='imagenet', include_top=False, pooling=pool)
stream = self.source(
folder=folder, target_size=self.target_size,
batch_size=self.batch_size, infinite=False)
batches = []
with tqdm.tqdm_notebook(total=stream.steps_per_epoch) as bar:
for x_batch, y_batch in stream:
x_preprocessed = self.preprocess_fn(x_batch)
batch = model.predict_on_batch(x_preprocessed)
batches.append(batch)
bar.update(1)
all_features = np.vstack(batches)
np.save(filename, all_features)
return filename
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment