path_finding_in_video.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1DPM9huhup7Z"
      },
      "outputs": [],
      "source": [
        "!pip install ultralytics"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import cv2\n",
        "import random\n",
        "import time\n",
        "from scipy.spatial.distance import euclidean\n",
        "from ultralytics import YOLO\n",
        "from google.colab import files\n",
        "import matplotlib.colors as mcolors  # Correct import for colors\n",
        "\n",
        "# Upload file via Colab's file upload functionality\n",
        "uploaded = files.upload()\n",
        "\n",
        "# Check if a file is uploaded and get the file name\n",
        "if uploaded:\n",
        "    video_path = next(iter(uploaded))  # Get the first uploaded file's name\n",
        "    print(f\"Uploaded file: {video_path}\")\n",
        "\n",
        "    # Load YOLO model\n",
        "    model = YOLO('/content/drive/MyDrive/best_p6.pt')\n",
        "\n",
        "    # Open the video file\n",
        "    cap = cv2.VideoCapture(video_path)\n",
        "\n",
        "    # Check if video opened successfully\n",
        "    if not cap.isOpened():\n",
        "        print(\"Error: Could not open video.\")\n",
        "        exit()\n",
        "\n",
        "    # Get the video frame width and height\n",
        "    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
        "    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
        "\n",
        "    # List of 10 different colors\n",
        "    color_list = [\n",
        "        \"red\", \"blue\", \"green\", \"orange\", \"purple\",\n",
        "        \"pink\", \"cyan\", \"brown\", \"magenta\", \"yellow\"\n",
        "    ]\n",
        "\n",
        "    # Notify user about valid number of agents\n",
        "    print(\"Please enter the number of agents (between 1 and 10):\")\n",
        "    num_agents = int(input(\"Enter the number of agents: \"))\n",
        "\n",
        "    # Validate the number of agents\n",
        "    if num_agents < 1 or num_agents > 10:\n",
        "        print(\"Invalid number of agents. Please enter a value between 1 and 10.\")\n",
        "        exit()\n",
        "\n",
        "    # Initialize agent parameters\n",
        "    collected_trash = set()\n",
        "    agent_positions = [(random.randint(0, frame_width), random.randint(0, frame_height)) for _ in range(num_agents)]  # Default resolution\n",
        "    agent_paths = [[] for _ in range(num_agents)]\n",
        "\n",
        "    # APF Parameters\n",
        "    attraction_strength = 1.0\n",
        "    repulsion_strength = 1.5\n",
        "    repulsion_distance = 50  # Distance at which agents start to repel each other\n",
        "    step_size = 5  # Step size for agents' movements\n",
        "\n",
        "    # Create a VideoWriter object to save the video\n",
        "    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec for mp4 format\n",
        "    out = cv2.VideoWriter('waste_with_agents.mp4', fourcc, 20.0, (frame_width, frame_height))  # 20 fps, resolution of video\n",
        "\n",
        "    # Collect trash for agents with APF and dynamic recalculation\n",
        "    def dynamic_path_recalculation(trash_coords, num_agents, frame):\n",
        "            # Create the frame for this step (draw agents and trash)\n",
        "            frame_with_agents = frame.copy()\n",
        "\n",
        "            # Draw agents with their unique color\n",
        "            for agent_idx in range(num_agents):\n",
        "                agent_color = color_list[agent_idx % len(color_list)]  # Assign color from color_list\n",
        "                color_rgb = mcolors.to_rgb(agent_color)  # Convert color to RGB format\n",
        "                color_bgr = tuple([int(c * 255) for c in color_rgb])  # Convert to BGR for OpenCV\n",
        "                cv2.circle(frame_with_agents, tuple(map(int, agent_positions[agent_idx])), 10, color_bgr, -1)  # Colored circles for agents\n",
        "\n",
        "            # Draw trash with a fixed color (green)\n",
        "            for trash in trash_coords:\n",
        "                cv2.circle(frame_with_agents, tuple(map(int, trash)), 5, (0, 255, 0), -1)  # Green circles for trash\n",
        "\n",
        "            # Draw the paths with their unique color (per agent)\n",
        "            for agent_idx in range(num_agents):\n",
        "                path_coords = [trash_coords[node] for node in agent_paths[agent_idx]]\n",
        "                for i in range(len(path_coords) - 1):\n",
        "                    start_point = path_coords[i]\n",
        "                    end_point = path_coords[i + 1]\n",
        "                    agent_color = color_list[agent_idx % len(color_list)]  # Assign color from color_list\n",
        "                    color_rgb = mcolors.to_rgb(agent_color)  # Convert color to RGB format\n",
        "                    color_bgr = tuple([int(c * 255) for c in color_rgb])  # Convert to BGR for OpenCV\n",
        "                    cv2.line(frame_with_agents, tuple(map(int, start_point)), tuple(map(int, end_point)), color_bgr, 2)  # Draw path with agent's color\n",
        "\n",
        "            # Write the frame to the video\n",
        "            out.write(frame_with_agents)\n",
        "\n",
        "            if len(collected_trash) == len(trash_coords):\n",
        "                break\n",
        "\n",
        "        return agent_paths\n",
        "\n",
        "    start_time = time.time()  # Start the timer\n",
        "    while cap.isOpened():\n",
        "        ret, frame = cap.read()\n",
        "        if not ret:\n",
        "            break\n",
        "\n",
        "        # Convert frame to RGB (for YOLO)\n",
        "        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
        "\n",
        "        # Perform inference to detect trash objects in the current frame\n",
        "        results = model.predict(source=frame)\n",
        "\n",
        "        # Initialize a list to store trash coordinates for this frame\n",
        "        trash_coords = []\n",
        "        for result in results:\n",
        "            boxes = result.boxes\n",
        "            for box in boxes:\n",
        "                class_id = int(box.cls)\n",
        "                if class_id == 1:\n",
        "                    x_center, y_center, _, _ = box.xywh[0]\n",
        "                    trash_coords.append((x_center, y_center))\n",
        "\n",
        "        # If no trash is found, continue to next frame\n",
        "        if not trash_coords:\n",
        "            continue\n",
        "\n",
        "        # Collect trash for agents in this frame\n",
        "        agent_paths = dynamic_path_recalculation(trash_coords, num_agents, frame_rgb)\n",
        "\n",
        "    # Release the video capture object and the VideoWriter object\n",
        "    cap.release()\n",
        "    out.release()\n",
        "    cv2.destroyAllWindows()\n",
        "\n",
        "    print(\"Video saved as waste_with_agents.mp4\")"
      ],
      "metadata": {
        "id": "ItOeaiYmutwd"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}