Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -14,7 +14,7 @@ from ultralytics import YOLO
|
|
14 |
import ultralytics
|
15 |
import time
|
16 |
import piexif
|
17 |
-
import
|
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
|
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
|
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
|
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
|
188 |
-
print(f"Input video: {frame_width}x{frame_height} ({input_resolution/1e6:.2f}MP), {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds
|
189 |
|
190 |
out_width, out_height = resize_width, resize_height
|
191 |
output_path = os.path.join(OUTPUT_DIR, "processed_output.mp4")
|
192 |
-
|
193 |
-
|
194 |
-
|
195 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
196 |
cap.release()
|
197 |
-
return
|
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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
-
|
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 |
-
|
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()
|