Update app.py

app.py

@@ -4,6 +4,7 @@ from PIL import Image, ImageDraw, ImageFont
 import json
 from paddleocr import PaddleOCR
 import gradio as gr
+import os
 
 # Initialize PaddleOCR
 ocr = PaddleOCR(use_angle_cls=True, lang='en')
@@ -27,10 +28,10 @@ def draw_boxes_on_image(image, data):
         bounding_box, (text, confidence) = item
 
         # Ensure bounding_box is a list of lists
-        if isinstance(bounding_box[0], list):
-            box = np.array(bounding_box).astype(int)
-        else:
-            box = np.array([bounding_box]).astype(int)
+        if not isinstance(bounding_box[0], list):
+            bounding_box = [bounding_box]
+
+        box = np.array(bounding_box).astype(int)
 
         # Draw the bounding box
         draw.line([tuple(box[0]), tuple(box[1])], fill="green", width=2)
@@ -44,26 +45,29 @@ def draw_boxes_on_image(image, data):
 
     return pil_image
 
-# Function to save OCR results to JSON
-def save_results_to_json(ocr_results):
-    results = []
-
-    for line in ocr_results:
-        for word_info in line:
-            bounding_box = word_info[0]
-            text, confidence = word_info[1]
-
-            # Ensure bounding_box is a list of lists
-            if not isinstance(bounding_box[0], list):
-                bounding_box = [bounding_box]
-
-            results.append({
-                "bounding_box": [list(map(float, coord)) for coord in bounding_box],
-                "text": text,
-                "confidence": confidence
-            })
-
-    return results
+# Function to convert OCR results to JSON
+def convert_to_json(results, output_file):
+    """
+    Converts the given results into a JSON file.
+
+    Args:
+        results: The list of results containing bounding box coordinates, text, and confidence.
+        output_file: The name of the output JSON file.
+    """
+    json_data = []
+    for result in results:
+        bounding_box = result[0]
+        text = result[1][0]
+        confidence = result[1][1]
+
+        json_data.append({
+            "bounding_box": [list(map(float, coord)) for coord in bounding_box],
+            "text": text,
+            "confidence": confidence
+        })
+
+    with open(output_file, "w") as f:
+        json.dump(json_data, f, indent=4)
 
 # Function to identify 'field', 'value' pairs
 def identify_field_value_pairs(ocr_results, fields):
@@ -86,16 +90,15 @@ def process_image(image):
     processed_image = draw_boxes_on_image(image, ocr_results[0])
     
     # Save OCR results to JSON
-    results_json = save_results_to_json(ocr_results[0])
     json_path = "ocr_results.json"
-    with open(json_path, 'w') as json_file:
-        json.dump(results_json, json_file, indent=4)
+    convert_to_json(ocr_results[0], json_path)
     
     # Identify field-value pairs
     fields = ["Scheme Name", "Folio Number", "Number of Units", "PAN", "Signature", "Tax Status", 
               "Mobile Number", "Email", "Address", "Bank Account Details"]
     field_value_pairs = identify_field_value_pairs(ocr_results[0], fields)
     field_value_json_path = "field_value_pairs.json"
+    
     with open(field_value_json_path, 'w') as json_file:
         json.dump(field_value_pairs, json_file, indent=4)
 
@@ -107,8 +110,8 @@ interface = gr.Interface(
     inputs="image", 
     outputs=[
         "image", 
-        gr.File(label="OCR Results JSON"), 
-        gr.File(label="Field-Value Pairs JSON")
+        gr.File(label="Download OCR Results JSON"), 
+        gr.File(label="Download Field-Value Pairs JSON")
     ],
     title="OCR Web Application",
     description="Upload an image and get OCR results with bounding boxes and two JSON outputs."