Update utility/utils.py

utility/utils.py

@@ -32,15 +32,16 @@ def draw_boxes(image, bounds, color='red', width=2):
         draw.line([*p0, *p1, *p2, *p3, *p0], fill=color, width=width)
     return image
 
-#Image Quality upscaling 
 # Load image using OpenCV
 def load_image(image_path):
-    return cv2.imread(image_path)
+    image = cv2.imread(image_path)
+    if image is None:
+        raise ValueError(f"Could not load image from {image_path}. It may be corrupted or the path is incorrect.")
+    return image
 
-# Function for upscaling image using OpenCV's INTER_CUBIC or ESRGAN (if available)
+# Function for upscaling image using OpenCV's INTER_CUBIC
 def upscale_image(image, scale=2):
     height, width = image.shape[:2]
-    # Simple upscaling using cubic interpolation
     upscaled_image = cv2.resize(image, (width * scale, height * scale), interpolation=cv2.INTER_CUBIC)
     return upscaled_image
 
@@ -58,11 +59,9 @@ def sharpen_image(image):
 
 # Function to increase contrast and enhance details without changing color
 def enhance_image(image):
-    # Convert from BGR to RGB for PIL processing, then back to BGR
     pil_img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
     enhancer = ImageEnhance.Contrast(pil_img)
     enhanced_image = enhancer.enhance(1.5)
-    # Convert back to BGR
     enhanced_image_bgr = cv2.cvtColor(np.array(enhanced_image), cv2.COLOR_RGB2BGR)
     return enhanced_image_bgr
 
@@ -70,59 +69,58 @@ def enhance_image(image):
 def process_image(image_path, scale=2):
     # Load the image
     image = load_image(image_path)
-    
+
     # Upscale the image
     upscaled_image = upscale_image(image, scale)
-    
+
     # Reduce noise
     denoised_image = reduce_noise(upscaled_image)
-    
+
     # Sharpen the image
     sharpened_image = sharpen_image(denoised_image)
-    
+
     # Enhance the image contrast and details without changing color
     final_image = enhance_image(sharpened_image)
-    
-    return final_image
 
+    return final_image
 
 def ocr_with_paddle(img):
     finaltext = ''
     model_dir = os.getenv('PADDLEOCR_MODEL_DIR', '/tmp/.paddleocr')
     ocr = PaddleOCR(lang='en', use_angle_cls=True, det_model_dir=model_dir)
-    # img_path = 'exp.jpeg'
     result = ocr.ocr(img)
-    
+
     for i in range(len(result[0])):
         text = result[0][i][1][0]
-        finaltext += ' '+ text
+        finaltext += ' ' + text
     return finaltext
 
 def extract_text_from_images(image_paths, RESULT_FOLDER):
     all_extracted_texts = {}
-    all_extracted_imgs={}
+    all_extracted_imgs = {}
     for image_path in image_paths:
-        # Enhance the image before OCR        
-        enhanced_image = process_image(image_path, scale=2)
-        #bounds = reader.readtext(enhanced_image)
-            # Draw boxes on the processed image
-        img_result = Image.fromarray(enhanced_image)
-        #draw_boxes(img_result, bounds)
-
-        result_image_path = os.path.join(RESULT_FOLDER, f'result_{os.path.basename(image_path)}')
-        img_result.save(result_image_path)  # Save the processed image
-
-        # Perform OCR on the enhanced image
-        result=ocr_with_paddle(enhanced_image)
-        # results = reader.readtext(enhanced_image)
-        # extracted_text = " ".join([res[1] for res in results])
-
-        all_extracted_texts[image_path] =result
-        all_extracted_imgs[image_path] = result_image_path
-     # Convert to JSON-compatible structure
-    all_extracted_imgs_json = {str(k): str(v) for k, v in all_extracted_imgs.items()}
-    return all_extracted_texts,all_extracted_imgs_json
+        try:
+            # Enhance the image before OCR        
+            enhanced_image = process_image(image_path, scale=2)
 
+            # Draw boxes on the processed image (optional, requires bounds)
+            img_result = Image.fromarray(enhanced_image)
+
+            result_image_path = os.path.join(RESULT_FOLDER, f'result_{os.path.basename(image_path)}')
+            img_result.save(result_image_path)  # Save the processed image
+
+            # Perform OCR on the enhanced image
+            result = ocr_with_paddle(enhanced_image)
+
+            all_extracted_texts[image_path] = result
+            all_extracted_imgs[image_path] = result_image_path
+        except ValueError as ve:
+            print(f"Error processing image {image_path}: {ve}")
+            continue  # Continue to the next image if there's an error
+
+    # Convert to JSON-compatible structure
+    all_extracted_imgs_json = {str(k): str(v) for k, v in all_extracted_imgs.items()}
+    return all_extracted_texts, all_extracted_imgs_json
 
 # Function to call the Gemma model and process the output as Json 
 def Data_Extractor(data, client=client):