Update app.py

YOLO5V with OpenPilot Trajectory Prediction

app.py

@@ -1,61 +1,155 @@
-from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
+from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
 import datetime
 import requests
 import pytz
 import yaml
 from tools.final_answer import FinalAnswerTool
-
+from ultralytics import YOLO  # YOLOv5 model
+import cv2
+import numpy as np
 from Gradio_UI import GradioUI
 
-# Below is an example of a tool that does nothing. Amaze us with your creativity !
 @tool
-def my_custom_tool(arg1:str, arg2:int)-> str: #it's import to specify the return type
-    #Keep this format for the description / args / args description but feel free to modify the tool
-    """A tool that does nothing yet 
-    Args:
-        arg1: the first argument
-        arg2: the second argument
+def get_yolov5_coco_detections(image_path: str) -> dict:
+    """Detects objects using YOLOv5 on the COCO dataset and provides structured outputs."""
+    model = YOLO("yolov5s.pt")
+    image = cv2.imread(image_path)
+    results = model(image)
+
+    detections = []
+    if results:
+        for r in results.pred[0]:  
+            x1, y1, x2, y2, conf, cls = r.tolist()
+            class_name = model.names[int(cls)]
+            detections.append({"object": class_name, "confidence": conf, "bbox": (x1, y1, x2, y2)})
+
+    return {"detected_objects": detections} if detections else {"detected_objects": []}
+
+@tool
+def detect_road_lanes(image_path: str) -> dict:
+    """Detects road lanes using a YOLOv5 model trained for lane detection."""
+    model = YOLO("yolov5-lane.pt")
+    image = cv2.imread(image_path)
+    results = model(image)
+
+    lane_detections = []
+    if results:
+        for r in results.pred[0]:  
+            x1, y1, x2, y2, conf, cls = r.tolist()
+            lane_detections.append({"lane": f"Lane {cls}", "confidence": conf, "bbox": (x1, y1, x2, y2)})
+
+    return {"detected_lanes": lane_detections} if lane_detections else {"detected_lanes": []}
+
+@tool
+def driving_situation_analyzer(image_path: str) -> dict:
+    """Analyzes road conditions by integrating object detections and lane information."""
+    objects_info = get_yolov5_coco_detections(image_path)
+    lanes_info = detect_road_lanes(image_path)
+
+    detected_objects = objects_info.get("detected_objects", [])
+    detected_lanes = lanes_info.get("detected_lanes", [])
+
+    situation = []
+
+    if any(obj["object"] in ["car", "truck", "bus"] for obj in detected_objects):
+        situation.append("Traffic detected ahead, maintain safe distance.")
+
+    if any(obj["object"] == "pedestrian" for obj in detected_objects):
+        situation.append("Pedestrian detected, be prepared to stop.")
+
+    if any(obj["object"] == "traffic light" for obj in detected_objects):
+        situation.append("Traffic light detected, slow down if red.")
+
+    if not detected_lanes:
+        situation.append("Lane markings not detected, potential risk of veering.")
+
+    elif len(detected_lanes) == 1:
+        situation.append("Single lane detected, ensure proper lane following.")
+
+    elif len(detected_lanes) >= 2:
+        situation.append("Multiple lanes detected, stay within lane boundaries.")
+
+    return {
+        "situation_summary": " | ".join(situation) if situation else "Road situation unclear, proceed with caution.",
+        "detected_objects": detected_objects,
+        "detected_lanes": detected_lanes
+    }
+
+@tool
+def predict_trajectory(image_path: str) -> dict:
+    """Predicts vehicle trajectory based on the driving situation analysis.
+
+    Uses OpenPilot-style motion prediction based on detected lanes, objects, and road conditions.
     """
-    return "What magic will you build ?"
+    analysis = driving_situation_analyzer(image_path)
+    detected_objects = analysis["detected_objects"]
+    detected_lanes = analysis["detected_lanes"]
+    summary = analysis["situation_summary"]
+
+    trajectory = []
+
+    # Define simple trajectory logic
+    if "Traffic detected" in summary:
+        trajectory.append("Reduce speed, maintain a safe distance.")
+
+    if "Pedestrian detected" in summary:
+        trajectory.append("Prepare for sudden braking or yielding.")
+
+    if "Traffic light detected" in summary:
+        trajectory.append("Adjust speed based on light status.")
+
+    if "Lane markings not detected" in summary:
+        trajectory.append("Risk of lane departure, drive cautiously.")
+
+    if len(detected_lanes) >= 2:
+        trajectory.append("Stay centered in the lane, adjust for merging vehicles.")
+
+    # Generate simple trajectory points (Mocked example)
+    future_positions = []
+    x, y = 0, 0  
+    for t in range(10):  # Predict 10 future steps
+        x += np.random.uniform(-0.5, 0.5)  # Small lateral deviation
+        y += 1  # Move forward
+        future_positions.append((x, y))
+
+    return {
+        "trajectory_recommendation": " | ".join(trajectory) if trajectory else "Maintain current path.",
+        "future_positions": future_positions
+    }
 
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
-    """A tool that fetches the current local time in a specified timezone.
-    Args:
-        timezone: A string representing a valid timezone (e.g., 'America/New_York').
-    """
+    """Fetches the current local time in a specified timezone."""
     try:
-        # Create timezone object
         tz = pytz.timezone(timezone)
-        # Get current time in that timezone
         local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
         return f"The current local time in {timezone} is: {local_time}"
     except Exception as e:
         return f"Error fetching time for timezone '{timezone}': {str(e)}"
 
-
 final_answer = FinalAnswerTool()
 
-# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
-# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 
-
 model = HfApiModel(
-max_tokens=2096,
-temperature=0.5,
-model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
-custom_role_conversions=None,
+    max_tokens=2096,
+    temperature=0.5,
+    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
+    custom_role_conversions=None,
 )
 
-
-# Import tool from Hub
 image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
 
 with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
-    
+
 agent = CodeAgent(
     model=model,
-    tools=[final_answer], ## add your tools here (don't remove final answer)
+    tools=[
+        final_answer,
+        get_yolov5_coco_detections,
+        detect_road_lanes,
+        driving_situation_analyzer,
+        predict_trajectory
+    ],
     max_steps=6,
     verbosity_level=1,
     grammar=None,
@@ -65,5 +159,4 @@ agent = CodeAgent(
     prompt_templates=prompt_templates
 )
 
-
 GradioUI(agent).launch()