nguyen-brat commited on
Commit
0cf770a
1 Parent(s): 1683dda
Files changed (1) hide show
  1. app.py +118 -0
app.py CHANGED
@@ -0,0 +1,118 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import os
3
+ from os.path import join as osp
4
+ import subprocess
5
+ import zipfile
6
+ import io
7
+ import shutil
8
+ import time
9
+
10
+ def run_bash_script(input_image_path, output_path, progress_placeholder, status_text):
11
+ bash_command = f"bash config/text_detection.sh -s {input_image_path} -t {output_path}"
12
+ process = subprocess.Popen(bash_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
13
+
14
+ progress = 0
15
+ for line in process.stdout:
16
+ st.text(line.strip())
17
+ progress += 0.1
18
+ progress_placeholder.progress(min(progress, 1.0))
19
+
20
+ # Capture and display stderr
21
+ stderr_output = process.stderr.read()
22
+ if stderr_output:
23
+ status_text.error("Error output:")
24
+ st.code(stderr_output, language="bash")
25
+
26
+ rc = process.wait()
27
+ return rc, stderr_output
28
+
29
+ def zip_result_files(result_folder):
30
+ zip_buffer = io.BytesIO()
31
+
32
+ with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
33
+ for root, _, files in os.walk(result_folder):
34
+ for file in files:
35
+ if file.endswith(".png"):
36
+ file_path = os.path.join(root, file)
37
+ arcname = os.path.relpath(file_path, result_folder)
38
+ zip_file.write(file_path, arcname)
39
+
40
+ return zip_buffer
41
+
42
+ def clear_folder(folder_path):
43
+ for filename in os.listdir(folder_path):
44
+ file_path = os.path.join(folder_path, filename)
45
+ try:
46
+ if os.path.isfile(file_path) or os.path.islink(file_path):
47
+ os.unlink(file_path) # Remove file or symlink
48
+ elif os.path.isdir(file_path):
49
+ shutil.rmtree(file_path, ignore_errors=True) # Remove directory and its contents
50
+ except Exception as e:
51
+ print(f'Failed to delete {file_path}. Reason: {e}')
52
+
53
+ st.title("Text Detection App")
54
+
55
+ uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
56
+
57
+ if uploaded_file is not None:
58
+ input_path = "test_folder"
59
+ output_path = "target_test"
60
+ input_file_path = os.path.join(input_path, uploaded_file.name)
61
+ with open(input_file_path, "wb") as f:
62
+ f.write(uploaded_file.getbuffer())
63
+
64
+ if uploaded_file is not None:
65
+ st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
66
+
67
+ # Create a temporary directory for processing
68
+
69
+ # Save the uploaded file temporarily
70
+ input_path = "test_folder"
71
+ output_path = "target_test"
72
+ os.makedirs(input_path, exist_ok=True)
73
+ os.makedirs(osp(output_path, "result"), exist_ok=True)
74
+ os.makedirs(osp(output_path, "mask"), exist_ok=True)
75
+
76
+ input_file_path = os.path.join(input_path, uploaded_file.name)
77
+
78
+ if st.button("Run Text Detection"):
79
+ progress_placeholder = st.empty()
80
+ status_text = st.empty()
81
+
82
+ try:
83
+ status_text.text("Running text detection...")
84
+ os.makedirs(input_path, exist_ok=True)
85
+ os.makedirs(osp(output_path, "result"), exist_ok=True)
86
+ os.makedirs(osp(output_path, "mask"), exist_ok=True)
87
+ rc, stderr_output = run_bash_script(input_path, output_path, progress_placeholder, status_text)
88
+
89
+ if rc == 0:
90
+ status_text.text("Text detection completed successfully!")
91
+ result_folder = os.path.join(output_path, "result")
92
+ if os.path.exists(result_folder):
93
+ st.write("You can now download the results.")
94
+
95
+ # Add download button
96
+ zip_buffer = zip_result_files(result_folder)
97
+ st.download_button(
98
+ label="Download Results",
99
+ data=zip_buffer.getvalue(),
100
+ file_name="text_detection_results.zip",
101
+ mime="application/zip"
102
+ )
103
+ else:
104
+ st.error("Result folder not found. The text detection might have failed.")
105
+ else:
106
+ st.error(f"Text detection failed with return code {rc}")
107
+ if stderr_output:
108
+ st.error("Error details:")
109
+ st.code(stderr_output, language="bash")
110
+ except Exception as e:
111
+ st.error(f"An error occurred: {str(e)}")
112
+ finally:
113
+ # Clean up temporary files
114
+ clear_folder(osp(output_path, "mask"))
115
+ progress_placeholder.empty()
116
+ status_text.empty()
117
+
118
+ st.write("Note: The download button will appear after running text detection.")