nguyen-brat commited on
Commit
063d7dc
·
1 Parent(s): aa8fb5b
Files changed (1) hide show
  1. app.py +24 -48
app.py CHANGED
@@ -5,28 +5,13 @@ import subprocess
5
  import zipfile
6
  import io
7
  import shutil
 
8
  import sys
9
  from PIL import Image
10
  import tempfile
11
- import uuid
12
 
13
  os.environ["HYDRA_FULL_ERROR"] = "1"
14
 
15
- def GET_PROJECT_ROOT():
16
- count = 0
17
- # goto the root folder of LogBar
18
- current_abspath = os.path.abspath(__file__)
19
- while True:
20
- if count > 1000:
21
- print("Can find root error")
22
- sys.exit()
23
- if os.path.split(current_abspath)[1] == 'text-remove':
24
- project_root = current_abspath
25
- break
26
- else:
27
- current_abspath = os.path.dirname(current_abspath)
28
- return project_root
29
-
30
  def run_bash_script(input_image_path, output_path, progress_placeholder, status_text):
31
  bash_command = f"bash config/text_detection.sh -s {input_image_path} -t {output_path}"
32
  process = subprocess.Popen(bash_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
@@ -37,7 +22,6 @@ def run_bash_script(input_image_path, output_path, progress_placeholder, status_
37
  progress += 0.1
38
  progress_placeholder.progress(min(progress, 1.0))
39
 
40
- # Capture and display stderr
41
  stderr_output = process.stderr.read()
42
  if stderr_output:
43
  status_text.error("Error output:")
@@ -59,52 +43,37 @@ def zip_result_files(result_folder):
59
 
60
  return zip_buffer
61
 
62
- def clear_folder(folder_path):
63
- for filename in os.listdir(folder_path):
64
- file_path = os.path.join(folder_path, filename)
65
- try:
66
- if os.path.isfile(file_path) or os.path.islink(file_path):
67
- os.unlink(file_path) # Remove file or symlink
68
- elif os.path.isdir(file_path):
69
- shutil.rmtree(file_path, ignore_errors=True) # Remove directory and its contents
70
- except Exception as e:
71
- print(f'Failed to delete {file_path}. Reason: {e}')
72
-
73
- def create_persistent_structure():
74
- # Create a unique directory name
75
- unique_id = str(uuid.uuid4())
76
-
77
- # Create a persistent directory in the current working directory
78
- persistent_dir = os.path.join(os.getcwd(), f"processing_{unique_id}")
79
 
80
  # Create test_folder
81
- test_folder = os.path.join(persistent_dir, "test_folder")
82
  os.makedirs(test_folder, exist_ok=True)
83
 
84
  # Create target_folder with mask and result subdirectories
85
- target_folder = os.path.join(persistent_dir, "target_folder")
86
  os.makedirs(os.path.join(target_folder, "mask"), exist_ok=True)
87
  os.makedirs(os.path.join(target_folder, "result"), exist_ok=True)
88
  os.makedirs(os.path.join(target_folder, "bbox"), exist_ok=True)
89
 
90
- return persistent_dir, test_folder, target_folder
91
 
92
  st.title("Text Detection App")
93
- # file_name = " || ".join(os.listdir('craft_pytorch'))
94
- # st.write(file_name)
95
  uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
96
 
97
  if uploaded_file is not None:
98
  st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
99
 
100
- # Create a persistent directory for processing
101
- persistent_dir, input_path, output_path = create_persistent_structure()
102
- st.write(f"Processing directory: {persistent_dir}")
103
 
104
  input_file_path = os.path.join(input_path, uploaded_file.name)
105
  image = Image.open(uploaded_file)
106
  image.save(input_file_path)
107
-
108
  if st.button("Run Text Detection"):
109
  progress_placeholder = st.empty()
110
  status_text = st.empty()
@@ -124,8 +93,7 @@ if uploaded_file is not None:
124
  label="Download Results",
125
  data=zip_buffer.getvalue(),
126
  file_name="text_detection_results.zip",
127
- mime="application/zip",
128
- on_click=lambda: clear_folder(persistent_dir)
129
  )
130
  else:
131
  st.error("Result folder not found. The text detection might have failed.")
@@ -141,9 +109,9 @@ if uploaded_file is not None:
141
  status_text.empty()
142
 
143
  # Display directory contents for debugging
144
- st.write(f"Contents of {persistent_dir}:")
145
- for root, dirs, files in os.walk(persistent_dir):
146
- level = root.replace(persistent_dir, '').count(os.sep)
147
  indent = ' ' * 4 * (level)
148
  st.write(f"{indent}{os.path.basename(root)}/")
149
  subindent = ' ' * 4 * (level + 1)
@@ -151,3 +119,11 @@ if uploaded_file is not None:
151
  st.write(f"{subindent}{f}")
152
 
153
  st.write("Note: The download button will appear after running text detection.")
 
 
 
 
 
 
 
 
 
5
  import zipfile
6
  import io
7
  import shutil
8
+ import time
9
  import sys
10
  from PIL import Image
11
  import tempfile
 
12
 
13
  os.environ["HYDRA_FULL_ERROR"] = "1"
14
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
15
  def run_bash_script(input_image_path, output_path, progress_placeholder, status_text):
16
  bash_command = f"bash config/text_detection.sh -s {input_image_path} -t {output_path}"
17
  process = subprocess.Popen(bash_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
 
22
  progress += 0.1
23
  progress_placeholder.progress(min(progress, 1.0))
24
 
 
25
  stderr_output = process.stderr.read()
26
  if stderr_output:
27
  status_text.error("Error output:")
 
43
 
44
  return zip_buffer
45
 
46
+ def create_temp_structure():
47
+ # Create a temporary directory
48
+ temp_dir = tempfile.mkdtemp()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
49
 
50
  # Create test_folder
51
+ test_folder = os.path.join(temp_dir, "test_folder")
52
  os.makedirs(test_folder, exist_ok=True)
53
 
54
  # Create target_folder with mask and result subdirectories
55
+ target_folder = os.path.join(temp_dir, "target_folder")
56
  os.makedirs(os.path.join(target_folder, "mask"), exist_ok=True)
57
  os.makedirs(os.path.join(target_folder, "result"), exist_ok=True)
58
  os.makedirs(os.path.join(target_folder, "bbox"), exist_ok=True)
59
 
60
+ return temp_dir, test_folder, target_folder
61
 
62
  st.title("Text Detection App")
63
+
 
64
  uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
65
 
66
  if uploaded_file is not None:
67
  st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
68
 
69
+ # Create a temporary directory for processing
70
+ temp_dir, input_path, output_path = create_temp_structure()
71
+ st.write(f"Temp dir: {temp_dir}")
72
 
73
  input_file_path = os.path.join(input_path, uploaded_file.name)
74
  image = Image.open(uploaded_file)
75
  image.save(input_file_path)
76
+
77
  if st.button("Run Text Detection"):
78
  progress_placeholder = st.empty()
79
  status_text = st.empty()
 
93
  label="Download Results",
94
  data=zip_buffer.getvalue(),
95
  file_name="text_detection_results.zip",
96
+ mime="application/zip"
 
97
  )
98
  else:
99
  st.error("Result folder not found. The text detection might have failed.")
 
109
  status_text.empty()
110
 
111
  # Display directory contents for debugging
112
+ st.write(f"Contents of temp directory:")
113
+ for root, dirs, files in os.walk(temp_dir):
114
+ level = root.replace(temp_dir, '').count(os.sep)
115
  indent = ' ' * 4 * (level)
116
  st.write(f"{indent}{os.path.basename(root)}/")
117
  subindent = ' ' * 4 * (level + 1)
 
119
  st.write(f"{subindent}{f}")
120
 
121
  st.write("Note: The download button will appear after running text detection.")
122
+
123
+ # Cleanup function to be called when the Streamlit script reruns
124
+ def cleanup():
125
+ if 'temp_dir' in locals():
126
+ shutil.rmtree(temp_dir, ignore_errors=True)
127
+
128
+ # Register the cleanup function
129
+ st.on_script_run.set(cleanup)