cat-vs-dog / app.py
eddydecena's picture
Create app.py
54c0802
raw history blame
No virus
1.71 kB
import os
import gradio as gr
import tensorflow as tf
from keras_tuner import HyperParameters
from huggingface_hub import hf_hub_download
from src.models import MakeHyperModel
from src.preprocessing import get_data_augmentation
from src.config import IMAGE_SIZE
data_augmentation = get_data_augmentation()
cache_dir = os.path.join('hf_hub')
for f in ['checkpoint', 'checkpoint.data-00000-of-00001', 'checkpoint.index']:
print(f)
old_name = hf_hub_download(repo_id="eddydecena/cat-vs-dog", filename=f"tuner_model/cat-vs-dog/trial_0484d8d758a5ef7b91ca97d334ba7870/checkpoints/epoch_0/{f}", cache_dir=cache_dir)
temp_value = old_name.split('/')
temp_value.pop(-1)
path = '/'.join(temp_value)
os.rename(old_name, os.path.join(path, f))
latest = tf.train.latest_checkpoint('./tuner_model/cat-vs-dog/trial_0484d8d758a5ef7b91ca97d334ba7870/checkpoints/epoch_0')
hypermodel = MakeHyperModel(input_shape=IMAGE_SIZE + (3,), num_classes=2, data_augmentation=data_augmentation)
model = hypermodel.build(hp=HyperParameters())
model.load_weights(latest).expect_partial()
def cat_vs_dog(image):
img_array = tf.constant(image, dtype=tf.float32)
img_array = tf.expand_dims(img_array, 0)
predictions = model.predict(img_array)
score = predictions[0]
return {'cat': float((1 - score)), 'dog': float(score)}
iface = gr.Interface(
cat_vs_dog,
gr.inputs.Image(shape=IMAGE_SIZE),
gr.outputs.Label(num_top_classes=2),
capture_session=True,
interpretation="default",
examples=[
["examples/cat1.jpg"],
["examples/cat2.jpg"],
["examples/dog1.jpeg"],
["examples/dog2.jpeg"]
])
if __name__ == "__main__":
iface.launch()