# 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__) # Initialize model globally model = None TEMP_IMAGES_DIR = None def initialize_model(): global model if model is None: try: logger.info("Initializing model...") # Use model directly from Hugging Face hub model_name = "microsoft/resnet-50" # Using a more general model for classification try: # First try to load from cache cache_dir = os.path.join(os.environ.get('TMPDIR', '/tmp'), 'model_cache') os.makedirs(cache_dir, exist_ok=True) config = AutoConfig.from_pretrained(model_name, cache_dir=cache_dir) model = AutoModelForImageClassification.from_pretrained( model_name, config=config, cache_dir=cache_dir ) logger.info(f"Model loaded from {model_name}") except Exception as e: logger.error(f"Error loading model from hub: {str(e)}") # Fallback to simpler processing if model fails to load return None 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") return model except Exception as e: logger.error(f"Error initializing model: {str(e)}") return None def image_preprocessing(image): try: images = [] for i in image: try: # Ensure image is in correct format if isinstance(i, str): # If i is a path, read the image i = cv2.imread(i, cv2.IMREAD_GRAYSCALE) if i is None: logger.error("Failed to read image from path") continue # Resize to model input size binary_image = cv2.resize(i, (224, 224)) # Convert to RGB (3 channels) binary_image = cv2.cvtColor(binary_image, cv2.COLOR_GRAY2RGB) # Normalize binary_image = binary_image.astype(np.float32) / 255.0 # Convert to tensor binary_image = torch.from_numpy(binary_image) binary_image = binary_image.permute(2, 0, 1) # Change to CxHxW format images.append(binary_image) except Exception as e: logger.error(f"Error preprocessing individual image: {str(e)}") continue if not images: logger.error("No images were successfully preprocessed") return None return images except Exception as e: logger.error(f"Error in image_preprocessing: {str(e)}") return None def predict_image(image_paths, model): try: if model is None: logger.warning("Model not initialized, using basic processing") return process_without_model(image_paths) preprocessed_imgs = image_preprocessing(image_paths) if not preprocessed_imgs: logger.warning("No preprocessed images, using basic processing") return process_without_model(image_paths) images = torch.stack(preprocessed_imgs) if torch.cuda.is_available(): images = images.to('cuda') with torch.no_grad(): predictions = model(images).logits.detach().cpu().numpy() return predictions except Exception as e: logger.error(f"Error in predict_image: {str(e)}") return process_without_model(image_paths) def process_without_model(image_paths): """Fallback processing when model is not available""" try: results = [] for path in image_paths: # Basic image processing to detect if image is struck through img = cv2.imread(path, cv2.IMREAD_GRAYSCALE) if img is None: continue # Use basic image processing to detect strike-through # This is a simplified approach thresh = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)[1] horizontal_lines = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, np.ones((1, 20), np.uint8)) # If there are significant horizontal lines, consider it struck if np.sum(horizontal_lines) > (img.shape[0] * img.shape[1] * 0.1): results.append(1) # Struck else: results.append(0) # Not struck return np.array(results) except Exception as e: logger.error(f"Error in process_without_model: {str(e)}") return np.zeros(len(image_paths)) # Return all as not struck def process_single_image(img): try: # Convert to grayscale if needed if len(img.shape) == 3: img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # Enhance contrast img = cv2.equalizeHist(img) # Denoise img = cv2.fastNlMeansDenoising(img) # Apply adaptive thresholding binary = cv2.adaptiveThreshold( img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 21, 15 ) # Remove noise and smooth edges kernel = np.ones((3,3), np.uint8) binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel) binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel) return binary except Exception as e: logger.error(f"Error in process_single_image: {str(e)}") return None def check_strike_through(img): """Check if an image contains strike-through lines""" try: # Convert to binary _, binary = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) # Look for horizontal lines horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (25, 1)) horizontal_lines = cv2.morphologyEx(binary, cv2.MORPH_OPEN, horizontal_kernel) # Count pixels in horizontal lines line_pixels = np.sum(horizontal_lines == 255) total_pixels = img.shape[0] * img.shape[1] # If more than 5% of pixels are part of horizontal lines, consider it struck through return (line_pixels / total_pixels) > 0.05 except Exception as e: logger.error(f"Error checking strike-through: {str(e)}") return False def struck_images(image_paths): """Process images and detect which ones are not struck through""" try: if not image_paths: logger.error("No image paths provided") return [] logger.info(f"Processing {len(image_paths)} images") not_struck = [] for i, img_path in enumerate(image_paths): try: # Read the image img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) if img is None: logger.error(f"Failed to read image: {img_path}") continue # Process the image processed = process_single_image(img) if processed is None: continue # Check if image is struck through if not check_strike_through(processed): not_struck.append(img_path) except Exception as e: logger.error(f"Error processing image {img_path}: {str(e)}") continue logger.info(f"Found {len(not_struck)} non-struck images") return not_struck except Exception as e: logger.error(f"Error in struck_images: {str(e)}") return [] # struck_images()