printing extracted text17

HTR/app.py

@@ -89,16 +89,13 @@ def extract_text_from_image(img_path):
         imgs = convert_image(img)
         if not imgs:
             log_print("No text regions detected, processing whole image...", "WARNING")
-            # Try processing the whole image as one region
-            temp_path = os.path.join(tempfile.gettempdir(), 'whole_image.png')
-            cv2.imwrite(temp_path, img)
-            imgs = [temp_path]
+            imgs = [img]
             
         log_print(f"Found {len(imgs)} text regions")
         
         log_print("Processing text regions...")
-        images_path = struck_images(imgs)
-        if not images_path:
+        processed_images = struck_images(imgs)
+        if not processed_images:
             error_msg = "No valid text regions after processing"
             log_print(error_msg, "ERROR")
             notification_queue.put({
@@ -108,7 +105,7 @@ def extract_text_from_image(img_path):
             return ""
             
         log_print("Extracting text from regions...")
-        t = text(images_path)
+        t = text(processed_images)
         if not t:
             error_msg = "No text could be extracted from image"
             log_print(error_msg, "ERROR")

HTR/strike.py

@@ -23,7 +23,7 @@ try:
     
     # 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))
+    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
@@ -37,7 +37,7 @@ try:
         raise PermissionError(f"No read permission for model path: {model_path}")
     
     # Check for required model files
-    required_files = ['config.json', 'pytorch_model.bin']
+    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):
@@ -49,9 +49,13 @@ try:
     
     logger.info(f"Loading model from: {model_path}")
     
-    # Load model from local path
+    # Load model from local path with safetensors support
     config = AutoConfig.from_pretrained(model_path)
-    model = AutoModelForImageClassification.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')
@@ -88,42 +92,55 @@ def image_preprocessing(image):
         logger.error(f"Error in image_preprocessing: {str(e)}")
         return None
 
-def predict_image(image_paths, model):
+def predict_image(images, 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()
+        preprocessed_img = image_preprocessing(images)
+        if preprocessed_img is None:
+            logger.error("Image preprocessing failed")
+            return None
+            
+        images_tensor = torch.stack(preprocessed_img)
+        images_tensor = images_tensor.permute(0, 3, 1, 2)
+        
+        with torch.no_grad():
+            predictions = model(images_tensor).logits
+            if torch.cuda.is_available():
+                predictions = predictions.cpu()
+            predictions = predictions.numpy()
+            
         return predictions
         
     except Exception as e:
         logger.error(f"Error in predict_image: {str(e)}")
-        return process_without_model(image_paths)
+        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
 
-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] 
+def struck_images(word_images):
     try:
-        predictions = predict_image(word__image, model)
-
-        not_struck =[]
+        if model is None:
+            logger.warning("Model not initialized, processing without model")
+            return word_images
+            
+        predictions = predict_image(word_images, model)
+        if predictions is None:
+            logger.warning("Predictions failed, processing without model")
+            return word_images
+            
+        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)
+            if predictions[i].argmax() == 0:  # Assuming 0 is the "not struck" class
+                not_struck.append(word_images[i])
+                
+        if not not_struck:
+            logger.warning("No non-struck images found, returning all images")
+            return word_images
+            
         return not_struck
+        
     except Exception as e:
-        logger.error(f"Error in process_without_model: {str(e)}")
-        return None  # Return all as not struck
+        logger.error(f"Error in struck_images: {str(e)}")
+        return word_images  # Return all images on error