update_new_new_new_new_new

HTR/strike.py

@@ -3,8 +3,9 @@ import numpy as np
 import torch
 import os
 import cv2
-from transformers import AutoModelForImageClassification
+from transformers import AutoModelForImageClassification, AutoConfig
 import logging
+from pathlib import Path
 
 logging.basicConfig(
     level=logging.INFO,
@@ -21,38 +22,93 @@ def initialize_model():
     if model is None:
         try:
             logger.info("Initializing model...")
-            model = AutoModelForImageClassification.from_pretrained("models/vit-base-beans")
+            # 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)}")
-            raise
+            return None
 
 def image_preprocessing(image):
     try:
         images = []
         for i in image:
-            binary_image = i
-            binary_image = cv2.resize(binary_image, (224, 224))
-            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)
+            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 []
+        return None
 
-def predict_image(image_path, model):
+def predict_image(image_paths, model):
     try:
-        preprocessed_img = image_preprocessing(image_path)
-        if not preprocessed_img:
-            return None
+        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_img)
-        images = images.permute(0, 3, 1, 2)
+        images = torch.stack(preprocessed_imgs)
         
         if torch.cuda.is_available():
             images = images.to('cuda')
@@ -60,9 +116,38 @@ def predict_image(image_path, model):
         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 None
+        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 struck_images(image_paths):
     try:
@@ -73,6 +158,9 @@ def struck_images(image_paths):
         logger.info(f"Processing {len(image_paths)} images")
         processed_paths = []
         
+        # Initialize model
+        model = initialize_model()
+        
         for i, img_path in enumerate(image_paths):
             try:
                 # Read the image from the path
@@ -81,12 +169,6 @@ def struck_images(image_paths):
                     logger.error(f"Failed to read image: {img_path}")
                     continue
                 
-                # Resize if image is too small
-                min_size = 800
-                if img.shape[0] < min_size or img.shape[1] < min_size:
-                    scale = min_size / min(img.shape[0], img.shape[1])
-                    img = cv2.resize(img, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
-                    
                 # Process the image
                 processed = process_single_image(img)
                 if processed is None:
@@ -101,8 +183,21 @@ def struck_images(image_paths):
                 logger.error(f"Error processing image {img_path}: {str(e)}")
                 continue
                 
-        logger.info(f"Successfully processed {len(processed_paths)} images")
-        return processed_paths
+        # Get predictions
+        predictions = predict_image(processed_paths, model)
+        
+        # Filter based on predictions
+        not_struck = []
+        for i, pred in enumerate(predictions):
+            if isinstance(pred, np.ndarray):
+                if pred.argmax() == 0:  # Not struck
+                    not_struck.append(processed_paths[i])
+            else:
+                if pred == 0:  # Not struck
+                    not_struck.append(processed_paths[i])
+                
+        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)}")