Update app.py

app.py

@@ -1,5 +1,3 @@
-import asyncio
-import platform
 import cv2
 import torch
 import gradio as gr
@@ -14,15 +12,14 @@ from typing import List, Dict, Any, Optional
 from ultralytics import YOLO
 import ultralytics
 import time
-import exiftool
-import csv
+import piexif
 
 # Set YOLO config directory
 os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics"
 
 # Set up logging
 logging.basicConfig(
-    filename="drone_app.log",
+    filename="app.log",
     level=logging.INFO,
     format="%(asctime)s - %(levelname)s - %(message)s"
 )
@@ -47,23 +44,8 @@ last_metrics: Dict[str, Any] = {}
 frame_count: int = 0
 SAVE_IMAGE_INTERVAL = 1  # Save every frame with detections
 
-# SOP Parameters from Annexure-I
-DRONE_SPEED_MS = 5  # 5 m/s (18 km/hr)
-MIN_SATELLITES = 12
-IMAGE_OVERLAP = 0.85  # 85% front and side overlap
-MIN_RESOLUTION_MP = 12  # Minimum 12 MP
-RECORDING_ANGLE = 90  # Nadir (90 degrees)
-IMAGE_FORMAT = "JPEG"
-
-# Annexure-III Operations and Maintenance parameters
-DETECTION_CLASSES = [
-    "Potholes", "Edge Drops", "Crack", "Raveling", "Rain Cut Embankments",
-    "Authorized Median Opening", "Unauthorized Median Opening",
-    "Intersection/Crossroads", "Temporary Encroachments", "Permanent Encroachments",
-    "Missing Lane Markings", "Missing Boundary Wall", "Damaged Boundary Wall",
-    "Open Drain", "Covered Drain", "Blocked Drain", "Unclean Drain",
-    "Missing Dissipation Basin"
-]
+# Detection classes (as per original code)
+DETECTION_CLASSES = ["Longitudinal", "Pothole", "Transverse", "Crack"]
 
 # Debug: Check environment
 print(f"Torch version: {torch.__version__}")
@@ -74,7 +56,7 @@ print(f"CUDA available: {torch.cuda.is_available()}")
 # Load custom YOLO model
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {device}")
-model = YOLO('./data/best.pt').to(device)  # Assumes model is trained for all DETECTION_CLASSES
+model = YOLO('./data/best.pt').to(device)
 if device == "cuda":
     model.half()  # Use half-precision (FP16)
 print(f"Model classes: {model.names}")
@@ -93,54 +75,51 @@ def generate_map(gps_coords: List[List[float]], items: List[Dict[str, Any]]) ->
 
 def write_geotag(image_path: str, gps_coord: List[float]) -> bool:
     try:
-        with exiftool.ExifToolHelper() as et:
-            et.set_tags(
-                [image_path],
-                {
-                    "EXIF:GPSLatitude": gps_coord[0],
-                    "EXIF:GPSLongitude": gps_coord[1],
-                    "EXIF:GPSLatitudeRef": "N" if gps_coord[0] >= 0 else "S",
-                    "EXIF:GPSLongitudeRef": "E" if gps_coord[1] >= 0 else "W"
-                }
-            )
+        lat = abs(gps_coord[0])
+        lon = abs(gps_coord[1])
+        lat_ref = "N" if gps_coord[0] >= 0 else "S"
+        lon_ref = "E" if gps_coord[1] >= 0 else "W"
+        exif_dict = piexif.load(image_path) if os.path.exists(image_path) else {"GPS": {}}
+        exif_dict["GPS"] = {
+            piexif.GPSIFD.GPSLatitudeRef: lat_ref,
+            piexif.GPSIFD.GPSLatitude: ((int(lat), 1), (0, 1), (0, 1)),
+            piexif.GPSIFD.GPSLongitudeRef: lon_ref,
+            piexif.GPSIFD.GPSLongitude: ((int(lon), 1), (0, 1), (0, 1))
+        }
+        piexif.insert(piexif.dump(exif_dict), image_path)
         return True
     except Exception as e:
         logging.error(f"Failed to geotag {image_path}: {str(e)}")
+        log_entries.append(f"Error: Failed to geotag {image_path}: {str(e)}")
         return False
 
 def write_flight_log(frame_count: int, gps_coord: List[float], timestamp: str) -> str:
-    log_path = os.path.join(FLIGHT_LOG_DIR, f"flight_log_{frame_count}.csv")
-    with open(log_path, 'w', newline='') as csvfile:
-        writer = csv.writer(csvfile)
-        writer.writerow(["Frame", "Timestamp", "Latitude", "Longitude", "Speed_ms", "Satellites", "Altitude_m"])
-        writer.writerow([frame_count, timestamp, gps_coord[0], gps_coord[1], DRONE_SPEED_MS, MIN_SATELLITES, 60])  # Example altitude
-    return log_path
-
-def check_sop_compliance(frame: np.ndarray, gps_coord: List[float], frame_count: int) -> bool:
+    log_path = os.path.join(FLIGHT_LOG_DIR, f"flight_log_{frame_count:06d}.csv")
+    try:
+        with open(log_path, 'w', newline='') as csvfile:
+            writer = csv.writer(csvfile)
+            writer.writerow(["Frame", "Timestamp", "Latitude", "Longitude", "Speed_ms", "Satellites", "Altitude_m"])
+            writer.writerow([frame_count, timestamp, gps_coord[0], gps_coord[1], 5.0, 12, 60])
+        return log_path
+    except Exception as e:
+        logging.error(f"Failed to write flight log {log_path}: {str(e)}")
+        log_entries.append(f"Error: Failed to write flight log {log_path}: {str(e)}")
+        return ""
+
+def check_image_quality(frame: np.ndarray) -> bool:
     height, width, _ = frame.shape
-    if width * height < MIN_RESOLUTION_MP * 1e6:  # Check resolution (12MP)
-        log_entries.append(f"Frame {frame_count}: Resolution below {MIN_RESOLUTION_MP}MP")
-        return False
-    if len(gps_coord) != 2 or not all(isinstance(x, float) for x in gps_coord):
-        log_entries.append(f"Frame {frame_count}: Invalid GPS coordinates")
+    if width * height < 12_000_000:  # 12 MP requirement
+        log_entries.append(f"Frame {frame_count}: Resolution below 12MP")
         return False
     return True
 
 def update_metrics(detections: List[Dict[str, Any]]) -> Dict[str, Any]:
     counts = Counter([det["label"] for det in detections])
-    metrics = {
+    return {
         "items": [{"type": k, "count": v} for k, v in counts.items()],
         "total_detections": len(detections),
-        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-        "sop_compliance": {
-            "drone_speed_ms": DRONE_SPEED_MS,
-            "image_overlap": IMAGE_OVERLAP,
-            "min_resolution_mp": MIN_RESOLUTION_MP,
-            "recording_angle_degrees": RECORDING_ANGLE,
-            "image_format": IMAGE_FORMAT
-        }
+        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     }
-    return metrics
 
 def generate_line_chart() -> Optional[str]:
     if not detected_counts:
@@ -157,7 +136,7 @@ def generate_line_chart() -> Optional[str]:
     plt.close()
     return chart_path
 
-async def process_video(video, resize_width=320, resize_height=240, frame_skip=5):
+def process_video(video, resize_width=320, resize_height=240, frame_skip=5):
     global frame_count, last_metrics, detected_counts, detected_issues, gps_coordinates, log_entries
     frame_count = 0
     detected_counts.clear()
@@ -219,7 +198,8 @@ async def process_video(video, resize_width=320, resize_height=240, frame_skip=5
     data_lake_submission = {
         "images": [],
         "flight_logs": [],
-        "analytics": []
+        "analytics": [],
+        "metrics": {}
     }
 
     while True:
@@ -233,6 +213,10 @@ async def process_video(video, resize_width=320, resize_height=240, frame_skip=5
         frame_start = time.time()
 
         frame = cv2.resize(frame, (out_width, out_height))
+        if not check_image_quality(frame):
+            log_entries.append(f"Frame {frame_count}: Skipped due to low resolution")
+            continue
+
         results = model(frame, verbose=False, conf=0.5, iou=0.7)
         annotated_frame = results[0].plot()
 
@@ -240,9 +224,7 @@ async def process_video(video, resize_width=320, resize_height=240, frame_skip=5
         timestamp_str = f"{int(frame_timestamp // 60)}:{int(frame_timestamp % 60):02d}"
 
         gps_coord = [17.385044 + (frame_count * 0.0001), 78.486671 + (frame_count * 0.0001)]
-        if not check_sop_compliance(frame, gps_coord, frame_count):
-            log_entries.append(f"Frame {frame_count}: SOP compliance check failed")
-            continue
+        gps_coordinates.append(gps_coord)
 
         frame_detections = []
         for detection in results[0].boxes:
@@ -250,7 +232,7 @@ async def process_video(video, resize_width=320, resize_height=240, frame_skip=5
             conf = float(detection.conf)
             box = detection.xyxy[0].cpu().numpy().astype(int).tolist()
             label = model.names[cls]
-            if label in DETECTION_CLASSES:
+            if label in DETECTION_CLASSES:  # Only process relevant classes
                 frame_detections.append({
                     "label": label,
                     "box": box,
@@ -266,10 +248,7 @@ async def process_video(video, resize_width=320, resize_height=240, frame_skip=5
             detection_frame_count += 1
             if detection_frame_count % SAVE_IMAGE_INTERVAL == 0:
                 captured_frame_path = os.path.join(CAPTURED_FRAMES_DIR, f"detected_{frame_count:06d}.jpg")
-                if not cv2.imwrite(captured_frame_path, annotated_frame):
-                    log_entries.append(f"Error: Failed to save {captured_frame_path}")
-                    logging.error(f"Failed to save {captured_frame_path}")
-                else:
+                if cv2.imwrite(captured_frame_path, annotated_frame):
                     if write_geotag(captured_frame_path, gps_coord):
                         detected_issues.append(captured_frame_path)
                         data_lake_submission["images"].append({
@@ -281,13 +260,17 @@ async def process_video(video, resize_width=320, resize_height=240, frame_skip=5
                         if len(detected_issues) > 100:
                             detected_issues.pop(0)
                     else:
-                        log_entries.append(f"Error: Failed to geotag {captured_frame_path}")
+                        log_entries.append(f"Frame {frame_count}: Geotagging failed")
+                else:
+                    log_entries.append(f"Error: Failed to save {captured_frame_path}")
+                    logging.error(f"Failed to save {captured_frame_path}")
 
         flight_log_path = write_flight_log(frame_count, gps_coord, timestamp_str)
-        data_lake_submission["flight_logs"].append({
-            "path": flight_log_path,
-            "frame": frame_count
-        })
+        if flight_log_path:
+            data_lake_submission["flight_logs"].append({
+                "path": flight_log_path,
+                "frame": frame_count
+            })
 
         out.write(annotated_frame)
         output_frame_count += 1
@@ -298,7 +281,6 @@ async def process_video(video, resize_width=320, resize_height=240, frame_skip=5
                 output_frame_count += 1
 
         detected_counts.append(len(frame_detections))
-        gps_coordinates.append(gps_coord)
         all_detections.extend(frame_detections)
 
         detection_summary = {
@@ -324,8 +306,14 @@ async def process_video(video, resize_width=320, resize_height=240, frame_skip=5
     data_lake_submission["gps_coordinates"] = gps_coordinates[-1] if gps_coordinates else [0, 0]
 
     submission_json_path = os.path.join(OUTPUT_DIR, "data_lake_submission.json")
-    with open(submission_json_path, 'w') as f:
-        json.dump(data_lake_submission, f, indent=2)
+    try:
+        with open(submission_json_path, 'w') as f:
+            json.dump(data_lake_submission, f, indent=2)
+        log_entries.append(f"Submission JSON saved: {submission_json_path}")
+        logging.info(f"Submission JSON saved: {submission_json_path}")
+    except Exception as e:
+        log_entries.append(f"Error: Failed to save submission JSON: {str(e)}")
+        logging.error(f"Failed to save submission JSON: {str(e)}")
 
     cap.release()
     out.release()
@@ -359,10 +347,10 @@ async def process_video(video, resize_width=320, resize_height=240, frame_skip=5
 
 # Gradio interface
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange")) as iface:
-    gr.Markdown("# NHAI Drone Analytics Dashboard")
+    gr.Markdown("# NHAI Road Defect Detection Dashboard")
     with gr.Row():
         with gr.Column(scale=3):
-            video_input = gr.Video(label="Upload Drone Video")
+            video_input = gr.Video(label="Upload Video")
             width_slider = gr.Slider(320, 640, value=320, label="Output Width", step=1)
             height_slider = gr.Slider(240, 480, value=240, label="Output Height", step=1)
             skip_slider = gr.Slider(1, 10, value=5, label="Frame Skip", step=1)
@@ -384,8 +372,5 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange")) as iface:
         outputs=[video_output, metrics_output, logs_output, issue_gallery, chart_output, map_output]
     )
 
-if platform.system() == "Emscripten":
-    asyncio.ensure_future(process_video())
-else:
-    if __name__ == "__main__":
-        iface.launch()
+if __name__ == "__main__":
+    iface.launch()