Spaces:
Runtime error
Runtime error
File size: 13,255 Bytes
5555b23 5510d67 5555b23 e4663b7 5555b23 cc54f39 5555b23 61f83f1 5555b23 6fceb8b 5555b23 5e5a95f 713cc4f 5e5a95f dd5a50a 713cc4f d8a1507 5e5a95f 5555b23 310acce 80164ec 5555b23 e4663b7 0fd84ac af47efa 0fd84ac af47efa e4663b7 9462bf0 e4663b7 a59ff1a cf747f8 e4663b7 5555b23 cf747f8 7206cfb 3c1fffc 5e5a95f 3c1fffc 5555b23 b64e8b0 83f55b5 13e34fb b64e8b0 3e03f25 83f55b5 5555b23 6fceb8b 5555b23 6fceb8b 5555b23 5510d67 5555b23 73ac9a9 5555b23 9cdef11 5555b23 9cdef11 5555b23 9cdef11 5555b23 9cdef11 5555b23 9cdef11 5555b23 9cdef11 5555b23 9cdef11 5555b23 9cdef11 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 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 |
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
import io
from skimage import filters
# Load the OCR model outside the function to prevent reloading it every time
ocr_model = load_model()
def main():
st.title("منصة تجريبية لمكتب أضاميم")
st.title("Arabic Manuscript OCR")
st.warning("يرجى تحميل صورة ذو دقة عالية للحصول على نتيجة أفضل (حوالي 5000*3000)")
# 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 get_dynamic_kernel(img_height):
# Set the kernel size to be 1% of the image height
kernel_height = int(0.001 * img_height)
# Ensure the kernel height is odd
if kernel_height % 2 == 0:
kernel_height += 1
# Set minimum and maximum limits
kernel_height = max(1, kernel_height) # Minimum limit
kernel_height = min(11, kernel_height) # Maximum limit
st.text(f"img_height : {img_height}")
st.text(f"kernel_height : {kernel_height}")
return np.ones((kernel_height, 1), np.uint8)
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, (3, 3), 0)
# Apply adaptive thresholding
# Let's assume `gray_image` is your grayscale image array
height, width = gray_image.shape
# Example: setting block size to 1/30th of the average image dimension, making sure it's odd.
block_size = ((height + width) // 2) // 30
block_size = block_size + 1 if block_size % 2 == 0 else block_size
# Example: setting offset to a small fraction of the global mean intensity.
offset = np.mean(gray_image) * 0.05
adaptive_threshold = filters.threshold_local(gray_image, block_size, offset=offset, method='mean')
binary_adaptive = gray_image > adaptive_threshold
thresh = binary_adaptive
# Apply OTSU's thresholding
# _, thresh = cv2.threshold(binary_adaptive, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
# kernel = get_dynamic_kernel(thresh.shape[0])
# kernel = np.ones((5, 1), np.uint8)
# thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=2)
# # Apply median blur
# thresh = cv2.medianBlur(thresh, 3)
# # Apply morphological thinning
# thresh = cv2.ximgproc.thinning(thresh)
# # Apply morphological closing
# kernel = np.ones((3,3), np.uint8)
# thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)
# 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)
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 = [] # For storing paths of the thresholded line images for OCR
original_line_paths = [] # For storing paths of the original line images for display
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 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")
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 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;'>{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() |