File size: 5,840 Bytes
0cf770a 13c174c 6d036be a5e19ff 0cf770a 503e7d9 9c365c8 a8ab437 0cf770a 0ad1f80 0cf770a a8ab437 0cf770a a8ab437 063d7dc a5e19ff 063d7dc a5e19ff 063d7dc a5e19ff 503e7d9 a5e19ff 063d7dc 43a3b0d 0cf770a a8ab437 0cf770a a8ab437 0cf770a 0ba87af a8ab437 a5e19ff a8ab437 0cf770a a8ab437 aa8fb5b 0cf770a |
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 |
import streamlit as st
import os
from os.path import join as osp
import subprocess
import zipfile
import io
import shutil
import sys
from PIL import Image
import tempfile
os.environ["HYDRA_FULL_ERROR"] = "1"
def GET_PROJECT_ROOT():
count = 0
# goto the root folder of LogBar
current_abspath = os.path.abspath(__file__)
while True:
if count > 1000:
print("Can find root error")
sys.exit()
if os.path.split(current_abspath)[1] == 'text-remove':
project_root = current_abspath
break
else:
current_abspath = os.path.dirname(current_abspath)
return project_root
def run_bash_script(input_image_path, output_path, progress_placeholder, status_text):
bash_command = f"bash config/text_detection.sh -s {input_image_path} -t {output_path}"
process = subprocess.Popen(bash_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
progress = 0
for line in process.stdout:
st.text(line.strip())
progress += 0.1
progress_placeholder.progress(min(progress, 1.0))
# Capture and display stderr
stderr_output = process.stderr.read()
if stderr_output:
status_text.error("Error output:")
st.code(stderr_output, language="bash")
rc = process.wait()
return rc, stderr_output
def zip_result_files(result_folder):
zip_buffer = io.BytesIO()
with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
for root, _, files in os.walk(result_folder):
for file in files:
if file.endswith(".png"):
file_path = os.path.join(root, file)
arcname = os.path.relpath(file_path, result_folder)
zip_file.write(file_path, arcname)
return zip_buffer
def clear_folder(folder_path):
for filename in os.listdir(folder_path):
file_path = os.path.join(folder_path, filename)
try:
if os.path.isfile(file_path) or os.path.islink(file_path):
os.unlink(file_path) # Remove file or symlink
elif os.path.isdir(file_path):
shutil.rmtree(file_path, ignore_errors=True) # Remove directory and its contents
except Exception as e:
print(f'Failed to delete {file_path}. Reason: {e}')
def create_temp_structure():
# Create a temporary directory
temp_dir = tempfile.mkdtemp()
# Create test_folder
test_folder = os.path.join(temp_dir, "test_folder")
os.makedirs(test_folder, exist_ok=True)
# Create target_folder with mask and result subdirectories
target_folder = os.path.join(temp_dir, "target_folder")
os.makedirs(os.path.join(target_folder, "mask"), exist_ok=True)
os.makedirs(os.path.join(target_folder, "result"), exist_ok=True)
os.makedirs(os.path.join(target_folder, "bbox"), exist_ok=True)
return temp_dir, test_folder, target_folder
st.title("Text Detection App")
# file_name = " || ".join(os.listdir('craft_pytorch'))
# st.write(file_name)
uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
# Create a temporary directory for processing
# Save the uploaded file temporarily
temp_dir, input_path, output_path = create_temp_structure()
# st.write(f"Temp dir: {temp_dir}")
# os.makedirs(input_path, exist_ok=True)
# os.makedirs(osp(output_path, "result"), exist_ok=True)
# os.makedirs(osp(output_path, "mask"), exist_ok=True)
input_file_path = os.path.join(input_path, uploaded_file.name)
image = Image.open(uploaded_file)
# image.save(os.path.join(PROJECT_ROOT, input_file_path))
image.save(input_file_path)
# file_name = " || ".join(os.listdir(f'{temp_dir}/target_folder'))
# st.write(file_name)
progress_placeholder = st.empty()
status_text = st.empty()
status_text.text("Running text detection...")
st.write(f'the input file path {input_path}')
st.write(f'the output file path {output_path}')
#os.makedirs(input_path, exist_ok=True)
#os.makedirs(osp(output_path, "result"), exist_ok=True)
#os.makedirs(osp(output_path, "mask"), exist_ok=True)
rc, stderr_output = run_bash_script(input_path, output_path, progress_placeholder, status_text)
mask_file_name = " || ".join(os.listdir('{temp_dir}/target_folder/mask'))
bbox_file_name = " || ".join(os.listdir('{temp_dir}/target_folder/bbox'))
result_file_name = " || ".join(os.listdir('{temp_dir}/target_folder/bbox'))
print(f'mask_file_name: {mask_file_name}')
print(f'bbox_file_name: {bbox_file_name}')
print(f'result_file_name: {result_file_name}')
if rc == 0:
st.write("Text detection completed successfully!")
status_text.text("Text detection completed successfully!")
result_folder = os.path.join(output_path, "result")
if os.path.exists(result_folder):
st.write("You can now download the results.")
# Add download button
zip_buffer = zip_result_files(result_folder)
st.download_button(
label="Download Results",
data=zip_buffer.getvalue(),
file_name="text_detection_results.zip",
mime="application/zip"
)
else:
st.error("Result folder not found. The text detection might have failed.")
else:
st.error(f"Text detection failed with return code {rc}")
if stderr_output:
st.error("Error details:")
st.code(stderr_output, language="bash")
st.write("Note: The download button will appear after running text detection.")
|