| 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 | |
| os.environ['CUDA_VISIBLE_DEVICES'] = '-1' | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| MODEL_REPO = "Ahmedhassan54/Image-Classification-Model" | |
| MODEL_FILE = "best_model.keras" | |
| 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}") | |
| with tf.device('/CPU:0'): | |
| 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 | |
| load_model() | |
| def classify_image(image): | |
| try: | |
| if image is None: | |
| return {"Cat": 0.5, "Dog": 0.5}, pd.DataFrame({'Class': ['Cat', 'Dog'], 'Confidence': [0.5, 0.5]}) | |
| if isinstance(image, np.ndarray): | |
| image = Image.fromarray(image.astype('uint8')) | |
| 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) | |
| if model is not None: | |
| with tf.device('/CPU:0'): | |
| pred = model.predict(img_array, verbose=0) | |
| confidence = float(pred[0][0]) | |
| else: | |
| confidence = 0.75 | |
| 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"Error: {str(e)}") | |
| return {"Error": str(e)}, pd.DataFrame({'Class': ['Cat', 'Dog'], 'Confidence': [0.5, 0.5]}) | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# 🐾 Cat vs Dog Classifier 🦮") | |
| with gr.Row(): | |
| with gr.Column(): | |
| img_input = gr.Image(type="pil") | |
| classify_btn = gr.Button("Classify", variant="primary") | |
| with gr.Column(): | |
| label_out = gr.Label(num_top_classes=2) | |
| plot_out = gr.BarPlot( | |
| pd.DataFrame({'Class': ['Cat', 'Dog'], 'Confidence': [0.5, 0.5]}), | |
| x="Class", y="Confidence", y_lim=[0,1] | |
| ) | |
| classify_btn.click( | |
| classify_image, | |
| inputs=img_input, | |
| outputs=[label_out, plot_out] | |
| ) | |
| 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=img_input, | |
| outputs=[label_out, plot_out], | |
| fn=classify_image, | |
| cache_examples=True | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0", server_port=7860) |