Spaces:
Runtime error
Runtime error
File size: 17,079 Bytes
07fc447 |
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 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 |
import logging
import joblib
import tensorflow as tf
from pathlib import Path
from sklearn.preprocessing import MultiLabelBinarizer
import cv2
import numpy as np
import logging
import cv2
import keras
from pathlib import Path
import tensorflow as tf
from typing import Optional, Tuple, List
from app.src.logger import setup_logger
# Configure logging
logger = setup_logger("vgg16_load")
def load_vgg_artifacts(model_path: Path, mlb_path: Path) -> tuple[tf.keras.Model, MultiLabelBinarizer]:
"""
Loads the VGG model and the MultiLabelBinarizer from specified paths.
Args:
model_path: Path to the VGG model file (.keras).
mlb_path: Path to the MultiLabelBinarizer file (.joblib).
Returns:
A tuple containing the loaded Keras model and MultiLabelBinarizer object.
Raises:
FileNotFoundError: If either the model file or the MLB file is not found.
Exception: If any other error occurs during loading.
"""
model = None
mlb = None
try:
logger.info(f"Attempting to load VGG model from {model_path}")
model = tf.keras.models.load_model(model_path)
logger.info("VGG model loaded successfully.")
except FileNotFoundError:
logger.error(f"Error: VGG model file not found at {model_path}")
raise
except Exception as e:
logger.error(f"An error occurred while loading the VGG model: {e}")
raise
try:
logger.info(f"Attempting to load MultiLabelBinarizer from {mlb_path}")
mlb = joblib.load(mlb_path)
logger.info("MultiLabelBinarizer loaded successfully.")
except FileNotFoundError:
logger.error(f"Error: MultiLabelBinarizer file not found at {mlb_path}")
raise
except Exception as e:
logger.error(f"An error occurred while loading the MultiLabelBinarizer: {e}")
raise
logger.info("Both VGG model and MultiLabelBinarizer loaded successfully.")
return model, mlb
def preprocess_image(image_path: Path, target_size: tuple[int, int] = (224, 224)) -> np.ndarray | None:
"""
Preprocesses an image for VGG model prediction.
Loads an image from the specified path, converts it to RGB, resizes it,
and normalizes pixel values. Includes robust error handling and logging
at each step.
Args:
image_path: Path to the image file.
target_size: A tuple (width, height) specifying the desired output size.
Returns:
A preprocessed NumPy array representing the image with pixel values
scaled between 0 and 1, or None if an error occurred during processing.
"""
try:
logger.info(f"Attempting to load image from {image_path}")
img = cv2.imread(str(image_path)) # cv2.imread expects a string or numpy array
if img is None:
logger.error(f"Error: Could not load image from {image_path}. cv2.imread returned None.")
return None
logger.info("Image loaded successfully.")
logger.info("Attempting to convert image to RGB.")
# Check if the image is already in a format that doesn't need BGR to RGB conversion
# cv2.imread loads in BGR format by default for color images.
# If the image is grayscale, it might be loaded as such.
# We want RGB for consistency with models trained on RGB data.
if len(img.shape) == 3 and img.shape[2] == 3: # Check if it's a color image (likely BGR)
try:
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
logger.info("Image converted to RGB successfully.")
except cv2.error as e:
logger.error(f"Error during BGR to RGB conversion for image {image_path}: {e}")
return None
elif len(img.shape) == 2: # Grayscale image
try:
img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
logger.info("Grayscale image converted to RGB successfully.")
except cv2.error as e:
logger.error(f"Error during Grayscale to RGB conversion for image {image_path}: {e}")
return None
else:
logger.warning(f"Unexpected image format for {image_path}. Attempting to proceed.")
# If it's not a standard color or grayscale, we might proceed but log a warning.
# Depending on requirements, you might want to return None here.
logger.info(f"Attempting to resize image to {target_size}.")
try:
img = cv2.resize(img, target_size)
if img is None or img.size == 0:
logger.error(f"Error: cv2.resize returned None or empty array for image {image_path}.")
return None
logger.info("Image resized successfully.")
except cv2.error as e:
logger.error(f"Error during image resizing for image {image_path} to size {target_size}: {e}")
return None
logger.info("Attempting to normalize pixel values.")
try:
# Ensure the image is the correct dtype before division
img = img.astype("float32") / 255.0
if img is None or img.size == 0 or np.max(img) > 1.0 or np.min(img) < 0.0:
logger.error(f"Error: Image normalization failed or resulted in unexpected values for image {image_path}.")
return None
logger.info("Pixel values normalized successfully.")
except Exception as e:
logger.error(f"Error during pixel normalization for image {image_path}: {e}")
return None
logger.info(f"Image preprocessing completed successfully for {image_path}.")
return img
except Exception as e:
logger.error(f"An unexpected error occurred during image preprocessing for {image_path}: {e}")
return None
class VGGDocumentClassifier:
"""
A class for classifying documents using a VGG16 model.
This class encapsulates the loading of the VGG16 model and its associated
MultiLabelBinarizer, provides a method to preprocess input images, and
performs document classification predictions.
"""
def __init__(self, model_path: Path, mlb_path: Path, target_size: Tuple[int, int] = (224, 224)) -> None:
"""
Initializes the VGGDocumentClassifier by loading the model and MLB.
Args:
model_path: Path to the VGG model file (.keras).
mlb_path: Path to the MultiLabelBinarizer file (.joblib).
target_size: The target size (width, height) for image preprocessing.
Defaults to (224, 224).
Raises:
FileNotFoundError: If either the model file or the MLB file is not found.
Exception: If any other error occurs during loading.
"""
logger.info("Initializing VGGDocumentClassifier.")
self.model: Optional[tf.keras.Model] = None
self.mlb: Optional[MultiLabelBinarizer] = None
self.target_size: Tuple[int, int] = target_size
try:
self._load_artifacts(model_path, mlb_path)
if self.model and self.mlb:
logger.info("VGGDocumentClassifier initialized successfully.")
else:
logger.critical("VGGDocumentClassifier failed to fully initialize due to artifact loading errors.")
raise RuntimeError("Failed to load all required artifacts for VGGDocumentClassifier.")
except Exception as e:
logger.critical(f"Failed to initialize VGGDocumentClassifier: {e}", exc_info=True)
raise # Re-raise the exception after logging
def _load_artifacts(self, model_path: Path, mlb_path: Path) -> None:
"""
Loads the VGG model and MultiLabelBinarizer with error handling and logging.
Args:
model_path: Path to the VGG model file (.keras).
mlb_path: Path to the MultiLabelBinarizer file (.joblib).
Raises:
FileNotFoundError: If either the model file or the MLB file is not found.
Exception: If any other unexpected error occurs during loading.
"""
logger.info("Starting artifact loading.")
model_loaded: bool = False
mlb_loaded: bool = False
# Load Model
try:
logger.info(f"Attempting to load VGG model from {model_path}")
self.model = tf.keras.models.load_model(model_path)
logger.info("VGG model loaded successfully.")
model_loaded = True
except FileNotFoundError:
logger.critical(f"Critical Error: VGG model file not found at {model_path}", exc_info=True)
raise # Re-raise to indicate a critical initialization failure
except Exception as e:
logger.critical(f"Critical Error: An unexpected error occurred while loading the VGG model from {model_path}: {e}", exc_info=True)
raise # Re-raise to indicate a critical initialization failure
# Load MLB
try:
logger.info(f"Attempting to load MultiLabelBinarizer from {mlb_path}")
self.mlb = joblib.load(mlb_path)
logger.info("MultiLabelBinarizer loaded successfully.")
mlb_loaded = True
except FileNotFoundError:
logger.critical(f"Critical Error: MultiLabelBinarizer file not found at {mlb_path}", exc_info=True)
raise # Re-raise to indicate a critical initialization failure
except Exception as e:
logger.critical(f"Critical Error: An unexpected error occurred while loading the MultiLabelBinarizer from {mlb_path}: {e}", exc_info=True)
raise # Re-raise to indicate a critical initialization failure
if model_loaded and mlb_loaded:
logger.info("All required VGG artifacts loaded successfully.")
else:
logger.error("One or more required VGG artifacts failed to load during _load_artifacts.")
def preprocess_image(self, image_path: Path) -> Optional[np.ndarray]:
"""
Preprocesses an image for VGG model prediction.
Loads an image from the specified path, converts it to RGB, resizes it,
and normalizes pixel values. Includes robust error handling and logging
at each step.
Args:
image_path: Path to the image file.
Returns:
A preprocessed NumPy array representing the image with pixel values
scaled between 0 and 1, or None if an error occurred during processing.
"""
try:
logger.info(f"Attempting to load image from {image_path}")
img = cv2.imread(str(image_path)) # cv2.imread expects a string or numpy array
if img is None:
logger.error(f"Error: Could not load image from {image_path}. cv2.imread returned None.")
return None
logger.info("Image loaded successfully.")
logger.info("Attempting to convert image to RGB.")
if len(img.shape) == 3 and img.shape[2] == 3: # Check if it's a color image (likely BGR)
try:
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
logger.info("Image converted to RGB successfully.")
except cv2.error as e:
logger.error(f"Error during BGR to RGB conversion for image {image_path}: {e}")
return None
elif len(img.shape) == 2: # Grayscale image
try:
img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
logger.info("Grayscale image converted to RGB successfully.")
except cv2.error as e:
logger.error(f"Error during Grayscale to RGB conversion for image {image_path}: {e}")
return None
else:
logger.warning(f"Unexpected image format for {image_path}. Attempting to proceed.")
logger.info(f"Attempting to resize image to {self.target_size}.")
try:
img = cv2.resize(img, self.target_size)
if img is None or img.size == 0:
logger.error(f"Error: cv2.resize returned None or empty array for image {image_path}.")
return None
logger.info("Image resized successfully.")
except cv2.error as e:
logger.error(f"Error during image resizing for image {image_path} to size {self.target_size}: {e}")
return None
logger.info("Attempting to normalize pixel values.")
try:
img = img.astype("float32") / 255.0
if img is None or img.size == 0 or np.max(img) > 1.0 or np.min(img) < 0.0:
logger.error(f"Error: Image normalization failed or resulted in unexpected values for image {image_path}.")
return None
logger.info("Pixel values normalized successfully.")
except Exception as e:
logger.error(f"Error during pixel normalization for image {image_path}: {e}")
return None
logger.info(f"Image preprocessing completed successfully for {image_path}.")
return img
except Exception as e:
logger.error(f"An unexpected error occurred during image preprocessing for {image_path}: {e}")
return None
def predict(self, image_path: Path) -> Optional[List[str]]:
"""
Predicts the class labels for a given image using the loaded VGG model.
The process involves loading and preprocessing the image, performing
inference with the model, and converting the prediction to class labels
using the MultiLabelBinarizer.
Args:
image_path: Path to the image file to classify.
Returns:
A list of predicted class labels (strings) if the prediction process
is successful. Returns None if any critical step (image loading,
preprocessing, model inference, or inverse transform) fails.
Returns an empty list if the prediction process is successful but
no labels are predicted.
"""
logger.info(f"Starting prediction process for image: {image_path}.")
if self.model is None or self.mlb is None:
logger.error("Model or MultiLabelBinarizer not loaded. Cannot perform prediction.")
return None
# Preprocess image
image = self.preprocess_image(image_path)
if image is None:
logger.error(f"Image preprocessing failed for {image_path}. Cannot perform prediction.")
return None
try:
logger.info(f"Performing model inference for {image_path}.")
# Add batch dimension to the image
image = np.expand_dims(image, axis=0)
prd = self.model.predict(image)
logger.info(f"Model inference completed for {image_path}. Prediction shape: {prd.shape}")
except Exception as e:
logger.error(f"An error occurred during model inference for {image_path}: {e}", exc_info=True)
return None
# Convert the prediction to a binary indicator format and get labels
try:
logger.info(f"Converting prediction to labels for {image_path}.")
# Assuming multi-class classification for now, taking the argmax
# If it's multi-label, you'd apply a sigmoid and thresholding here
pred_id = np.argmax(prd, axis=1)
# Create a zero array with the shape (1, number of classes)
binary_prediction = np.zeros((1, len(self.mlb.classes_)))
# Set the index of the predicted class to 1
binary_prediction[0, pred_id] = 1
predicted_labels_tuple_list: List[Tuple[str, ...]] = self.mlb.inverse_transform(binary_prediction)
logger.info(f"Prediction processed for {image_path}. Predicted labels (raw tuple list): {predicted_labels_tuple_list}")
if predicted_labels_tuple_list and len(predicted_labels_tuple_list) > 0:
final_labels: List[str] = list(predicted_labels_tuple_list[0])
logger.info(f"Final predicted labels for {image_path}: {final_labels}")
return final_labels
else:
logger.warning(f"MLB inverse_transform returned an empty list for {image_path}. No labels predicted.")
return []
except Exception as e:
logger.error(f"An error occurred during inverse transform or label processing for {image_path}: {e}", exc_info=True)
return None
|