Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 import gradio as gr
-from PIL import Image, ImageTk
+from PIL import Image
 import pandas as pd
 from ultralytics import YOLO
 import easyocr
@@ -9,22 +9,27 @@ from tqdm import tqdm  # For progress tracking
 from datetime import datetime  # Import for handling dates
 
 # Initialize variables
-output_excel = "results.xlsx"
 annotations = []
 current_index = 0
+output_excel = "results.xlsx"
+
 # Directories for saving original and YOLO-ed images
 original_dir = "original_images"
 yoloed_dir = "yoloed_images"
+
 # Ensure directories exist
 os.makedirs(original_dir, exist_ok=True)
 os.makedirs(yoloed_dir, exist_ok=True)
+
 # Load the trained YOLO model
 model_path = "best.pt"
 if not os.path.exists(model_path):
     raise FileNotFoundError(f"YOLO model not found at {model_path}. Ensure the model is in the correct location.")
 model = YOLO(model_path)
+
 # Initialize EasyOCR reader
 reader = easyocr.Reader(['en'])
+
 # Global counter for cropped image filenames
 global_crop_counter = 0
 
@@ -79,68 +84,18 @@ def perform_ocr(image):
     numbers = ''.join(filter(str.isdigit, ''.join(result)))
     return numbers.strip()
 
-# Function to load images and process them one by one
-def load_images_with_progress(folder_path):
-    global annotations, current_index
-    try:
-        # Clear previous annotations
-        annotations.clear()
-        current_index = 0
-        # Get list of images
-        image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
-        total_images = len(image_files)
-        # Process images one by one
-        for i, image_file in enumerate(tqdm(image_files, desc="Processing Images")):
-            image_path = os.path.join(folder_path, image_file)
-            # Save the original image
-            original_filename = os.path.basename(image_path)
-            original_path = os.path.join(original_dir, original_filename)
-            original_image = Image.open(image_path)
-            original_image.save(original_path)
-            # Detect objects using YOLO
-            detections = detect_objects(image_path)
-            if detections:
-                # Crop the image using YOLO bounding boxes
-                cropped_images = crop_image(image_path, detections)
-                # Perform OCR on each cropped image
-                for j, cropped_image in enumerate(cropped_images):
-                    annotation = {
-                        "image_path": image_path,
-                        "cropped_image": cropped_image,
-                        "meter_value": perform_ocr(cropped_image),
-                        "original_image": original_image,
-                        "room_number": ""  # Initialize room number as empty
-                    }
-                    annotations.append(annotation)
-            else:
-                # If no detections, add the original image with no OCR result
-                annotation = {
-                    "image_path": image_path,
-                    "cropped_image": None,
-                    "meter_value": "",
-                    "original_image": original_image,
-                    "room_number": ""  # Initialize room number as empty
-                }
-                annotations.append(annotation)
-        # Update GUI
-        if annotations:
-            current_index = 0
-            return update_gui()
-    except Exception as e:
-        return str(e), None, None, None, None
-
 # Function to update GUI
 def update_gui():
     global current_index
     if not annotations:
-        return "No images processed.", None, None, None, None
+        return "No images processed.", None, "", ""
     annotation = annotations[current_index]
     # Return images and annotations
     original_image = annotation["original_image"]
     cropped_image = annotation["cropped_image"]
     meter_value = annotation.get("meter_value", "")
     room_number = annotation.get("room_number", "")
-    return original_image, cropped_image, meter_value, room_number, ""
+    return original_image, cropped_image, meter_value, room_number
 
 # Function to move to the next image
 def next_image():
@@ -148,7 +103,7 @@ def next_image():
     if current_index < len(annotations) - 1:
         current_index += 1
         return update_gui()
-    return "No more images.", None, None, None, None
+    return "No more images.", None, "", ""
 
 # Function to move to the previous image
 def prev_image():
@@ -156,14 +111,14 @@ def prev_image():
     if current_index > 0:
         current_index -= 1
         return update_gui()
-    return "No previous images.", None, None, None, None
+    return "No previous images.", None, "", ""
 
 # Function to export to Excel
-def export_to_excel(room_numbers, meter_values):
+def export_to_excel(room_numbers_state, meter_values_state):
     try:
         # Prepare data for export
         data = []
-        for room_number, meter_value in zip(room_numbers, meter_values):
+        for room_number, meter_value in zip(room_numbers_state, meter_values_state):
             data.append({
                 "room_number": room_number,
                 "meter_value": meter_value
@@ -175,18 +130,19 @@ def export_to_excel(room_numbers, meter_values):
         # Create a folder with today's date
         today_date = datetime.now().strftime("%Y-%m-%d")  # Format: YYYY-MM-DD
         folder_name = f"electricity_meter_values_{today_date}"
-        os.makedirs(folder_name, exist_ok=True)  # Create the folder if it doesn't exist
+        os.makedirs(folder_name, exist_ok=True)
+
         # Save original images renamed to room numbers
-        for annotation, room_number in zip(annotations, room_numbers):
+        for annotation, room_number in zip(annotations, room_numbers_state):
             if room_number:  # Only proceed if room number is not empty
                 original_image_path = annotation["image_path"]
                 original_image = Image.open(original_image_path)
-                # Construct the new filename using the room number
-                new_image_name = f"{room_number}.jpg"  # Use room number as the filename
+                new_image_name = f"{room_number}.jpg"
                 new_image_path = os.path.join(folder_name, new_image_name)
-                # Save the image in the new folder
                 original_image.save(new_image_path)
+
         return f"Excel file and images saved successfully in '{folder_name}' folder."
+
     except Exception as e:
         return str(e)
 
@@ -202,33 +158,21 @@ def update_meter_value(meter_value):
         annotations[current_index]["meter_value"] = meter_value
     return meter_value
 
+# Function to process uploaded images
 def process_images(folder_input):
     global annotations, current_index
-
     if folder_input is None:
-        # Return default values if no folder is selected
-        return (
-            None,  # original_image_output
-            None,  # cropped_image_output
-            None,  # meter_value_output
-            None,  # room_number_output
-            [],    # room_numbers (state)
-            []     # meter_values (state)
-        )
-
+        return None, None, "", "", [], []
     try:
-        folder_path = folder_input['name']  # Extract path from folder input
+        folder_path = folder_input['name']  # Get actual path
         annotations.clear()
         current_index = 0
-
         image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
         total_images = len(image_files)
-
         for i, image_file in enumerate(tqdm(image_files, desc="Processing Images")):
             image_path = os.path.join(folder_path, image_file)
             original_image = Image.open(image_path)
             detections = detect_objects(image_path)
-
             if detections:
                 cropped_images = crop_image(image_path, detections)
                 for cropped_image in cropped_images:
@@ -249,26 +193,83 @@ def process_images(folder_input):
                     "room_number": ""
                 }
                 annotations.append(annotation)
-
         if annotations:
             current_index = 0
             result = update_gui()
             room_numbers = [a["room_number"] for a in annotations]
             meter_values = [a["meter_value"] for a in annotations]
             return (*result, room_numbers, meter_values)
+        return None, None, "", "", [], []
+    except Exception as e:
+        return str(e), None, "", "", [], []
 
-        # Default return in case no annotations were created
-        return (None, None, "", "", [], [])
+# Define Gradio Interface
+with gr.Blocks() as demo:
+    gr.Markdown("## Electricity Meter Reader")
 
-    except Exception as e:
-        # Ensure all 6 outputs are returned even on error
-        return (
-            f"Error: {str(e)}",  # original_image_output (textual error display)
-            None,               # cropped_image_output
-            None,               # meter_value_output
-            None,               # room_number_output
-            [],                 # room_numbers (state)
-            []                  # meter_values (state)
-        )
+    with gr.Row():
+        folder_input = gr.File(label="Upload Folder", file_types=["directory"])
+
+    with gr.Row():
+        original_image_output = gr.Image(label="Original Image")
+        cropped_image_output = gr.Image(label="Cropped Meter Region")
+
+    with gr.Row():
+        meter_value_output = gr.Textbox(label="Meter Value", interactive=True)
+        room_number_output = gr.Textbox(label="Room Number", interactive=True)
+
+    with gr.Row():
+        prev_button = gr.Button("Previous")
+        next_button = gr.Button("Next")
+        export_button = gr.Button("Export to Excel")
+
+    # Hidden states for storing lists of room numbers and meter values
+    room_numbers_state = gr.State([])
+    meter_values_state = gr.State([])
+
+    # Actions
+    folder_input.change(
+        fn=process_images,
+        inputs=[folder_input],
+        outputs=[
+            original_image_output,
+            cropped_image_output,
+            meter_value_output,
+            room_number_output,
+            room_numbers_state,
+            meter_values_state,
+        ]
+    )
+
+    prev_button.click(
+        fn=prev_image,
+        inputs=[],
+        outputs=[original_image_output, cropped_image_output, meter_value_output, room_number_output]
+    )
+
+    next_button.click(
+        fn=next_image,
+        inputs=[],
+        outputs=[original_image_output, cropped_image_output, meter_value_output, room_number_output]
+    )
+
+    room_number_output.change(
+        fn=update_room_number,
+        inputs=[room_number_output],
+        outputs=[room_number_output]
+    )
+
+    meter_value_output.change(
+        fn=update_meter_value,
+        inputs=[meter_value_output],
+        outputs=[meter_value_output]
+    )
+
+    export_button.click(
+        fn=export_to_excel,
+        inputs=[room_numbers_state, meter_values_state],
+        outputs=[gr.Textbox(label="Status")]
+    )
 
+# Launch the app
 demo.launch()