printing extracted text18

HTR/hcr.py

@@ -6,37 +6,12 @@ import torch
 import sys
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 from utils import notification_queue, log_print
-
-# Global variables for model and processor
-processor = None
-model = None
-
-def initialize_model():
-    """Initialize the TrOCR model and processor"""
-    global processor, model
-    MODEL_NAME = "microsoft/trocr-large-handwritten"
-    try:
-        log_print("Initializing TrOCR model...")
-        processor = TrOCRProcessor.from_pretrained(MODEL_NAME)
-        model = VisionEncoderDecoderModel.from_pretrained(MODEL_NAME)
-        if torch.cuda.is_available():
-            model = model.to('cuda')
-            log_print("Model moved to CUDA")
-        log_print("TrOCR model initialized successfully")
-    except Exception as e:
-        error_msg = str(e)
-        log_print(f"Error initializing TrOCR model: {error_msg}", "ERROR")
-        raise
+from all_models import models
 
 def text(image_cv):
     try:
-        # Initialize model if not already initialized
-        if processor is None or model is None:
-            log_print("TrOCR model not initialized, initializing now...")
-            initialize_model()
-            
-        if processor is None or model is None:
-            raise RuntimeError("Failed to initialize TrOCR model")
+        # Get model instance from singleton
+        model, processor = models.get_trocr_model()
             
         if not isinstance(image_cv, list):
             image_cv = [image_cv]
@@ -61,7 +36,7 @@ def text(image_cv):
                 # Get pixel values
                 pixel_values = processor(image, return_tensors="pt").pixel_values
                 if torch.cuda.is_available():
-                    pixel_values = pixel_values.to('cuda')
+                    pixel_values = pixel_values.to(models.device)
                     
                 # Generate text
                 generated_ids = model.generate(pixel_values)
@@ -93,4 +68,7 @@ def text(image_cv):
             "message": error_msg
         })
         return ""
+    finally:
+        # Release model reference
+        models.release_trocr_model()
     

HTR/strike.py

@@ -6,6 +6,9 @@ import cv2
 from transformers import AutoModelForImageClassification, AutoConfig
 import logging
 from pathlib import Path
+import sys
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from all_models import models
 
 logging.basicConfig(
     level=logging.INFO,
@@ -13,75 +16,12 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# Global variables
-model = None
-TEMP_IMAGES_DIR = None
-
-# Initialize model at module load time
-try:
-    logger.info("Initializing model...")
-    
-    # Get the absolute path to the model
-    current_dir = os.path.dirname(os.path.abspath(__file__))
-    project_root = os.path.dirname(current_dir)  # Changed to one level up
-    model_path = os.path.join(project_root, "models", "vit-base-beans")
-    
-    # Check if model path exists and has proper permissions
-    if not os.path.exists(model_path):
-        logger.error(f"Model path does not exist: {model_path}")
-        raise FileNotFoundError(f"Model path does not exist: {model_path}")
-    
-    # Check if we have read permissions
-    if not os.access(model_path, os.R_OK):
-        logger.error(f"No read permission for model path: {model_path}")
-        raise PermissionError(f"No read permission for model path: {model_path}")
-    
-    # Check for required model files
-    required_files = ['config.json', 'model.safetensors']  # Updated for safetensors
-    for file in required_files:
-        file_path = os.path.join(model_path, file)
-        if not os.path.exists(file_path):
-            logger.error(f"Required model file missing: {file}")
-            raise FileNotFoundError(f"Required model file missing: {file}")
-        if not os.access(file_path, os.R_OK):
-            logger.error(f"No read permission for model file: {file}")
-            raise PermissionError(f"No read permission for model file: {file}")
-    
-    logger.info(f"Loading model from: {model_path}")
-    
-    # Load model from local path with safetensors support
-    config = AutoConfig.from_pretrained(model_path)
-    model = AutoModelForImageClassification.from_pretrained(
-        model_path,
-        local_files_only=True,
-        use_safetensors=True
-    )
-    
-    if torch.cuda.is_available():
-        model = model.to('cuda')
-        logger.info("Model moved to CUDA")
-    else:
-        logger.info("Running on CPU")
-        
-    model.eval()  # Set to evaluation mode
-    logger.info("Model initialized successfully")
-except Exception as e:
-    logger.error(f"Error initializing model: {str(e)}")
-    model = None
-
 def image_preprocessing(image):
     try:
         images = []
         for i in image:
-            # print(i)
-            # img = cv2.imread(i)
             binary_image = i
-            # converting into grayscale
-            # gray_image = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
-            # convrting into binaryimage
-            # _, binary_image = cv2.threshold(gray_image, 200, 255, cv2.THRESH_BINARY)
             binary_image = cv2.resize(binary_image, (224, 224))
-            # binary_image = np.expand_dims(binary_image, axis=-1)
             binary_image = cv2.merge([binary_image, binary_image, binary_image])
             binary_image = binary_image/255
             binary_image = torch.from_numpy(binary_image)
@@ -92,8 +32,11 @@ def image_preprocessing(image):
         logger.error(f"Error in image_preprocessing: {str(e)}")
         return None
 
-def predict_image(images, model):
+def predict_image(images):
     try:
+        # Get model instance from singleton
+        model, processor = models.get_vit_model()
+        
         preprocessed_img = image_preprocessing(images)
         if preprocessed_img is None:
             logger.error("Image preprocessing failed")
@@ -113,19 +56,13 @@ def predict_image(images, model):
     except Exception as e:
         logger.error(f"Error in predict_image: {str(e)}")
         return None
-
-def process_without_model(images):
-    """Fallback function when model prediction fails"""
-    logger.warning("Processing without model - returning all images as not struck")
-    return images  # Return all images as not struck
+    finally:
+        # Release model reference
+        models.release_vit_model()
 
 def struck_images(word_images):
     try:
-        if model is None:
-            logger.warning("Model not initialized, processing without model")
-            return word_images
-            
-        predictions = predict_image(word_images, model)
+        predictions = predict_image(word_images)
         if predictions is None:
             logger.warning("Predictions failed, processing without model")
             return word_images

all_models.py

@@ -36,10 +36,16 @@ class ModelSingleton:
                 self.similarity_model = None
                 self.flan_tokenizer = None
                 self.flan_model = None
+                self.trocr_processor = None
+                self.trocr_model = None
+                self.vit_model = None
+                self.vit_processor = None
                 
                 # Initialize reference counts
                 self._reference_counts['similarity'] = 0
                 self._reference_counts['flan'] = 0
+                self._reference_counts['trocr'] = 0
+                self._reference_counts['vit'] = 0
                 
                 self._initialized = True
                 logger.info("Model singleton initialized")
@@ -96,6 +102,93 @@ class ModelSingleton:
             logger.error(f"Error loading Flan-T5 model: {e}")
             raise
 
+    def get_trocr_model(self):
+        """Get TrOCR model with reference counting"""
+        try:
+            if self.trocr_model is None:
+                from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+                logger.info("Loading TrOCR model...")
+                MODEL_NAME = "microsoft/trocr-large-handwritten"
+                self.trocr_processor = TrOCRProcessor.from_pretrained(MODEL_NAME)
+                self.trocr_model = VisionEncoderDecoderModel.from_pretrained(MODEL_NAME)
+                self.trocr_model.to(self.device)
+                logger.info("TrOCR model loaded successfully")
+            
+            self._reference_counts['trocr'] += 1
+            return self.trocr_model, self.trocr_processor
+        except Exception as e:
+            logger.error(f"Error loading TrOCR model: {e}")
+            raise
+
+    def get_vit_model(self):
+        """Get ViT model with reference counting"""
+        try:
+            if self.vit_model is None:
+                from transformers import ViTImageProcessor, ViTModel, AutoModelForImageClassification, AutoConfig
+                logger.info("Loading ViT model...")
+                
+                # Get model path - fix to use project root
+                import os
+                project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+                model_path = os.path.join(project_root, 'models', 'vit-base-beans')
+                
+                logger.info(f"Looking for model at: {model_path}")
+                
+                if not os.path.exists(model_path):
+                    raise FileNotFoundError(f"Model path does not exist: {model_path}")
+                    
+                # Check for model files
+                model_files = os.listdir(model_path)
+                logger.info(f"Found model files: {model_files}")
+                
+                if 'model.safetensors' not in model_files or 'config.json' not in model_files:
+                    raise FileNotFoundError(f"Required model files missing in {model_path}")
+                
+                # Create processor with explicit settings
+                self.vit_processor = ViTImageProcessor(
+                    do_resize=True,
+                    size=224,
+                    do_normalize=True,
+                    image_mean=[0.5, 0.5, 0.5],
+                    image_std=[0.5, 0.5, 0.5]
+                )
+                
+                try:
+                    # Load model with explicit settings
+                    logger.info("Loading ViT model with safetensors...")
+                    self.vit_model = ViTModel.from_pretrained(
+                        model_path,
+                        local_files_only=True,
+                        use_safetensors=True,
+                        trust_remote_code=False
+                    )
+                    self.vit_model.to(self.device)
+                    self.vit_model.eval()
+                    logger.info("ViT model loaded successfully")
+                except Exception as model_error:
+                    logger.error(f"Error loading model: {model_error}")
+                    # Try alternative loading method
+                    try:
+                        logger.info("Attempting alternative model loading...")
+                        self.vit_model = AutoModelForImageClassification.from_pretrained(
+                            model_path,
+                            local_files_only=True,
+                            use_safetensors=True,
+                            trust_remote_code=False
+                        )
+                        self.vit_model.to(self.device)
+                        self.vit_model.eval()
+                        logger.info("ViT model loaded successfully using alternative method")
+                    except Exception as alt_error:
+                        logger.error(f"Alternative loading also failed: {alt_error}")
+                        raise
+            
+            self._reference_counts['vit'] += 1
+            return self.vit_model, self.vit_processor
+        except Exception as e:
+            logger.error(f"Error loading ViT model: {e}")
+            raise
+
     def release_similarity_model(self):
         """Release reference to similarity model"""
         self._reference_counts['similarity'] -= 1
@@ -108,6 +201,18 @@ class ModelSingleton:
         if self._reference_counts['flan'] <= 0:
             self._cleanup_flan_model()
 
+    def release_trocr_model(self):
+        """Release reference to TrOCR model"""
+        self._reference_counts['trocr'] -= 1
+        if self._reference_counts['trocr'] <= 0:
+            self._cleanup_trocr_model()
+
+    def release_vit_model(self):
+        """Release reference to ViT model"""
+        self._reference_counts['vit'] -= 1
+        if self._reference_counts['vit'] <= 0:
+            self._cleanup_vit_model()
+
     def _cleanup_similarity_model(self):
         """Clean up similarity model resources"""
         if self.similarity_model is not None:
@@ -126,16 +231,53 @@ class ModelSingleton:
             torch.cuda.empty_cache()
             logger.info("Flan-T5 model resources cleaned up")
 
+    def _cleanup_trocr_model(self):
+        """Clean up TrOCR model resources"""
+        if self.trocr_model is not None:
+            del self.trocr_model
+            del self.trocr_processor
+            self.trocr_model = None
+            self.trocr_processor = None
+            torch.cuda.empty_cache()
+            logger.info("TrOCR model resources cleaned up")
+
+    def _cleanup_vit_model(self):
+        """Clean up ViT model resources"""
+        if self.vit_model is not None:
+            del self.vit_model
+            del self.vit_processor
+            self.vit_model = None
+            self.vit_processor = None
+            torch.cuda.empty_cache()
+            logger.info("ViT model resources cleaned up")
+
     def cleanup(self):
         """Clean up all model resources"""
         try:
-            self._cleanup_similarity_model()
-            self._cleanup_flan_model()
-            self._reference_counts['similarity'] = 0
-            self._reference_counts['flan'] = 0
-            logger.info("All model resources cleaned up successfully")
+            logger.info("Starting model cleanup...")
+            
+            # Clean up each model type
+            if self._reference_counts.get('similarity', 0) > 0:
+                self._cleanup_similarity_model()
+            if self._reference_counts.get('flan', 0) > 0:
+                self._cleanup_flan_model()
+            if self._reference_counts.get('trocr', 0) > 0:
+                self._cleanup_trocr_model()
+            if self._reference_counts.get('vit', 0) > 0:
+                self._cleanup_vit_model()
+            
+            # Reset reference counts
+            for model_type in self._reference_counts:
+                self._reference_counts[model_type] = 0
+            
+            # Force CUDA cache cleanup
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            
+            logger.info("Model cleanup completed successfully")
         except Exception as e:
-            logger.error(f"Error during cleanup: {e}")
+            logger.error(f"Error during model cleanup: {e}")
+            # Continue cleanup even if there's an error
 
 # Create global instance
 models = ModelSingleton()

main.py

@@ -5,6 +5,8 @@ import psutil
 import time
 import logging
 import queue
+import torch
+from all_models import ModelSingleton
 
 # Set up logging first
 logging.basicConfig(
@@ -74,7 +76,6 @@ from dotenv import load_dotenv
 warnings.filterwarnings('ignore')
 
 # Import ML libraries
-import torch
 import nltk
 import gensim
 from gensim.models import FastText
@@ -137,6 +138,9 @@ log_file = os.path.join(log_dir, 'app.log')  # Add log file path
 global_models = {}
 initialization_complete = Event()
 
+# Initialize model singleton
+models = ModelSingleton()
+
 def ensure_directory(path):
     """Create directory and ensure full permissions with better error handling"""
     if os.path.exists(path):
@@ -178,14 +182,20 @@ def get_or_load_model(model_name):
                 from gensim.models import KeyedVectors
                 log_print(f"Loading {model_name} model...")
                 model_path = os.path.join(gensim_data_dir, 'fasttext-wiki-news-subwords-300', 'fasttext-wiki-news-subwords-300.gz')
+                model_dir = os.path.dirname(model_path)
+                
                 try:
-                    if not os.path.exists(model_path):
+                    # Create model directory if it doesn't exist
+                    os.makedirs(model_dir, exist_ok=True)
+                    
+                    if os.path.exists(model_path):
+                        log_print("Loading fasttext model from cache...")
+                        model = KeyedVectors.load_word2vec_format(model_path)
+                    else:
+                        # Only download if file doesn't exist
                         from gensim.downloader import load
                         log_print("Downloading fasttext model...")
                         model = load('fasttext-wiki-news-subwords-300')
-                    else:
-                        log_print("Loading fasttext model from cache...")
-                        model = KeyedVectors.load_word2vec_format(model_path)
                     
                     # Move model to CPU explicitly if it has the 'to' method
                     if hasattr(model, 'to'):
@@ -523,8 +533,8 @@ def cleanup_memory():
         # Clear Python garbage collection
         gc.collect()
         
-        # Clear model caches
-        if hasattr(models, 'cleanup'):
+        # Clean up models
+        if models:
             models.cleanup()
             
         # Log memory usage

similarity_check/semantic_meaning_check/semantic.py

@@ -26,8 +26,13 @@ os.environ['GENSIM_DATA_DIR'] = gensim_data_dir
 
 # Load fasttext with error handling
 try:
-    print("Loading fasttext model...")
-    fasttext = load('fasttext-wiki-news-subwords-300')
+    model_path = os.path.join(gensim_data_dir, 'fasttext-wiki-news-subwords-300', 'fasttext-wiki-news-subwords-300.gz')
+    if os.path.exists(model_path):
+        print("Loading fasttext model from cache...")
+        fasttext = KeyedVectors.load_word2vec_format(model_path)
+    else:
+        print("Loading fasttext model...")
+        fasttext = load('fasttext-wiki-news-subwords-300')
 except Exception as e:
     print(f"Error loading fasttext model: {e}")
     # Provide a fallback for similarity calculations