nguyen-brat commited on
Commit
35ca994
·
1 Parent(s): 063d7dc
Files changed (1) hide show
  1. app.py +31 -19
app.py CHANGED
@@ -59,15 +59,33 @@ def create_temp_structure():
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)
@@ -93,7 +111,8 @@ if uploaded_file is not None:
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,21 +128,14 @@ if uploaded_file is not None:
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)
118
- for f in files:
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)
 
59
 
60
  return temp_dir, test_folder, target_folder
61
 
62
+ def clear_temp_folder(temp_dir):
63
+ if temp_dir and os.path.exists(temp_dir):
64
+ shutil.rmtree(temp_dir, ignore_errors=True)
65
+ st.success("Temporary files have been cleared.")
66
+
67
  st.title("Text Detection App")
68
 
69
+ # Use session state to store the temporary directory path
70
+ if 'temp_dir' not in st.session_state:
71
+ st.session_state.temp_dir = None
72
+
73
  uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
74
 
75
  if uploaded_file is not None:
76
  st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
77
 
78
+ # Create a temporary directory for processing if it doesn't exist
79
+ if not st.session_state.temp_dir:
80
+ temp_dir, input_path, output_path = create_temp_structure()
81
+ st.session_state.temp_dir = temp_dir
82
+ st.session_state.input_path = input_path
83
+ st.session_state.output_path = output_path
84
+ else:
85
+ temp_dir = st.session_state.temp_dir
86
+ input_path = st.session_state.input_path
87
+ output_path = st.session_state.output_path
88
+
89
  st.write(f"Temp dir: {temp_dir}")
90
 
91
  input_file_path = os.path.join(input_path, uploaded_file.name)
 
111
  label="Download Results",
112
  data=zip_buffer.getvalue(),
113
  file_name="text_detection_results.zip",
114
+ mime="application/zip",
115
+ on_click=lambda: clear_temp_folder(st.session_state.temp_dir)
116
  )
117
  else:
118
  st.error("Result folder not found. The text detection might have failed.")
 
128
  status_text.empty()
129
 
130
  # Display directory contents for debugging
131
+ if st.session_state.temp_dir:
132
+ st.write(f"Contents of temp directory:")
133
+ for root, dirs, files in os.walk(st.session_state.temp_dir):
134
+ level = root.replace(st.session_state.temp_dir, '').count(os.sep)
135
+ indent = ' ' * 4 * (level)
136
+ st.write(f"{indent}{os.path.basename(root)}/")
137
+ subindent = ' ' * 4 * (level + 1)
138
+ for f in files:
139
+ st.write(f"{subindent}{f}")
140
 
141
  st.write("Note: The download button will appear after running text detection.")