nagasurendra commited on
Commit
1724c7a
·
verified ·
1 Parent(s): a128de1

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +49 -65
app.py CHANGED
@@ -14,7 +14,7 @@ from ultralytics import YOLO
14
  import ultralytics
15
  import time
16
  import piexif
17
- import zipfile
18
 
19
  # Set YOLO config directory
20
  os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics"
@@ -44,9 +44,9 @@ detected_issues: List[str] = []
44
  gps_coordinates: List[List[float]] = []
45
  last_metrics: Dict[str, Any] = {}
46
  frame_count: int = 0
47
- SAVE_IMAGE_INTERVAL = 1
48
 
49
- # Detection classes
50
  DETECTION_CLASSES = ["Longitudinal", "Pothole", "Transverse"]
51
 
52
  # Debug: Check environment
@@ -60,24 +60,9 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
60
  print(f"Using device: {device}")
61
  model = YOLO('./data/best.pt').to(device)
62
  if device == "cuda":
63
- model.half()
64
  print(f"Model classes: {model.names}")
65
 
66
- def zip_directory(folder_path: str, zip_path: str) -> str:
67
- """Zip all files in a directory."""
68
- try:
69
- with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
70
- for root, _, files in os.walk(folder_path):
71
- for file in files:
72
- file_path = os.path.join(root, file)
73
- arcname = os.path.relpath(file_path, folder_path)
74
- zipf.write(file_path, arcname)
75
- return zip_path
76
- except Exception as e:
77
- logging.error(f"Failed to zip {folder_path}: {str(e)}")
78
- log_entries.append(f"Error: Failed to zip {folder_path}: {str(e)}")
79
- return ""
80
-
81
  def generate_map(gps_coords: List[List[float]], items: List[Dict[str, Any]]) -> str:
82
  map_path = os.path.join(OUTPUT_DIR, "map_temp.png")
83
  plt.figure(figsize=(4, 4))
@@ -126,9 +111,9 @@ def write_flight_log(frame_count: int, gps_coord: List[float], timestamp: str) -
126
  def check_image_quality(frame: np.ndarray, input_resolution: int) -> bool:
127
  height, width, _ = frame.shape
128
  frame_resolution = width * height
129
- if frame_resolution < 12_000_000:
130
  log_entries.append(f"Frame {frame_count}: Resolution {width}x{height} ({frame_resolution/1e6:.2f}MP) below 12MP, non-compliant")
131
- if frame_resolution < input_resolution:
132
  log_entries.append(f"Frame {frame_count}: Output resolution {width}x{height} below input resolution")
133
  return False
134
  return True
@@ -156,6 +141,10 @@ def generate_line_chart() -> Optional[str]:
156
  plt.close()
157
  return chart_path
158
 
 
 
 
 
159
  def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
160
  global frame_count, last_metrics, detected_counts, detected_issues, gps_coordinates, log_entries
161
  frame_count = 0
@@ -168,14 +157,14 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
168
  if video is None:
169
  log_entries.append("Error: No video uploaded")
170
  logging.error("No video uploaded")
171
- return None, json.dumps({"error": "No video uploaded"}, indent=2), "\n".join(log_entries), [], None, None, None, None, None
172
 
173
  start_time = time.time()
174
  cap = cv2.VideoCapture(video)
175
  if not cap.isOpened():
176
  log_entries.append("Error: Could not open video file")
177
  logging.error("Could not open video file")
178
- return None, json.dumps({"error": "Could not open video file"}, indent=2), "\n".join(log_entries), [], None, None, None, None, None
179
 
180
  frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
181
  frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
@@ -183,18 +172,32 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
183
  fps = cap.get(cv2.CAP_PROP_FPS)
184
  total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
185
  expected_duration = total_frames / fps if fps > 0 else 0
186
- log_entries.append(f"Input video: {frame_width}x{frame_height} ({input_resolution/1e6:.2f}MP), {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds, Frame skip: {frame_skip}")
187
- logging.info(f"Input video: {frame_width}x{frame_height} ({input_resolution/1e6:.2f}MP), {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds, Frame skip: {frame_skip}")
188
- print(f"Input video: {frame_width}x{frame_height} ({input_resolution/1e6:.2f}MP), {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds, Frame skip: {frame_skip}")
189
 
190
  out_width, out_height = resize_width, resize_height
191
  output_path = os.path.join(OUTPUT_DIR, "processed_output.mp4")
192
- out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (out_width, out_height))
193
- if not out.isOpened():
194
- log_entries.append("Error: Failed to initialize mp4v codec")
195
- logging.error("Failed to initialize mp4v codec")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
196
  cap.release()
197
- return None, json.dumps({"error": "mp4v codec failed"}, indent=2), "\n".join(log_entries), [], None, None, None, None, None
198
 
199
  processed_frames = 0
200
  all_detections = []
@@ -292,16 +295,20 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
292
 
293
  frame_time = (time.time() - frame_start) * 1000
294
  frame_times.append(frame_time)
295
- log_entries.append(f"Frame {frame_count}: Processed in {frame_time:.2f} ms")
 
 
 
 
 
 
 
 
 
 
296
  if len(log_entries) > 50:
297
  log_entries.pop(0)
298
 
299
- # Timeout check
300
- if time.time() - start_time > 600:
301
- log_entries.append("Error: Processing timeout after 600 seconds")
302
- logging.error("Processing timeout after 600 seconds")
303
- break
304
-
305
  while output_frame_count < total_frames and last_annotated_frame is not None:
306
  out.write(last_annotated_frame)
307
  output_frame_count += 1
@@ -342,10 +349,6 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
342
  chart_path = generate_line_chart()
343
  map_path = generate_map(gps_coordinates[-5:], all_detections)
344
 
345
- # Zip images and logs
346
- images_zip = zip_directory(CAPTURED_FRAMES_DIR, os.path.join(OUTPUT_DIR, "captured_frames.zip"))
347
- logs_zip = zip_directory(FLIGHT_LOG_DIR, os.path.join(OUTPUT_DIR, "flight_logs.zip"))
348
-
349
  return (
350
  output_path,
351
  json.dumps(last_metrics, indent=2),
@@ -353,9 +356,7 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
353
  detected_issues,
354
  chart_path,
355
  map_path,
356
- submission_json_path,
357
- images_zip,
358
- logs_zip
359
  )
360
 
361
  # Gradio interface
@@ -378,30 +379,13 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange")) as iface:
378
  map_output = gr.Image(label="Issue Locations Map")
379
  with gr.Row():
380
  logs_output = gr.Textbox(label="Logs", lines=5, interactive=False)
381
- with gr.Row():
382
- gr.Markdown("## Download Results")
383
- with gr.Row():
384
- json_download = gr.File(label="Download Data Lake JSON")
385
- images_zip_download = gr.File(label="Download Geotagged Images (ZIP)")
386
- logs_zip_download = gr.File(label="Download Flight Logs (ZIP)")
387
- video_download = gr.File(label="Download Processed Video")
388
 
389
  process_btn.click(
390
  process_video,
391
  inputs=[video_input, width_slider, height_slider, skip_slider],
392
- outputs=[
393
- video_output,
394
- metrics_output,
395
- logs_output,
396
- issue_gallery,
397
- chart_output,
398
- map_output,
399
- json_download,
400
- images_zip_download,
401
- logs_zip_download,
402
- video_download
403
- ]
404
  )
405
 
406
  if __name__ == "__main__":
407
- iface.launch()
 
14
  import ultralytics
15
  import time
16
  import piexif
17
+ import shutil
18
 
19
  # Set YOLO config directory
20
  os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics"
 
44
  gps_coordinates: List[List[float]] = []
45
  last_metrics: Dict[str, Any] = {}
46
  frame_count: int = 0
47
+ SAVE_IMAGE_INTERVAL = 1 # Save every frame with detections
48
 
49
+ # Detection classes (aligned with model classes, excluding 'Crocodile')
50
  DETECTION_CLASSES = ["Longitudinal", "Pothole", "Transverse"]
51
 
52
  # Debug: Check environment
 
60
  print(f"Using device: {device}")
61
  model = YOLO('./data/best.pt').to(device)
62
  if device == "cuda":
63
+ model.half() # Use half-precision (FP16)
64
  print(f"Model classes: {model.names}")
65
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
66
  def generate_map(gps_coords: List[List[float]], items: List[Dict[str, Any]]) -> str:
67
  map_path = os.path.join(OUTPUT_DIR, "map_temp.png")
68
  plt.figure(figsize=(4, 4))
 
111
  def check_image_quality(frame: np.ndarray, input_resolution: int) -> bool:
112
  height, width, _ = frame.shape
113
  frame_resolution = width * height
114
+ if frame_resolution < 12_000_000: # NHAI requires 12 MP
115
  log_entries.append(f"Frame {frame_count}: Resolution {width}x{height} ({frame_resolution/1e6:.2f}MP) below 12MP, non-compliant")
116
+ if frame_resolution < input_resolution: # Ensure output is not below input
117
  log_entries.append(f"Frame {frame_count}: Output resolution {width}x{height} below input resolution")
118
  return False
119
  return True
 
141
  plt.close()
142
  return chart_path
143
 
144
+ def generate_download_zip():
145
+ shutil.make_archive("outputs_bundle", 'zip', OUTPUT_DIR)
146
+ return "outputs_bundle.zip"
147
+
148
  def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
149
  global frame_count, last_metrics, detected_counts, detected_issues, gps_coordinates, log_entries
150
  frame_count = 0
 
157
  if video is None:
158
  log_entries.append("Error: No video uploaded")
159
  logging.error("No video uploaded")
160
+ return "processed_output.mp4", json.dumps({"error": "No video uploaded"}, indent=2), "\n".join(log_entries), [], None, None, None
161
 
162
  start_time = time.time()
163
  cap = cv2.VideoCapture(video)
164
  if not cap.isOpened():
165
  log_entries.append("Error: Could not open video file")
166
  logging.error("Could not open video file")
167
+ return "processed_output.mp4", json.dumps({"error": "Could not open video file"}, indent=2), "\n".join(log_entries), [], None, None, None
168
 
169
  frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
170
  frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
172
  fps = cap.get(cv2.CAP_PROP_FPS)
173
  total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
174
  expected_duration = total_frames / fps if fps > 0 else 0
175
+ log_entries.append(f"Input video: {frame_width}x{frame_height} ({input_resolution/1e6:.2f}MP), {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds")
176
+ logging.info(f"Input video: {frame_width}x{frame_height} ({input_resolution/1e6:.2f}MP), {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds")
177
+ print(f"Input video: {frame_width}x{frame_height} ({input_resolution/1e6:.2f}MP), {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds")
178
 
179
  out_width, out_height = resize_width, resize_height
180
  output_path = os.path.join(OUTPUT_DIR, "processed_output.mp4")
181
+ codecs = [('mp4v', '.mp4'), ('XVID', '.avi'), ('MJPG', '.avi')] # Prioritize mp4v
182
+ out = None
183
+ for codec, ext in codecs:
184
+ fourcc = cv2.VideoWriter_fourcc(*codec)
185
+ temp_output_path = os.path.join(OUTPUT_DIR, f"processed_output{ext}")
186
+ out = cv2.VideoWriter(temp_output_path, fourcc, fps, (out_width, out_height))
187
+ if out.isOpened():
188
+ output_path = temp_output_path
189
+ log_entries.append(f"Using codec: {codec}, output: {output_path}")
190
+ logging.info(f"Using codec: {codec}, output: {output_path}")
191
+ break
192
+ else:
193
+ log_entries.append(f"Failed to initialize codec: {codec}")
194
+ logging.warning(f"Failed to initialize codec: {codec}")
195
+
196
+ if not out or not out.isOpened():
197
+ log_entries.append("Error: All codecs failed to initialize video writer")
198
+ logging.error("All codecs failed to initialize video writer")
199
  cap.release()
200
+ return "processed_output.mp4", json.dumps({"error": "All codecs failed"}, indent=2), "\n".join(log_entries), [], None, None, None
201
 
202
  processed_frames = 0
203
  all_detections = []
 
295
 
296
  frame_time = (time.time() - frame_start) * 1000
297
  frame_times.append(frame_time)
298
+
299
+ detection_summary = {
300
+ "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
301
+ "video_timestamp": timestamp_str,
302
+ "frame": frame_count,
303
+ "gps": gps_coord,
304
+ "processing_time_ms": frame_time,
305
+ "detections": {label: sum(1 for det in frame_detections if det["label"] == label) for label in DETECTION_CLASSES}
306
+ }
307
+ data_lake_submission["analytics"].append(detection_summary)
308
+ log_entries.append(json.dumps(detection_summary, indent=2))
309
  if len(log_entries) > 50:
310
  log_entries.pop(0)
311
 
 
 
 
 
 
 
312
  while output_frame_count < total_frames and last_annotated_frame is not None:
313
  out.write(last_annotated_frame)
314
  output_frame_count += 1
 
349
  chart_path = generate_line_chart()
350
  map_path = generate_map(gps_coordinates[-5:], all_detections)
351
 
 
 
 
 
352
  return (
353
  output_path,
354
  json.dumps(last_metrics, indent=2),
 
356
  detected_issues,
357
  chart_path,
358
  map_path,
359
+ generate_download_zip() # Provide the zip link for all outputs
 
 
360
  )
361
 
362
  # Gradio interface
 
379
  map_output = gr.Image(label="Issue Locations Map")
380
  with gr.Row():
381
  logs_output = gr.Textbox(label="Logs", lines=5, interactive=False)
382
+ zip_download = gr.File(label="Download All Outputs (ZIP)")
 
 
 
 
 
 
383
 
384
  process_btn.click(
385
  process_video,
386
  inputs=[video_input, width_slider, height_slider, skip_slider],
387
+ outputs=[video_output, metrics_output, logs_output, issue_gallery, chart_output, map_output, zip_download]
 
 
 
 
 
 
 
 
 
 
 
388
  )
389
 
390
  if __name__ == "__main__":
391
+ iface.launch()