Spaces:
Runtime error
Runtime error
File size: 11,288 Bytes
5555b23 6fceb8b 5555b23 6fceb8b 5555b23 6fceb8b 5555b23 73ac9a9 5555b23 73ac9a9 5555b23 73ac9a9 5555b23 73ac9a9 5555b23 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 |
import streamlit as st
import pandas as pd
import numpy as np
import cv2
from PIL import Image
import os
import zipfile
import gdown
from tempfile import TemporaryDirectory
from Utils import Loadlines, decode_predictions, load_model
import shutil
import tempfile
import subprocess
# Load the OCR model outside the function to prevent reloading it every time
ocr_model = load_model()
def main():
st.title("Arabic Manuscript OCR")
# download_folder()
# Load sample images
sample_images_dir = "sample_images"
sample_images = [os.path.join(sample_images_dir, img) for img in os.listdir(sample_images_dir) if img.endswith(('.png', '.jpg', '.jpeg'))]
# Display images in a grid and let user select
col1, col2, col3 = st.columns(3) # Adjust based on how many images you have; this example assumes 3
# Normalized size for display
display_size = (200, 300)
# Placeholder for selected image
selected_image_pil = None
processed_img_path = None
with col1:
if st.button("Select Image 1", key="img1"):
selected_image_pil = display_image(sample_images[0])
col1.image(resize_image(sample_images[0], display_size), use_column_width=True)
with col2:
if st.button("Select Image 2", key="img2"):
selected_image_pil = display_image(sample_images[1])
col2.image(resize_image(sample_images[1], display_size), use_column_width=True)
with col3:
if st.button("Select Image 3", key="img3"):
selected_image_pil = display_image(sample_images[2])
col3.image(resize_image(sample_images[2], display_size), use_column_width=True)
# Option to upload a new image
uploaded_image = st.file_uploader("Or upload a new image of the Arabic manuscript", type=["jpg", "jpeg", "png"])
if uploaded_image:
selected_image_pil = Image.open(uploaded_image)
st.image(selected_image_pil, caption="Uploaded Image", use_column_width=True)
if selected_image_pil:
thresh_pil = process_image(selected_image_pil)
st.image(thresh_pil, caption="Thresholded Image", use_column_width=True)
# processed_img_path = process_with_yolo(selected_image_pil)
processed_img_path = process_with_yolo(thresh_pil)
if processed_img_path:
st.image(processed_img_path, caption="Processed with YOLO", use_column_width=True)
txt_file_path = os.path.join('yolov3/runs/detect/mainlinedection/labels', os.path.basename(processed_img_path).replace(".jpg", ".txt"))
if os.path.exists(txt_file_path):
if uploaded_image:
original_img_path = uploaded_image # If you have saved it to a location
else:
original_img_path = next((img for img in sample_images if Image.open(img) == selected_image_pil), None) # Match selected image with sample images
display_detected_lines(original_img_path, processed_img_path)
else:
st.error("Annotation file (.txt) not found!")
else:
st.error("Error displaying the processed image.")
# display_files_in_directory('/home/user/app/detected_lines/')
# Function to display files in a directory
def display_files_in_directory(path="."):
if os.path.exists(path):
files = os.listdir(path)
st.write(f"Files in directory: {path}")
for file in files:
st.write(file)
else:
st.write(f"Directory {path} does not exist!")
def resize_image(image_path, size):
"""Function to resize an image"""
with Image.open(image_path) as img:
img = img.resize(size)
return img
def display_image(img):
"""Function to display an image. If img is a path, open it. Otherwise, just display it."""
if isinstance(img, str): # img is a file path
img = Image.open(img)
st.image(img, caption="Selected Image", use_column_width=True)
return img
# --------------------------------------------
def process_image(selected_image):
# Convert PIL image to OpenCV format
opencv_image = np.array(selected_image)
# Check if the image is grayscale or RGB
if len(opencv_image.shape) == 3 and opencv_image.shape[2] == 3: # RGB image
gray_image = cv2.cvtColor(opencv_image, cv2.COLOR_RGB2GRAY)
else: # Image is already grayscale
gray_image = opencv_image
# Ensure the image is 8-bit grayscale
if gray_image.dtype != np.uint8:
gray_image = (gray_image * 255).astype(np.uint8)
# Optionally apply Gaussian Blur
blurred_img = cv2.GaussianBlur(gray_image, (5, 5), 0)
# Apply OTSU's thresholding
_, thresh = cv2.threshold(blurred_img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
# Convert thresholded image back to PIL format to display in Streamlit
thresh_pil = Image.fromarray(thresh)
return thresh_pil
def get_detected_boxes(txt_path, img_width, img_height):
with open(txt_path, 'r') as f:
lines = f.readlines()
boxes = []
for line in lines:
parts = list(map(float, line.strip().split()))
# Assuming the format is: class x_center y_center width height confidence
x_center, y_center, width, height = parts[1:5]
# Convert to pixel values
x_center *= img_width
y_center *= img_height
width *= img_width
height *= img_height
boxes.append([x_center, y_center, width, height])
return boxes
# ---------------------------------------------------------------------------------
def download_ultralytics_yolov3_folder():
st.text("Downloading Ultralytics YOLOv3 folder from Google Drive. This may take a while...")
# url = 'https://drive.google.com/file/d/1n6YcqHl5Y2xRpWw7DQPrZ2FoAqfcf12I/view?usp=share_link'
url = 'https://drive.google.com/uc?id=1n6YcqHl5Y2xRpWw7DQPrZ2FoAqfcf12I'
output = 'ultralytics_yolov3.zip'
gdown.download(url, output, quiet=False)
# Extracting the zip file
with zipfile.ZipFile(output, 'r') as zip_ref:
zip_ref.extractall('.')
st.text("Folder extraction complete!")
os.remove(output) # Optional: remove the downloaded zip file after extraction.
st.text("Download and extraction completed!")
def process_with_yolo(img_pil):
with st.spinner('Downloading YOLOv3 folder...'):
download_ultralytics_yolov3_folder()
# display_files_in_directory('/home/user/app/')
# Save the image to a temporary file
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".jpg")
img_pil.save(temp_file.name)
if os.path.exists('yolov3/runs/detect/mainlinedection'):
shutil.rmtree('yolov3/runs/detect/mainlinedection')
cmd = [
'python', 'yolov3/detect.py',
'--source', temp_file.name, # use the temp file path here
'--weights', 'yolov3/runs/train/mainline/weights/best.pt',
'--save-txt',
'--save-conf',
'--imgsz', '672',
'--name', 'mainlinedection'
]
process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
stdout, stderr = process.communicate()
if process.returncode != 0:
st.error(f"YOLOv3 command failed with code {process.returncode}.")
# After processing the image, at the end of the function...
output_path = os.path.join('yolov3/runs/detect/mainlinedection', os.path.basename(temp_file.name))
if os.path.exists(output_path):
return output_path
else:
st.error("Processed image not found!")
return None
# Optional: You can print the output to Streamlit, though it might be extensive.
st.write(stdout.decode())
if stderr:
st.error(stderr.decode())
# Optional: Close and delete the temporary file
temp_file.close()
os.unlink(temp_file.name)
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
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
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 thresholded image instead of the original
extracted_line = thresholded_image.crop((x_min, y_min, x_max, y_max))
# Save the detected line image to the temporary directory
detected_line_path = os.path.join(temp_dir, f"detected_line_{index}.jpg")
extracted_line.save(detected_line_path)
detected_line_paths.append(detected_line_path)
# Perform OCR on detected lines
recognized_texts = perform_ocr_on_detected_lines(detected_line_paths)
# print("Decoded OCR Results:", recognized_texts)
# st.text(f"Detected Line: {recognized_texts}")
# Display the results
for img_path, text in zip(detected_line_paths, recognized_texts):
# st.image(img_path, caption=f"Detected Line: {text}", use_column_width=True)
st.image(img_path, use_column_width=True)
st.markdown(
f"<p style='font-size: 18px; font-weight: bold;'>{text}</p>",
unsafe_allow_html=True
)
# Add a small break for better spacing
st.markdown("<br>", unsafe_allow_html=True)
else:
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.
Args:
- detected_line_paths: List of paths to the detected line images.
Returns:
- A list of recognized text for each image.
"""
# Load the saved detected lines for OCR processing
test_dataset = Loadlines(detected_line_paths)
prediction_texts = []
for batch in test_dataset:
preds = ocr_model(batch)
pred_texts = decode_predictions(preds)
# st.text(f"Decoded OCR Results : {pred_texts}")
prediction_texts.extend(pred_texts)
return prediction_texts
if __name__ == "__main__":
main() |