Spaces:
Sleeping
Sleeping
| from tensorflow.keras.models import load_model | |
| from tensorflow.keras.preprocessing.image import load_img, img_to_array | |
| import numpy as np | |
| import json | |
| import os | |
| from PIL import Image | |
| import io | |
| # Load class names | |
| try: | |
| with open("models/class_names.json", "r") as f: | |
| CLASS_NAMES = json.load(f) | |
| except FileNotFoundError: | |
| print("Warning: class_names.json not found. Make sure to train the model first.") | |
| CLASS_NAMES = [] | |
| def load_and_prepare_image(img_file): | |
| """ | |
| Load and preprocess an image file for prediction. | |
| Handles both file paths and file-like objects. | |
| """ | |
| try: | |
| if isinstance(img_file, str): | |
| # If img_file is a path | |
| img = load_img(img_file, target_size=(256, 256)) | |
| else: | |
| # If img_file is a file-like object (e.g., from Streamlit upload) | |
| img = Image.open(img_file) | |
| img = img.resize((256, 256)) | |
| # Convert to array and normalize | |
| img_array = img_to_array(img) | |
| img_array = img_array / 255.0 | |
| return np.expand_dims(img_array, axis=0) | |
| except Exception as e: | |
| raise Exception(f"Error processing image: {str(e)}") | |
| def predict(model_path, img_file): | |
| """ | |
| Make a prediction using the loaded model. | |
| Returns the predicted class and confidence score. | |
| """ | |
| try: | |
| # Check if model exists | |
| if not os.path.exists(model_path): | |
| raise FileNotFoundError(f"Model file not found at {model_path}") | |
| # Load model | |
| model = load_model(model_path) | |
| # Prepare image | |
| image = load_and_prepare_image(img_file) | |
| # Make prediction | |
| preds = model.predict(image, verbose=0) | |
| # Get predicted class and confidence | |
| predicted_class_idx = np.argmax(preds[0]) | |
| confidence = preds[0][predicted_class_idx] | |
| # Get class name | |
| if CLASS_NAMES and predicted_class_idx < len(CLASS_NAMES): | |
| predicted_class = CLASS_NAMES[predicted_class_idx] | |
| else: | |
| predicted_class = f"Class_{predicted_class_idx}" | |
| return predicted_class, confidence | |
| except Exception as e: | |
| raise Exception(f"Prediction error: {str(e)}") | |
| def get_class_probabilities(model_path, img_file): | |
| """ | |
| Get probability distribution across all classes. | |
| Returns a dictionary of class names and their probabilities. | |
| """ | |
| try: | |
| model = load_model(model_path) | |
| image = load_and_prepare_image(img_file) | |
| preds = model.predict(image, verbose=0)[0] | |
| probabilities = {} | |
| for idx, prob in enumerate(preds): | |
| class_name = CLASS_NAMES[idx] if CLASS_NAMES and idx < len(CLASS_NAMES) else f"Class_{idx}" | |
| probabilities[class_name] = float(prob) | |
| return probabilities | |
| except Exception as e: | |
| raise Exception(f"Error getting class probabilities: {str(e)}") | |