HTR/strike.py

# import pandas as pd
import numpy as np
import torch
import os
import cv2
from transformers import AutoModelForImageClassification, AutoConfig
import logging
from pathlib import Path

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
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(os.path.dirname(current_dir))
    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', 'pytorch_model.bin']
    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
    config = AutoConfig.from_pretrained(model_path)
    model = AutoModelForImageClassification.from_pretrained(model_path)
    
    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)
            images.append(binary_image)
        return images
        
    except Exception as e:
        logger.error(f"Error in image_preprocessing: {str(e)}")
        return None

def predict_image(image_paths, model):
    try:
        preprocessed_img = image_preprocessing(image_path)
        images = torch.stack(preprocessed_img)
        images = images.permute(0, 3, 1, 2)
        predictions = model(images).logits.detach().numpy()
        return predictions
        
    except Exception as e:
        logger.error(f"Error in predict_image: {str(e)}")
        return process_without_model(image_paths)


def struck_images(word__image):
    # folder_path = 'images'
    # images_path = []
    # for filename in os.listdir(folder_path):
    #     file_path = os.path.join(folder_path, filename)
    #     images_path.append(file_path)
    # # print()
    # images_path.sort(key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
  
    # # print(images_path)  
    
    # # images_path = images_path[:2] 
    try:
        predictions = predict_image(word__image, model)

        not_struck =[]
        for i in range(len(predictions)):
            if predictions[i].argmax().item() == 0:
                # not_struck.append(images_path[i])
                not_struck.append(word__image[i])

        # print(not_struck)
        return not_struck
    except Exception as e:
        logger.error(f"Error in process_without_model: {str(e)}")
        return None  # Return all as not struck