|
|
|
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
|
|
|
|
|
|
MODEL_REPO = "Ahmedhassan54/Image-Classification"
|
|
MODEL_FILE = "best_model.h5"
|
|
|
|
|
|
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}")
|
|
|
|
|
|
model = tf.keras.models.load_model(MODEL_FILE)
|
|
print("Model loaded successfully!")
|
|
return model
|
|
except Exception as e:
|
|
print(f"Error loading model: {str(e)}")
|
|
raise
|
|
|
|
|
|
model = load_model_from_hf()
|
|
|
|
|
|
def classify_image(image):
|
|
try:
|
|
|
|
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])
|
|
|
|
if confidence > 0.5:
|
|
return {
|
|
"Dog": confidence * 100,
|
|
"Cat": (1 - confidence) * 100
|
|
}
|
|
else:
|
|
return {
|
|
"Cat": (1 - confidence) * 100,
|
|
"Dog": confidence * 100
|
|
}
|
|
except Exception as e:
|
|
return f"Error processing image: {str(e)}"
|
|
|
|
|
|
demo = gr.Interface(
|
|
fn=classify_image,
|
|
inputs=gr.Image(type="pil", label="Upload Image"),
|
|
outputs=gr.Label(num_top_classes=2, label="Predictions"),
|
|
title="๐ฑ Cat vs Dog Classifier ๐ถ",
|
|
description="Upload an image to classify whether it's a cat or dog",
|
|
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"]
|
|
],
|
|
allow_flagging="never"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
demo.launch(debug=True, server_port=7860) |