Spaces:
Sleeping
Sleeping
File size: 4,397 Bytes
e76a638 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
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
|