Spaces:
Sleeping
Sleeping
import streamlit as st | |
import numpy as np | |
import tensorflow as tf | |
import tensorflow_hub as hub | |
import cv2 | |
from PIL import Image | |
import time | |
# Set page title and favicon | |
st.set_page_config(page_title="Cat and Dog Classifier", page_icon="๐ฑ๐ถ") | |
# Load the pre-trained model | |
mobilenet_model = 'https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4' | |
pretrained_model = hub.KerasLayer(mobilenet_model, input_shape=(224, 224, 3), trainable=False) | |
num_of_classes = 2 | |
model = tf.keras.Sequential([ | |
pretrained_model, | |
tf.keras.layers.Dense(num_of_classes) | |
]) | |
model.compile( | |
optimizer='adam', | |
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), | |
metrics=['acc'] | |
) | |
model.load_weights("cat_dog_classifier.h5") | |
# Define functions for image resizing and classification | |
def preprocess_image(image): | |
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
image = cv2.resize(image, (224, 224)) | |
image = image / 255.0 | |
image = np.expand_dims(image, axis=0) | |
return image | |
def classify_image(image): | |
image = preprocess_image(image) | |
prediction = model.predict(image) | |
return np.argmax(prediction) | |
# Sidebar | |
st.sidebar.header("Cat and Dog Classifier") | |
uploaded_image = st.sidebar.file_uploader("Upload an image", type=["jpg", "png", "jpeg"]) | |
# Main content | |
st.title("Cat/Dog Image Classification") | |
if uploaded_image: | |
with st.spinner("Uploading image..."): | |
time.sleep(2) # Simulate image upload process, replace with actual image upload logic | |
st.success("Image upload complete!") | |
image = Image.open(uploaded_image) | |
st.image(image, caption="Uploaded Image", use_column_width=True) | |
if st.button("Classify"): | |
image = np.array(image) | |
pred_label = classify_image(image) | |
if pred_label == 0: | |
st.write('<div style="font-size: 24px; color: white;">Prediction: It\'s a Cat ๐บ</div>', unsafe_allow_html=True) | |
else: | |
st.write('<div style="font-size: 24px; color: white;">Prediction: It\'s a Dog ๐ถ</div>', unsafe_allow_html=True) | |
# Add a footer with CSS for positioning | |
st.markdown( | |
""" | |
<div style="position: fixed; bottom: 0; right: 10px; padding: 10px; color: white;"> | |
<a href="https://github.com/sg-sparsh-goyal" target="_blank" style="color: white; text-decoration: none;"> | |
โจ Github | |
</a><br> | |
By Sparsh | |
</div> | |
""", | |
unsafe_allow_html=True | |
) | |