Update app.py

app.py

@@ -302,54 +302,51 @@ def process_with_yolo(img_pil):
 
 
 def display_detected_lines(original_path, output_path):
-    
     # Derive the txt_path from the output_path
     txt_path = os.path.join('yolov3/runs/detect/mainlinedection/labels', os.path.basename(output_path).replace(".jpg", ".txt"))
 
     if os.path.exists(txt_path):
         # Load both original and thresholded images
         original_image = Image.open(original_path)
-        thresholded_image = process_image(original_image)  # This is your function that returns a thresholded PIL image
+        thresholded_image = process_image(original_image)
         boxes = get_detected_boxes(txt_path, original_image.width, original_image.height)
-        
+
         if not boxes:
             st.warning("No lines detected by YOLOv3.")
             return
 
         # Create a temporary directory to store the detected lines
         with TemporaryDirectory() as temp_dir:
-            
-            detected_line_paths = []  # List to store paths of the detected line images
+            detected_line_paths = []  # For storing paths of the thresholded line images for OCR
+            original_line_paths = []  # For storing paths of the original line images for display
 
-            # Extract lines and perform OCR
             for index, box in enumerate(boxes):
                 x_center, y_center, width, height = box
                 x_min = int(x_center - (width / 2))
                 y_min = int(y_center - (height / 2))
                 x_max = int(x_center + (width / 2))
                 y_max = int(y_center + (height / 2))
-                
-                # Crop the ORIGINAL image here
-                extracted_line = original_image.crop((x_min, y_min, x_max, y_max))
-                
-                # Use BytesIO to convert the image to a streamable format for Streamlit
-                buffer = io.BytesIO()
-                extracted_line.save(buffer, format="JPEG")
-                buffer.seek(0)
-                
-                # Display the extracted line image
-                st.image(buffer, use_column_width=True)
-                
-                # Save the thresholded line image for OCR
+
+                # Crop the ORIGINAL image and save
+                original_line = original_image.crop((x_min, y_min, x_max, y_max))
+                original_line_path = os.path.join(temp_dir, f"original_line_{index}.jpg")
+                original_line.save(original_line_path)
+                original_line_paths.append(original_line_path)
+
+                # Crop the THRESHOLDED image and save for OCR
+                extracted_line = thresholded_image.crop((x_min, y_min, x_max, y_max))
                 detected_line_path = os.path.join(temp_dir, f"detected_line_{index}.jpg")
-                thresholded_image.crop((x_min, y_min, x_max, y_max)).save(detected_line_path)
-                
-                # Perform OCR on this specific detected line (thresholded one)
-                recognized_text = perform_ocr_on_detected_lines([detected_line_path])[0]
+                extracted_line.save(detected_line_path)
+                detected_line_paths.append(detected_line_path)
+
+            # Perform OCR on thresholded lines
+            recognized_texts = perform_ocr_on_detected_lines(detected_line_paths)
 
-                # Display the OCR result
+            # Display the results
+            for original_img_path, text in zip(original_line_paths, recognized_texts):
+                st.image(original_img_path, use_column_width=True)
                 st.markdown(
-                    f"<p style='font-size: 18px; font-weight: bold;'>{recognized_text}</p>",
+                    f"<p style='font-size: 18px; font-weight: bold;'>{text}</p>",
                     unsafe_allow_html=True
                 )
                 # Add a small break for better spacing
@@ -359,6 +356,7 @@ def display_detected_lines(original_path, output_path):
         st.error("Annotation file (.txt) not found!")
 
 
+
 def perform_ocr_on_detected_lines(detected_line_paths):
     """
     Performs OCR on the provided list of detected line image paths.