Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import pickle | |
| import numpy as np | |
| import tensorflow as tf | |
| from PIL import Image | |
| import os | |
| # --- Loading the Model and Class Indices --- | |
| MODEL_PATH = 'fracture_detection_model.pkl' | |
| CLASS_INDICES_PATH = 'class_indices.pkl' | |
| # Check if model files exist before loading | |
| if not os.path.exists(MODEL_PATH) or not os.path.exists(CLASS_INDICES_PATH): | |
| raise RuntimeError(f"Model or class indices files not found. Please ensure '{MODEL_PATH}' and '{CLASS_INDICES_PATH}' are uploaded to your Space.") | |
| # Load the trained model | |
| try: | |
| with open(MODEL_PATH, 'rb') as f: | |
| model = pickle.load(f) | |
| except Exception as e: | |
| raise RuntimeError(f"Error loading the model: {e}") | |
| # Load the class indices | |
| try: | |
| with open(CLASS_INDICES_PATH, 'rb') as f: | |
| class_indices = pickle.load(f) | |
| # Invert the dictionary to map index to class name | |
| class_names = {v: k for k, v in class_indices.items()} | |
| except Exception as e: | |
| raise RuntimeError(f"Error loading class indices: {e}") | |
| print("--- Model and class indices loaded successfully! ---") | |
| # --- Prediction Function (Adapted for Gradio) --- | |
| def predict_fracture(img): | |
| """ | |
| Takes a PIL Image from Gradio, preprocesses it, | |
| and returns a dictionary of class probabilities. | |
| """ | |
| if img is None: | |
| return "Please upload an image." | |
| try: | |
| # 1. Ensure image is in RGB format | |
| if img.mode != "RGB": | |
| img = img.convert("RGB") | |
| # 2. Resize the image (matching your original 150x150 target) | |
| img = img.resize((150, 150)) | |
| # 3. Convert image to numpy array and scale pixel values | |
| img_array = tf.keras.preprocessing.image.img_to_array(img) | |
| img_array = img_array / 255.0 | |
| # 4. Expand dimensions to create a batch of 1 | |
| img_batch = np.expand_dims(img_array, axis=0) | |
| # 5. Make a prediction | |
| prediction = model.predict(img_batch) | |
| # 6. Format output for Gradio's Label component | |
| # This creates a nice bar graph of results: {"Fractured": 0.85, "Normal": 0.15} | |
| result = {} | |
| for i, class_name in class_names.items(): | |
| result[class_name] = float(prediction[0][i]) | |
| return result | |
| except Exception as e: | |
| return f"Prediction failed: {str(e)}" | |
| # --- Create the Gradio Web Interface --- | |
| demo = gr.Interface( | |
| fn=predict_fracture, | |
| inputs=gr.Image(type="pil", label="Upload X-ray Image"), | |
| # gr.Label automatically highlights the highest probability class | |
| outputs=gr.Label(num_top_classes=3, label="Analysis Result"), | |
| title="🩻 X-ray Fracture Detection System", | |
| description="Upload a bone X-ray image to identify potential fractures. Powered by TensorFlow and Gradio.", | |
| examples=[img for img in os.listdir('.') if img.endswith(('.jpg', '.png', '.jpeg'))][:2] if any(img.endswith(('.jpg', '.png', '.jpeg')) for img in os.listdir('.')) else None | |
| ) | |
| # Launch the app (Hugging Face looks for this to start the container) | |
| if __name__ == "__main__": | |
| demo.launch() |