Spaces:
Sleeping
Sleeping
| import os | |
| import sys | |
| import warnings | |
| # ⚡ LAZY IMPORT: TensorFlow is only imported when actually needed | |
| # This prevents slow startup time when TensorFlow is not required | |
| # Import preprocessing functions | |
| from app.captcha_solver.preprocess import preprocess_image, decode_predictions, get_vocab_info | |
| from app.core.logger import setup_logging | |
| logger = setup_logging() | |
| # Lazy import tensorflow | |
| _tensorflow_imported = False | |
| _tf = None | |
| _keras = None | |
| _CTCLayer = None | |
| def _lazy_import_tensorflow(): | |
| """Lazy import TensorFlow only when needed""" | |
| global _tensorflow_imported, _tf, _keras, _CTCLayer | |
| if not _tensorflow_imported: | |
| try: | |
| import tensorflow as tf | |
| from tensorflow import keras | |
| _tf = tf | |
| _keras = keras | |
| # Suppress warnings | |
| warnings.filterwarnings('ignore') | |
| os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' | |
| # Define CTCLayer after TensorFlow is imported | |
| class CTCLayer(keras.layers.Layer): | |
| def __init__(self, name=None, **kwargs): | |
| super().__init__(name=name, **kwargs) | |
| self.loss_fn = keras.backend.ctc_batch_cost | |
| def call(self, y_true, y_pred, input_length, label_length): | |
| loss = self.loss_fn(y_true, y_pred, input_length, label_length) | |
| self.add_loss(loss) | |
| return y_pred | |
| def get_config(self): | |
| config = super().get_config() | |
| return config | |
| def from_config(cls, config): | |
| return cls(**config) | |
| _CTCLayer = CTCLayer | |
| _tensorflow_imported = True | |
| logger.info("TensorFlow loaded (lazy import)") | |
| except ImportError as e: | |
| logger.error(f"Failed to import TensorFlow: {e}") | |
| raise | |
| return _tf, _keras, _CTCLayer | |
| # Model configuration | |
| BASE_DIR = os.path.dirname(os.path.abspath(__file__)) | |
| MODEL_PATH = os.path.join(BASE_DIR, 'my_model.h5') | |
| IMG_HEIGHT, IMG_WIDTH = 50, 200 | |
| class CaptchaPredictor: | |
| def __init__(self, model_path=MODEL_PATH): | |
| self.model_path = model_path | |
| self.model = None | |
| self.prediction_model = None | |
| self.vocab_info = get_vocab_info() | |
| self._load_model() | |
| def _load_model(self): | |
| """Load the trained model""" | |
| try: | |
| # Lazy import TensorFlow | |
| tf, keras, CTCLayer = _lazy_import_tensorflow() | |
| self.model = keras.models.load_model( | |
| self.model_path, | |
| custom_objects={'CTCLayer': CTCLayer} | |
| ) | |
| self._create_prediction_model() | |
| except Exception as e: | |
| if os.environ.get('DEBUG'): | |
| print(f"Error loading model: {e}") | |
| print("Please ensure the model file exists and is valid.") | |
| logger.error(f"Error loading captcha model: {e}") | |
| sys.exit(1) | |
| def _create_prediction_model(self): | |
| """Create a lighter prediction model (exclude CTC loss layer)""" | |
| try: | |
| tf, keras, _ = _lazy_import_tensorflow() | |
| dense_layer = None | |
| for layer in reversed(self.model.layers): | |
| if 'dense' in layer.name.lower() and 'ctc' not in layer.name.lower(): | |
| dense_layer = layer | |
| break | |
| if dense_layer: | |
| self.prediction_model = keras.models.Model( | |
| inputs=self.model.inputs[0], | |
| outputs=dense_layer.output | |
| ) | |
| else: | |
| self.prediction_model = keras.models.Model( | |
| inputs=self.model.inputs[0], | |
| outputs=self.model.layers[-2].output | |
| ) | |
| except Exception as e: | |
| if os.environ.get('DEBUG'): | |
| print(f"Error creating prediction model: {e}") | |
| logger.error(f"Error creating prediction model: {e}") | |
| sys.exit(1) | |
| def predict(self, image_content): | |
| """Run prediction on a single image""" | |
| try: | |
| if not image_content: | |
| raise ValueError("Image content is empty or None") | |
| preprocessed_image = preprocess_image(image_content) | |
| predictions = self.prediction_model.predict(preprocessed_image, verbose=0) | |
| decoded_texts = decode_predictions(predictions) | |
| return decoded_texts[0] if decoded_texts else "" | |
| except Exception as e: | |
| if os.environ.get('DEBUG'): | |
| print(f"Error during prediction: {e}") | |
| logger.error(f"Error during captcha prediction: {e}") | |
| return "" | |
| def predict(image_content): | |
| """Predict captcha text from image content""" | |
| predictor = CaptchaPredictor() | |
| result = predictor.predict(image_content) | |
| return result | |