Spaces:
Sleeping
Sleeping
import gradio as gr | |
import tensorflow as tf | |
import cv2 | |
import numpy as np | |
# from transformers import pipeline | |
from keras_preprocessing.image import img_to_array | |
# pipeline = pipeline(task="dogs-cat-classification", model="carlos-pino/dogs-cats") | |
# Get model | |
model = './saves/dogs-cats.h5' | |
CNN_MODEL = tf.keras.models.load_model(model) | |
weight_model = CNN_MODEL.get_weights() | |
CNN_MODEL.set_weights(weight_model) | |
IMAGE_SIZE = 100 | |
def predict(frame): | |
img_frame = cv2.imread(frame) | |
# Parse to gray | |
img = cv2.cvtColor(img_frame, cv2.COLOR_BGR2GRAY) | |
img = cv2.resize(img, (IMAGE_SIZE, IMAGE_SIZE), interpolation=cv2.INTER_CUBIC) | |
# Img normalized | |
img = np.array(img).astype(float) / 255 | |
# Parse to 2D array | |
image = img_to_array(img) | |
image = np.expand_dims(image, axis=0) | |
# Predict | |
prediction = CNN_MODEL.predict(image) | |
prediction = prediction[0][0] | |
print(prediction) | |
prediction_cat = 0 | |
prediction_dog = 0 | |
predictions = [] | |
# Classification | |
if prediction <= 0.5: | |
prediction_cat = (1 - prediction) / 0.5 | |
prediction_dog = 1 - prediction_cat | |
else: | |
prediction_cat = (prediction - 1) / 0.5 | |
prediction_dog = (prediction - 0.5) / 0.5 | |
predictions.append(dict(zip(["label", "score"], ["Cat", str(prediction_cat)]))) | |
predictions.append(dict(zip(["label", "score"], ["Dog", str(prediction_dog)]))) | |
return {p["label"]: p["score"] for p in predictions} | |
gr.Interface( | |
predict, | |
inputs=gr.inputs.Image(label="Upload cat or dog candidate", type="filepath"), | |
outputs=gr.outputs.Label(num_top_classes=2), | |
title="Dog or Cat?", | |
).launch() |