Spaces:
Sleeping
Sleeping
| import tensorflow as tf | |
| import numpy as np | |
| from PIL import Image | |
| import joblib | |
| import os | |
| from tensorflow.keras.models import Sequential | |
| from tensorflow.keras.layers import Conv2D, MaxPooling2D, Dense, Flatten | |
| class BoundingBoxPredictor: | |
| def __init__(self): | |
| self.cnn_model = None | |
| self.knn_model = None | |
| self.input_shape = (128, 128) | |
| # Updated normalization values based on actual dataset analysis | |
| self.y_min = np.array([0., 0., 30., 30.]) # min values for x, y, width, height | |
| self.y_max = np.array([97., 97., 49., 49.]) # max values for x, y, width, height | |
| def create_cnn_model(self): | |
| model = Sequential([ | |
| Conv2D(32, (3, 3), activation='relu', input_shape=(*self.input_shape, 1)), | |
| MaxPooling2D((2, 2)), | |
| Conv2D(64, (3, 3), activation='relu'), | |
| MaxPooling2D((2, 2)), | |
| Flatten(), | |
| Dense(128, activation='relu'), | |
| Dense(4, activation='sigmoid') | |
| ]) | |
| model.compile(optimizer='adam', loss='mse', metrics=['mse']) | |
| return model | |
| def load_models(self, cnn_path, knn_path): | |
| # Check if files exist | |
| if not os.path.exists(cnn_path): | |
| raise FileNotFoundError(f"CNN model file not found: {cnn_path}") | |
| if not os.path.exists(knn_path): | |
| raise FileNotFoundError(f"KNN model file not found: {knn_path}") | |
| print(f"Loading CNN model from: {cnn_path}") | |
| try: | |
| self.cnn_model = tf.keras.models.load_model(cnn_path) | |
| except Exception as e: | |
| print(f"Error loading CNN model: {str(e)}") | |
| print("Creating new CNN model...") | |
| self.cnn_model = self.create_cnn_model() | |
| try: | |
| self.cnn_model.load_weights(cnn_path) | |
| print("Successfully loaded CNN weights") | |
| except: | |
| print("Could not load CNN weights, using uninitialized model") | |
| print("CNN model ready") | |
| print(f"Loading KNN model from: {knn_path}") | |
| self.knn_model = joblib.load(knn_path) | |
| print("KNN model loaded successfully") | |
| def preprocess_image(self, image): | |
| # Convert to grayscale if needed | |
| if image.mode != 'L': | |
| image = image.convert('L') | |
| # Resize image | |
| image = image.resize(self.input_shape) | |
| # Convert to numpy array and normalize | |
| img_array = np.array(image) | |
| img_array = img_array / 255.0 | |
| return img_array | |
| def predict(self, image, model_type='cnn'): | |
| # Check if models are loaded | |
| if model_type == 'cnn' and self.cnn_model is None: | |
| raise ValueError("CNN model not loaded. Please call load_models first.") | |
| if model_type == 'knn' and self.knn_model is None: | |
| raise ValueError("KNN model not loaded. Please call load_models first.") | |
| # Preprocess the image | |
| processed_image = self.preprocess_image(image) | |
| try: | |
| if model_type == 'cnn': | |
| # Reshape for CNN input | |
| img_array = processed_image.reshape(1, *self.input_shape, 1) | |
| # Get normalized predictions (between 0 and 1) | |
| prediction = self.cnn_model.predict(img_array)[0] | |
| # Denormalize predictions to original scale | |
| prediction = prediction * (self.y_max - self.y_min) + self.y_min | |
| else: # KNN | |
| # Flatten for KNN input | |
| img_array = processed_image.flatten().reshape(1, -1) | |
| # Get normalized predictions | |
| prediction = self.knn_model.predict(img_array)[0] | |
| # Denormalize predictions to original scale | |
| prediction = prediction * (self.y_max - self.y_min) + self.y_min | |
| # Ensure predictions are within valid ranges | |
| prediction = np.clip(prediction, self.y_min, self.y_max) | |
| return { | |
| 'x': float(prediction[0]), | |
| 'y': float(prediction[1]), | |
| 'width': float(prediction[2]), | |
| 'height': float(prediction[3]) | |
| } | |
| except Exception as e: | |
| print(f"Error during prediction with {model_type.upper()} model: {str(e)}") | |
| raise | |