rosebe commited on
Commit
16b8b27
·
1 Parent(s): 30fec0d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +15 -3
app.py CHANGED
@@ -136,6 +136,7 @@ def detect_image(image):
136
  results = model(image)
137
  return results.render()[0]
138
 
 
139
  def detect_video(video_path):
140
  video = cv2.VideoCapture(video_path)
141
  frame_rate = video.get(cv2.CAP_PROP_FPS)
@@ -143,6 +144,12 @@ def detect_video(video_path):
143
  # Create a directory to store the frames
144
  frames_dir = 'frames'
145
  os.makedirs(frames_dir, exist_ok=True)
 
 
 
 
 
 
146
 
147
  frame_count = 0
148
  process_frame_count = 0
@@ -162,6 +169,12 @@ def detect_video(video_path):
162
  # Process the frames with object detection and save the results
163
  results_dir = 'results'
164
  os.makedirs(results_dir, exist_ok=True)
 
 
 
 
 
 
165
 
166
  for i in range(process_frame_count):
167
  frame_path = os.path.join(frames_dir, f'frame_{i:04d}.jpg')
@@ -188,12 +201,11 @@ def detect_video(video_path):
188
  video_writer.release()
189
 
190
  # Clean up the temporary directories
191
- os.rmdir(frames_dir)
192
- os.rmdir(results_dir)
193
 
194
  return video_output_path
195
 
196
-
197
  # Create Gradio interfaces for different modes
198
  img_interface = gr.Interface(
199
  fn=detect_image,
 
136
  results = model(image)
137
  return results.render()[0]
138
 
139
+
140
  def detect_video(video_path):
141
  video = cv2.VideoCapture(video_path)
142
  frame_rate = video.get(cv2.CAP_PROP_FPS)
 
144
  # Create a directory to store the frames
145
  frames_dir = 'frames'
146
  os.makedirs(frames_dir, exist_ok=True)
147
+
148
+ # Remove existing contents of the frames directory
149
+ for file_name in os.listdir(frames_dir):
150
+ file_path = os.path.join(frames_dir, file_name)
151
+ if os.path.isfile(file_path):
152
+ os.remove(file_path)
153
 
154
  frame_count = 0
155
  process_frame_count = 0
 
169
  # Process the frames with object detection and save the results
170
  results_dir = 'results'
171
  os.makedirs(results_dir, exist_ok=True)
172
+
173
+ # Remove existing contents of the results directory
174
+ for file_name in os.listdir(results_dir):
175
+ file_path = os.path.join(results_dir, file_name)
176
+ if os.path.isfile(file_path):
177
+ os.remove(file_path)
178
 
179
  for i in range(process_frame_count):
180
  frame_path = os.path.join(frames_dir, f'frame_{i:04d}.jpg')
 
201
  video_writer.release()
202
 
203
  # Clean up the temporary directories
204
+ shutil.rmtree(frames_dir)
205
+ shutil.rmtree(results_dir)
206
 
207
  return video_output_path
208
 
 
209
  # Create Gradio interfaces for different modes
210
  img_interface = gr.Interface(
211
  fn=detect_image,