|
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
|
|
import pandas as pd
|
|
import logging
|
|
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
MODEL_REPO = "Ahmedhassan54/Image-Classification"
|
|
MODEL_FILE = "best_model.h5"
|
|
|
|
|
|
model = None
|
|
|
|
def load_model():
|
|
global model
|
|
try:
|
|
logger.info("โณ Downloading model...")
|
|
model_path = hf_hub_download(
|
|
repo_id=MODEL_REPO,
|
|
filename=MODEL_FILE,
|
|
cache_dir=".",
|
|
force_download=True
|
|
)
|
|
logger.info(f"๐ Model path: {model_path}")
|
|
|
|
|
|
if not os.path.exists(model_path):
|
|
raise FileNotFoundError(f"Model file not found at {model_path}")
|
|
|
|
logger.info("๐ Loading TensorFlow model...")
|
|
model = tf.keras.models.load_model(model_path)
|
|
logger.info("โ
Model loaded successfully!")
|
|
|
|
except Exception as e:
|
|
logger.error(f"โ Model loading failed: {str(e)}")
|
|
model = None
|
|
raise gr.Error(f"Model loading failed. Check logs for details.")
|
|
|
|
|
|
load_model()
|
|
|
|
def classify_image(image):
|
|
try:
|
|
if image is None:
|
|
raise gr.Error("Please upload an image first")
|
|
|
|
logger.info("๐ผ๏ธ Processing image...")
|
|
|
|
|
|
if isinstance(image, np.ndarray):
|
|
image = Image.fromarray(image)
|
|
|
|
|
|
image = image.resize((150, 150))
|
|
img_array = np.array(image) / 255.0
|
|
if len(img_array.shape) == 3:
|
|
img_array = np.expand_dims(img_array, axis=0)
|
|
|
|
logger.info(f"๐ Input shape: {img_array.shape}")
|
|
|
|
if model is None:
|
|
raise gr.Error("Model not loaded - using demo mode")
|
|
return {"Cat": 0.5, "Dog": 0.5}, pd.DataFrame({'Class': ['Cat', 'Dog'], 'Confidence': [0.5, 0.5]})
|
|
|
|
pred = model.predict(img_array, verbose=0)
|
|
confidence = float(pred[0][0])
|
|
logger.info(f"๐ฎ Prediction confidence: {confidence}")
|
|
|
|
results = {
|
|
"Cat": round(1 - confidence, 4),
|
|
"Dog": round(confidence, 4)
|
|
}
|
|
|
|
plot_data = pd.DataFrame({
|
|
'Class': ['Cat', 'Dog'],
|
|
'Confidence': [1 - confidence, confidence]
|
|
})
|
|
|
|
return results, plot_data
|
|
|
|
except Exception as e:
|
|
logger.error(f"๐ฅ Classification error: {str(e)}")
|
|
raise gr.Error(f"Error processing image: {str(e)}")
|
|
|
|
css = """
|
|
.gradio-container { max-width: 900px; margin: auto; }
|
|
footer { visibility: hidden; }
|
|
.progress-bar { color: #ff4d4d !important; }
|
|
"""
|
|
|
|
with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
|
|
gr.Markdown("""
|
|
# ๐พ Cat vs Dog Classifier ๐ฆฎ
|
|
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")
|
|
with gr.Row():
|
|
submit_btn = gr.Button("Classify ๐", variant="primary")
|
|
clear_btn = gr.Button("Clear ๐๏ธ")
|
|
|
|
with gr.Column():
|
|
label_output = gr.Label(label="Predictions")
|
|
confidence_bar = gr.BarPlot(
|
|
pd.DataFrame({'Class': ['Cat', 'Dog'], 'Confidence': [0.5, 0.5]}),
|
|
x="Class", y="Confidence", y_lim=[0,1],
|
|
title="Confidence Scores", width=400, height=300
|
|
)
|
|
|
|
|
|
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,
|
|
outputs=[label_output, confidence_bar],
|
|
fn=classify_image,
|
|
cache_examples=True
|
|
)
|
|
|
|
|
|
submit_btn.click(
|
|
fn=classify_image,
|
|
inputs=image_input,
|
|
outputs=[label_output, confidence_bar],
|
|
api_name="predict"
|
|
)
|
|
|
|
clear_btn.click(
|
|
fn=lambda: [None, pd.DataFrame({'Class': ['Cat', 'Dog'], 'Confidence': [0.5, 0.5]})],
|
|
inputs=None,
|
|
outputs=[image_input, confidence_bar]
|
|
)
|
|
|
|
if __name__ == "__main__":
|
|
demo.launch(debug=True) |