Spaces:
Sleeping
Sleeping
import gradio as gr | |
import tensorflow as tf | |
import numpy as np | |
from PIL import Image | |
from huggingface_hub import hf_hub_download | |
import os | |
# Configuration | |
MODEL_REPO = "Ahmedhassan54/Image-Classification" # Replace with your HF username and repo | |
MODEL_FILE = "best_model.h5" | |
# Download model from Hugging Face Hub | |
def load_model_from_hf(): | |
try: | |
if not os.path.exists(MODEL_FILE): | |
print("Downloading model from Hugging Face Hub...") | |
model_path = hf_hub_download( | |
repo_id=MODEL_REPO, | |
filename=MODEL_FILE, | |
cache_dir="." | |
) | |
os.system(f"cp {model_path} {MODEL_FILE}") | |
return tf.keras.models.load_model(MODEL_FILE) | |
except Exception as e: | |
raise gr.Error(f"Model loading failed: {str(e)}") | |
model = load_model_from_hf() | |
def classify_image(image): | |
try: | |
image = Image.fromarray(image) if isinstance(image, np.ndarray) else image | |
image = image.resize((150, 150)) | |
image_array = np.array(image) / 255.0 | |
image_array = np.expand_dims(image_array, axis=0) | |
prediction = model.predict(image_array) | |
confidence = float(prediction[0][0]) | |
return { | |
"Dog": confidence, | |
"Cat": 1 - confidence | |
} | |
except Exception as e: | |
raise gr.Error(f"Classification error: {str(e)}") | |
# Custom CSS for better UI | |
css = """ | |
.gradio-container { | |
background: linear-gradient(to right, #f5f7fa, #c3cfe2); | |
} | |
footer { | |
visibility: hidden | |
} | |
""" | |
# Build the interface | |
with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo: | |
gr.Markdown("# 🐾 Cat vs Dog Classifier 🦮") | |
gr.Markdown("Upload an image to classify whether it's a cat or dog") | |
with gr.Row(): | |
with gr.Column(): | |
image_input = gr.Image(label="Upload Image", type="pil") | |
submit_btn = gr.Button("Classify", variant="primary") | |
with gr.Column(): | |
label_output = gr.Label(label="Predictions", num_top_classes=2) | |
confidence_bar = gr.BarPlot( | |
x=["Cat", "Dog"], | |
y=[0.5, 0.5], | |
y_lim=[0,1], | |
title="Confidence Scores", | |
width=400, | |
height=300 | |
) | |
# Example images | |
gr.Examples( | |
examples=[ | |
["https://upload.wikimedia.org/wikipedia/commons/1/15/Cat_August_2010-4.jpg"], | |
["https://upload.wikimedia.org/wikipedia/commons/d/d9/Collage_of_Nine_Dogs.jpg"] | |
], | |
inputs=image_input | |
) | |
# Button action | |
submit_btn.click( | |
fn=classify_image, | |
inputs=image_input, | |
outputs=[label_output, confidence_bar], | |
api_name="classify" | |
) | |
if __name__ == "__main__": | |
demo.launch() |