Traffic Flow Prediction with YOLOv8

Detect vehicles in real traffic footage and predict flow with a GRU network.

Introduction

30 Second Summary

Traffic cameras record thousands of hours of footage at busy intersections every day. Knowing what vehicles are on the road right now is useful, but predicting what traffic will look like in the next few minutes is what keeps cities from grinding to a halt.

In this project, you will build a complete traffic flow prediction pipeline using real traffic camera footage. You will use YOLOv8 to detect vehicles in real traffic frames, then train a GRU neural network to forecast future traffic intensity from those detections.

What You'll Build

When you run the final cell in your Google Colab notebook, a chart appears showing predicted traffic values tracking right alongside the actual values, proving your model learned to anticipate real traffic flow patterns.

By the end of this project, you'll have:

  • A YOLOv8 vehicle detection pipeline that processes real traffic camera frames, drawing bounding boxes around every car, motorcycle, bus, and truck it finds.
  • A PCU-weighted traffic flow chart showing intensity rising and falling across hundreds of frames, where buses and trucks carry more weight than cars.
  • A trained GRU model that predicts future traffic values from a sliding window of history, with a chart showing predictions tracking alongside actual values.
  • An optional Secret Mission to push your skills further.

Are there any prerequisites?

This project assumes basic Python knowledge and some familiarity with deep learning concepts like training loops and loss functions. You will also need a Google account for Colab access and a Kaggle account with an API token to download the dataset.

Before We Start

Before we dive in, let's lock in what you're about to build and why it matters.

Set Up Colab and Load UA-DETRAC from Kaggle

Before you can detect a single vehicle, three things need to be in place: a GPU to run the detection model quickly, the YOLOv8 detection library, and real traffic video frames to process.

In this step, you will enable the T4 GPU in Google Colab, install the ultralytics package, authenticate with the Kaggle API, and download real traffic camera frames from the UA-DETRAC dataset.

In this step, get ready to:
  • Enable the T4 GPU and install the YOLOv8 detection library.
  • Authenticate Kaggle and download UA-DETRAC traffic data.
  • Load and display sample traffic frames from the dataset.
Enable the T4 GPU and install ultralytics

Object detection models like YOLOv8 process hundreds of frames. A GPU speeds this up from minutes to seconds.

  • In your Colab notebook, click Runtime in the top menu bar.
  • Click Change runtime type.
  • Under Hardware accelerator, select T4 GPU.
  • Click Save. Colab will reconnect to a GPU-enabled runtime.
  • Install the ultralytics package (which provides YOLOv8) by running this command in a new code cell:
!pip install ultralytics==8.4.126

What does this command do?

The ! prefix tells Colab to run this as a terminal command instead of Python code. ultralytics==8.4.126 pins the exact version so your results are reproducible.

  • Wait for the install to finish. You should see a success message ending with the installed package version.
  • Now import the libraries and verify your GPU by pasting this code into a new code cell:
import os
import glob
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from sklearn.preprocessing import MinMaxScaler
from ultralytics import YOLO
import cv2
from google.colab.patches import cv2_imshow

# Verify GPU
print("CUDA available:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("GPU:", torch.cuda.get_device_name(0))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

What does this code do?

  • The import block loads everything the full pipeline needs: PyTorch for the neural network, OpenCV for reading frames, NumPy for arrays, and YOLO for object detection.
  • torch.cuda.is_available() checks whether the T4 GPU runtime you just enabled is accessible to PyTorch.
  • device is set to cuda so all model training later runs on the GPU automatically.
  • Run the cell. You should see CUDA available: True and your GPU name printed (e.g. Tesla T4).

Seeing CUDA available: False?

  • Go back to Runtime > Change runtime type and confirm T4 GPU is selected. Click Save again if it wasn't.
  • If the runtime disconnected, click Connect in the top-right corner, then re-run the cell.

Why is my GPU not showing up in Colab?

Authenticate Kaggle and download UA-DETRAC

The UA-DETRAC dataset contains real traffic camera footage captured at road intersections. It is hosted on Kaggle, so you need to authenticate with the Kaggle API before you can download it.

What is kaggle.json?

Your kaggle.json file is an API token that proves your identity to Kaggle's servers. You can download it from your Kaggle account by going to Settings > API > Create New API Token.

  • Upload your kaggle.json and move it into position by pasting this code into a new code cell:
from google.colab import files

files.upload()  # Upload kaggle.json when prompted

os.makedirs(os.path.expanduser("~/.kaggle"), exist_ok=True)
os.rename("kaggle.json", os.path.expanduser("~/.kaggle/kaggle.json"))
os.chmod(os.path.expanduser("~/.kaggle/kaggle.json"), 0o600)

What does this code do?

  • files.upload() opens a file picker in your browser. You will see a Choose Files button. Select your kaggle.json file from your computer.
  • The remaining lines move the file to ~/.kaggle/kaggle.json and set permissions to 0o600 (owner-only read/write). The Kaggle CLI looks for this exact path.
  • Run the cell. When the Choose Files button appears, click it and select your kaggle.json file.
  • You should see a confirmation showing the uploaded filename and size.

Getting a FileNotFoundError?

  • Make sure you uploaded a file named exactly kaggle.json. If you renamed it or uploaded the wrong file, re-run the cell and upload again.
  • If you do not have a kaggle.json file yet, go to kaggle.com, sign in, open Settings, scroll to API, and click Create New API Token.

Help me set up Kaggle API authentication in Colab.

Now download the UA-DETRAC dataset. The full dataset is 9.89 GB, so this will take a few minutes.

  • Download the dataset by running this command in a new code cell:
!kaggle datasets download -d dtrnngc/ua-detrac-dataset

What is UA-DETRAC?

UA-DETRAC is a benchmark dataset for traffic analysis. It contains 100 video sequences captured at real intersections, stored as 960x540 JPEG frames at 25 fps. The full dataset has over 140,000 frames with 1.21 million bounding box annotations.

You are downloading the entire zip, but in the next chunk you will extract only one sequence folder to keep processing fast.

  • Wait for the download to complete. You should see a progress bar followed by a message confirming the download finished.
Extract and load the traffic frames

The downloaded zip contains all 100 sequences, but you only need one. Each sequence folder (named like MVI_XXXXX) holds sequential JPEG frames from a single traffic camera.

  • Extract the first sequence folder from the zip by pasting this code into a new code cell:
import zipfile

DATASET_ZIP = "ua-detrac-dataset.zip"
EXTRACT_DIR = "ua_detrac_data"
os.makedirs(EXTRACT_DIR, exist_ok=True)

with zipfile.ZipFile(DATASET_ZIP, "r") as zf:
    all_names = zf.namelist()
    seq_folders = sorted(
        set(
            name.split("/")[1]
            for name in all_names
            if len(name.split("/")) > 2 and name.split("/")[1].startswith("MVI_")
        )
    )
    print(f"Found {len(seq_folders)} sequences. Using first sequence: {seq_folders[0]}")

    target_seq = seq_folders[0]
    for name in all_names:
        parts = name.split("/")
        if len(parts) > 1 and parts[1] == target_seq:
            zf.extract(name, EXTRACT_DIR)

What does this code do?

  • It scans the zip file's contents to find all folders starting with MVI_. These are the individual video sequences.
  • It picks the first sequence alphabetically and extracts only that folder's frames into ua_detrac_data/. This keeps storage and processing time manageable.
  • Run the cell. You should see a message like Found 100 sequences. Using first sequence: MVI_20011 (the exact sequence name may differ).

With the frames extracted, load them into memory and display one to confirm everything is working.

  • Load the extracted frames and display a sample by pasting this code into a new code cell:
frame_dir_candidates = glob.glob(
    os.path.join(EXTRACT_DIR, "**", target_seq), recursive=True
)
FRAME_DIR = frame_dir_candidates[0] if frame_dir_candidates else EXTRACT_DIR

frame_paths = sorted(glob.glob(os.path.join(FRAME_DIR, "*.jpg")))
print(f"Loaded {len(frame_paths)} frames from sequence {target_seq}")

frames = [cv2.imread(p) for p in frame_paths]

if frames:
    print(f"Frame shape: {frames[0].shape}")
    cv2_imshow(frames[0])

What does this code do?

  • glob.glob() finds the extracted sequence folder, then collects all .jpg frame paths and sorts them by filename. Sorting preserves the temporal order of the video.
  • cv2.imread() loads each JPEG into a NumPy array. The frames list holds every frame in memory, ready for YOLOv8 in the next step.
  • cv2_imshow() is Colab's version of OpenCV's display function. It renders the image inline in the notebook.

Before you run the cell, what do you expect to see? Think about what a traffic camera frame would look like.

  • Run the cell. You should see a count like Loaded 590 frames (the exact number depends on the sequence), a frame shape of (540, 960, 3), and a traffic camera image showing a road with vehicles.

Not seeing any frames or getting an empty list?

  • Make sure the Kaggle download completed successfully in the previous cell. Re-run the download command if it was interrupted.
  • Check that the extraction cell printed a sequence name. If it shows an error about the zip file, the download may have been incomplete.
  • Verify the extracted files exist by running !ls ua_detrac_data/ in a new cell. You should see a folder structure containing JPEG files.

Help me load UA-DETRAC frames in Colab.

✔️ Awesome, I've got everything!

Great. Double check you have all your code cells saved and that each one ran without errors.

ⓧ I'd like to double check the full code

Scroll through your notebook and confirm you have these cells, in order:

  • The !pip install ultralytics==8.4.126 cell.
  • The imports and GPU verification cell (starts with import os, ends with the device assignment).
  • The Kaggle authentication cell (starts with from google.colab import files, ends with os.chmod).
  • The !kaggle datasets download cell.
  • The extraction cell (starts with import zipfile, ends with the zf.extract loop).
  • The frame loading cell (starts with frame_dir_candidates, ends with cv2_imshow(frames[0])).

Your Colab environment is fully loaded with a GPU, the detection library, and real traffic frames ready to process. Next up, you will feed these frames through YOLOv8 to detect and count every vehicle in every frame.

Detect and Count Vehicles with YOLOv8

Your Colab notebook has hundreds of traffic camera frames loaded, but right now they are just raw pixels. To build a traffic flow time series, you need structured data: how many cars, buses, trucks, and motorcycles appear in each frame.

In this step, you will run YOLOv8 (You Only Look Once) on every frame to detect vehicles. YOLOv8 processes each image in a single forward pass, returning bounding boxes and class labels for every object it finds. You will filter for vehicle classes and collect per-frame counts.

In this step, get ready to:
  • Load the pretrained YOLOv8n model and run inference on every frame.
  • Filter detections to keep only vehicle classes (car, motorcycle, bus, truck) and count each type per frame.
  • Visualize detection results with bounding boxes drawn on a sample frame.
Load YOLOv8n and detect vehicles across all frames

YOLOv8 is pretrained on the COCO dataset, which includes 80 object classes. Only four of those classes are vehicles: car (class ID 2), motorcycle (3), bus (5), and truck (7). You will map these IDs to readable names and filter out everything else.

  • Create a new code cell in your notebook. Define the vehicle class mapping and run detection on every frame by pasting this code:
# --- Cell 3: YOLOv8 Vehicle Detection ---
# Map COCO class IDs to vehicle type names
VEHICLE_CLASSES = {
    2: "car",
    3: "motorcycle",
    5: "bus",
    7: "truck",
}

# Load pretrained YOLOv8 nano model
model = YOLO("yolov8n.pt")

# Collect per-frame vehicle counts
all_counts = []

for idx, frame in enumerate(frames):
    results = model(frame, verbose=False)
    counts = {"car": 0, "motorcycle": 0, "bus": 0, "truck": 0}

    for result in results:
        for box in result.boxes:
            cls_id = int(box.cls[0])
            if cls_id in VEHICLE_CLASSES:
                counts[VEHICLE_CLASSES[cls_id]] += 1

    all_counts.append(counts)
    if idx % 100 == 0:
        print(f"Frame {idx}/{len(frames)}: {counts}")

print(f"\nProcessed {len(all_counts)} frames")
print("Sample counts:", all_counts[:5])

What does this code do?

VEHICLE_CLASSES maps the four COCO integer class IDs to human-readable names. This dictionary acts as a filter: any detection whose class ID is not in this map gets ignored.

YOLO("yolov8n.pt") loads the YOLOv8 "nano" variant, the smallest and fastest model. It downloads the pretrained weights automatically on first run.

The loop runs inference on every frame with verbose=False to suppress per-frame logging. For each detected bounding box, it checks if the class ID is a vehicle, and if so, increments the count. A progress print fires every 100 frames so you know it is working.

  • Run the cell by pressing Shift+Enter.

This cell takes a few minutes

YOLOv8 is running inference on every single frame in your sequence. You will see progress prints every 100 frames. On a T4 GPU, expect this to take 2-5 minutes depending on the sequence length.

When the cell finishes, you should see a summary line showing the total number of processed frames and the first five frame counts. Each count is a dictionary with keys car, motorcycle, bus, and truck.

Seeing a CUDA out of memory error?

  • Make sure you enabled the T4 GPU runtime in Step 1. Go to Runtime > Change runtime type and confirm T4 GPU is selected.
  • If you still hit memory issues, restart the runtime (Runtime > Restart session) and re-run all cells from the top.

Why is YOLOv8 not detecting any vehicles in my frames?

Visualize detections with bounding boxes

Counts alone do not tell you whether the detections are accurate. Drawing bounding boxes on a sample frame lets you visually confirm that YOLOv8 is correctly identifying vehicles in the scene.

  • In the same cell, add the following code at the bottom to annotate and display the first frame:
# Run detection on the first frame and display with bounding boxes
results_vis = model(frames[0])
annotated = results_vis[0].plot()
cv2_imshow(annotated)

What does this code do?

results_vis[0].plot() returns a NumPy array of the frame with bounding boxes, class labels, and confidence scores drawn directly on it. cv2_imshow is Colab's version of OpenCV's display function, rendering the annotated image inline in your notebook.

Before you re-run the cell, think about what you expect to see. Will every object in the frame have a bounding box, or only the vehicles?

  • Re-run the cell by pressing Shift+Enter.

Below the frame counts, you should now see the first frame with colored bounding boxes around detected objects. YOLOv8 detects all 80 COCO classes, so you may see boxes on non-vehicle objects too (like people). The boxes include class labels and confidence scores at the top of each rectangle.

Why does the annotated frame show non-vehicle boxes?

The .plot() method draws all detections, not just vehicles. Your counting code in the loop above already filters by VEHICLE_CLASSES, so the vehicle counts are correct. The visualization is just showing you the full picture of what YOLOv8 sees.

Not seeing an annotated image?

  • Make sure you added the visualization code at the bottom of the same cell, after the detection loop. If you placed it in a separate cell, the model and frames variables should still be available, but confirm by running both cells in order.
  • Check that you are using cv2_imshow (from google.colab.patches), not cv2.imshow. The standard OpenCV display function does not work in Colab.

My annotated frame is not displaying in Colab. How do I fix this?

✔️ Awesome, I've got everything!

Great. Confirm you have saved your notebook. Your Cell 3 should contain both the detection loop and the visualization code.

ⓧ I'd like to double check the full code

Scroll up through your Cell 3 and confirm these two sections are present in order:

  • The VEHICLE_CLASSES dictionary, model = YOLO("yolov8n.pt"), and the detection loop that populates all_counts with per-frame vehicle counts, ending with two print statements.
  • The visualization block: results_vis = model(frames[0]), annotated = results_vis[0].plot(), and cv2_imshow(annotated).

You have turned raw traffic camera pixels into structured per-frame vehicle counts. Next up, you will convert these counts into a meaningful traffic flow signal using Passenger Car Unit (PCU) weighting, and shape it into a time series ready for the GRU model.

Build the PCU-Weighted Traffic Time Series

Your all_counts list now holds raw vehicle counts for every frame. You know exactly how many cars, motorcycles, buses, and trucks YOLOv8 found in each one.

Those raw counts treat every vehicle the same. In reality, a single bus takes up far more road space and slows traffic much more than a motorcycle. In this step, you will convert raw counts into Passenger Car Unit (PCU) weighted values, visualize the resulting time series, and shape the data into sliding window sequences for the GRU model.

In this step, get ready to:
  • Define PCU weights and compute weighted traffic values per frame.
  • Normalize and plot the traffic flow time series.
  • Create sliding window train/test sequences as PyTorch tensors.
Define PCU weights and compute traffic values

Traffic engineers use a measure called Passenger Car Unit (PCU) to reflect how much road capacity different vehicle types consume. A car is the baseline at 1.0, while a bus might be 3.0 because it occupies roughly three times the road space.

What is PCU?

PCU is a standard traffic engineering unit that converts mixed vehicle counts into a single comparable value. Instead of saying "5 cars and 1 bus", you say "8.0 PCU" (5 x 1.0 + 1 x 3.0). This gives a much more realistic picture of how congested a road actually is.

  • Create a new code cell in your notebook and paste the following to define the PCU weights and compute weighted values for every frame:
# --- Cell 4: PCU-Weighted Traffic Time Series ---
PCU_WEIGHTS = {
    "car": 1.0,
    "motorcycle": 0.5,
    "bus": 3.0,
    "truck": 2.5,
}

# Compute weighted traffic value for each frame
pcu_values = []
for counts in all_counts:
    pcu = sum(counts[vtype] * PCU_WEIGHTS[vtype] for vtype in PCU_WEIGHTS)
    pcu_values.append(pcu)

pcu_array = np.array(pcu_values, dtype=np.float32).reshape(-1, 1)

What does this code do?

  • PCU_WEIGHTS maps each vehicle type to its road-impact weight. A motorcycle (0.5) counts for half a car, while a bus (3.0) counts for three.
  • The loop walks through every frame's count dict from all_counts and multiplies each vehicle count by its PCU weight, summing them into a single traffic intensity value per frame.
  • pcu_array reshapes the list into a NumPy column vector (N, 1) so MinMaxScaler can fit it in the next substep.
  • Run the cell by pressing Shift+Enter.
  • Verify the PCU values by running this in a new cell below:
print(pcu_array[:5])

You should see five floating-point numbers. Frames with buses or trucks will show higher values than frames with only cars or motorcycles.

Seeing all zeros?

  • Make sure you ran Cell 3 (the YOLOv8 detection cell) before this cell. The all_counts list must already be populated.
  • Check that your PCU_WEIGHTS keys exactly match the strings in VEHICLE_CLASSES: "car", "motorcycle", "bus", "truck".

Why are my PCU values all zero even though all_counts has data?

Normalize and plot the time series

Neural networks learn more effectively when input values are scaled to a small, consistent range. MinMaxScaler from scikit-learn squeezes the PCU values into the 0-to-1 range while preserving the shape of the signal.

  • Create a new code cell and paste the following to normalize the time series and plot both the raw and normalized versions:
# Normalize with MinMaxScaler
scaler = MinMaxScaler()
pcu_scaled = scaler.fit_transform(pcu_array).flatten()

# Plot raw and normalized time series
fig, axes = plt.subplots(2, 1, figsize=(12, 6), sharex=True)
axes[0].plot(pcu_array, color="steelblue")
axes[0].set_ylabel("PCU Value")
axes[0].set_title("Raw PCU-Weighted Traffic Flow")
axes[1].plot(pcu_scaled, color="coral")
axes[1].set_ylabel("Normalized")
axes[1].set_xlabel("Frame Index")
axes[1].set_title("Normalized Traffic Flow")
plt.tight_layout()
plt.show()

What does this code do?

  • scaler.fit_transform(pcu_array) learns the min and max of the PCU values and rescales every value into the 0-to-1 range. The scaler object remembers these bounds so you can convert predictions back to the original scale later.
  • .flatten() converts the column vector back to a 1D array, which is easier to slice into sequences.
  • The two-panel plot shows the raw PCU values on top and the normalized version on the bottom, sharing the same x-axis so you can compare them directly.

Before you run this cell, what do you think the traffic values will look like over time? Smooth and steady, or fluctuating?

  • Run the cell by pressing Shift+Enter.

You should see a two-panel plot. The top panel shows raw PCU values fluctuating as vehicles enter and leave the scene. Frames with buses or trucks spike noticeably higher than frames with only cars. The bottom panel shows the same pattern squeezed into the 0-to-1 range.

Plot looks flat or empty?

  • A flat line at zero means pcu_array is all zeros. Go back and confirm the PCU weights cell ran successfully with non-zero output.
  • If you see a NameError for pcu_array, the previous cell did not run. Run it first, then re-run this cell.

My PCU time series plot looks wrong or empty. Help me debug.

Create sliding window sequences for the GRU

The GRU needs structured input: a window of recent traffic values as input, and the very next value as the target to predict. A sliding window of size 10 means the model sees 10 consecutive frames of traffic data and learns to forecast the 11th.

What is a sliding window?

Imagine placing a magnifying glass over 10 frames of your time series. The values inside the glass become one training sample's input (X), and the value just past the glass becomes the target (y). Then you slide the glass forward by one frame and create the next sample. This produces hundreds of overlapping (X, y) pairs from a single time series.

  • Create a new code cell and paste the following to define the sliding window function and generate the sequences:
# Create sliding window sequences
WINDOW_SIZE = 10


def create_sequences(data, window_size):
    X, y = [], []
    for i in range(len(data) - window_size):
        X.append(data[i : i + window_size])
        y.append(data[i + window_size])
    return np.array(X), np.array(y)


X, y = create_sequences(pcu_scaled, WINDOW_SIZE)

What does this code do?

  • WINDOW_SIZE = 10 sets how many past frames the model looks at to make each prediction.
  • create_sequences() slides through the entire normalized time series, collecting each window of 10 values as an input sample and the immediately following value as its target.
  • X holds all the input windows and y holds all the corresponding next-step target values.
  • Run the cell, then verify the array shapes by running this in a new cell below:
print(f"X shape: {X.shape}, y shape: {y.shape}")

You should see X shape: (N, 10) and y shape: (N,), where N is the number of frames minus 10. Each row of X is one 10-frame window, and each entry in y is the next value to predict.

Now you need to split this data into training and test sets, then convert everything into PyTorch tensors that the GRU can consume.

  • Create a new code cell and paste the following to split the data 80/20 and convert to tensors:
# Train/test split (80/20)
split_idx = int(len(X) * 0.8)
X_train, X_test = X[:split_idx], X[split_idx:]
y_train, y_test = y[:split_idx], y[split_idx:]

# Convert to PyTorch tensors: (batch, seq_len, features)
X_train_t = torch.FloatTensor(X_train).unsqueeze(-1).to(device)
y_train_t = torch.FloatTensor(y_train).unsqueeze(-1).to(device)
X_test_t = torch.FloatTensor(X_test).unsqueeze(-1).to(device)
y_test_t = torch.FloatTensor(y_test).unsqueeze(-1).to(device)

print(f"Training samples: {X_train_t.shape[0]}, Test samples: {X_test_t.shape[0]}")
print(f"Input shape: {X_train_t.shape}")

What does this code do?

  • split_idx takes the first 80% of samples for training and reserves the last 20% for testing. Time series data is split sequentially (not shuffled) to preserve temporal order.
  • .unsqueeze(-1) adds a feature dimension, reshaping from (batch, 10) to (batch, 10, 1). The GRU expects input shaped as (batch, sequence length, number of features).
  • .to(device) moves the tensors to your T4 GPU so training runs fast.

Before you run this cell, how many training samples do you expect? Think about how many total frames you have, subtract the window size to get total samples, then take 80% of that.

  • Run the cell by pressing Shift+Enter.

You should see two lines printed. The first shows the number of training and test samples. The second confirms the input tensor shape is (batch, 10, 1), which is exactly what the GRU expects: batches of 10-step sequences, each with 1 feature.

Seeing a CUDA error?

  • If you see a RuntimeError about CUDA, go back to Cell 1 and confirm torch.cuda.is_available() prints True. If it prints False, re-enable the T4 GPU runtime under Runtime > Change runtime type.
  • If you see a NameError for X or y, the previous cell (sliding window) did not run. Run it first, then re-run this cell.

Help me fix a CUDA or tensor shape error when converting my sliding window data to PyTorch tensors.

Why split time series sequentially instead of randomly?

Shuffling would leak future information into the training set. The model would see traffic patterns from later frames during training and get unrealistically good scores. Sequential splitting preserves the real-world constraint: you can only use past data to predict the future.

✔️ Awesome, I've got everything!

Great. Double check you have all four cells from this step saved in your notebook.

ⓧ I'd like to double check the full code

The full Cell 4 is spread across the chunks you built above. Scroll up and confirm these blocks are in your notebook cells, in order:

  • PCU_WEIGHTS dictionary, the pcu_values loop, and pcu_array creation.
  • MinMaxScaler normalization and the two-panel matplotlib plot.
  • WINDOW_SIZE, the create_sequences() function definition, and the call that produces X and y.
  • Train/test split with split_idx, PyTorch tensor conversion (X_train_t, y_train_t, X_test_t, y_test_t), and the two print statements.

Your traffic data is shaped and ready for deep learning. Next up, you will build a GRU neural network, train it on these sequences, and watch it learn to predict future traffic flow.

Build, Train, and Evaluate the GRU Model

You now have a clean, normalized traffic time series sliced into sliding windows and loaded into PyTorch tensors. The data is ready, but you have no way to predict what happens next in the traffic flow.

In this step, you will build a GRU (Gated Recurrent Unit) neural network that learns temporal patterns from the PCU-weighted traffic data and forecasts future traffic intensity.

In this step, get ready to:
  • Define a GRU model class that takes a window of traffic values and predicts the next one.
  • Train the model on the training set and visualize the loss curve.
  • Evaluate the trained model on the test set and plot predicted vs actual traffic values.
Define the TrafficGRU model class

A GRU is a type of recurrent neural network that processes sequences step by step, maintaining a hidden state that carries information from earlier time steps. It is lighter and faster to train than an LSTM because it uses two gates instead of three.

  • Create a new code cell in your notebook and define the TrafficGRU class by pasting this code:
# --- Cell 5: GRU Model Definition, Training, and Evaluation ---
class TrafficGRU(nn.Module):
    def __init__(self, input_size=1, hidden_size=64, num_layers=1):
        super(TrafficGRU, self).__init__()
        self.gru = nn.GRU(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
        )
        self.fc = nn.Linear(hidden_size, 1)

    def forward(self, x):
        # x shape: (batch, seq_len, input_size)
        gru_out, h_n = self.gru(x)
        # Use the last time step output
        last_output = gru_out[:, -1, :]  # (batch, hidden_size)
        prediction = self.fc(last_output)  # (batch, 1)
        return prediction

What does this code do?

  • nn.GRU creates a GRU layer with input_size=1 (one feature: PCU value), hidden_size=64 (64 memory units), and batch_first=True so input tensors follow the (batch, sequence, features) shape.
  • nn.Linear(hidden_size, 1) maps the 64-dimensional hidden state down to a single output: the predicted PCU traffic value.
  • The forward method passes the input through the GRU, takes only the last time step's output, and feeds it into the linear layer to produce one prediction per sample.
  • Run the cell. If no errors appear, the class is defined and ready to use.

Seeing a NameError?

  • Make sure you ran Cell 1 (the imports cell) earlier in this session. nn comes from import torch.nn as nn in that cell.
  • Check that the class name is TrafficGRU with exact capitalization.

Why isn't my TrafficGRU class definition running?

Why GRU instead of LSTM?

GRU uses two gates (reset and update) instead of LSTM's three (input, forget, output). This means fewer parameters and faster training. For short sequences like our 10-step traffic window, GRU performs comparably to LSTM with less computational overhead.

Train the model on the traffic time series

With the model class defined, you can now instantiate it, set up a loss function and optimizer, and run the training loop.

  • In the same cell, paste the following code below the class definition:
gru_model = TrafficGRU(input_size=1, hidden_size=64, num_layers=1).to(device)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(gru_model.parameters(), lr=0.001)

# Training loop
NUM_EPOCHS = 50
train_losses = []

for epoch in range(NUM_EPOCHS):
    gru_model.train()
    optimizer.zero_grad()
    predictions = gru_model(X_train_t)
    loss = criterion(predictions, y_train_t)
    loss.backward()
    optimizer.step()
    train_losses.append(loss.item())

    if (epoch + 1) % 10 == 0:
        print(f"Epoch [{epoch+1}/{NUM_EPOCHS}], Loss: {loss.item():.6f}")

# Plot training loss
plt.figure(figsize=(8, 4))
plt.plot(train_losses, color="steelblue")
plt.xlabel("Epoch")
plt.ylabel("MSE Loss")
plt.title("GRU Training Loss")
plt.show()

What does this code do?

  • TrafficGRU(...).to(device) creates a model instance and moves it to the GPU.
  • nn.MSELoss() measures how far predictions are from actual values. torch.optim.Adam adjusts the model weights to minimize that error.
  • Each epoch runs one forward pass (predict), computes the loss, runs backpropagation, and updates the weights. The loss is printed every 10 epochs so you can watch the model improve.
  • The final plot shows the training loss curve. A curve that trends downward confirms the model is learning.
  • Run the cell.

You should see five lines of loss output (one for every 10 epochs), with the loss value decreasing over time. After the text output, a loss curve chart appears showing the downward trend.

Loss not decreasing?

  • Make sure X_train_t and y_train_t exist from Cell 4. If you restarted the runtime, re-run all previous cells first.
  • Check that device matches where your tensors live. Both the model and the data must be on the same device (GPU or CPU).
  • If loss stays flat at a high value, your data may not have enough variation. Confirm your PCU time series plot from Step 3 shows a fluctuating signal.

My GRU training loss is not going down. Help me debug.

Evaluate on the test set

The loss curve shows the model learned from the training data. The real test is whether it can predict traffic values it has never seen before.

Before you run the next cell, think about this: do you expect the predicted values to follow the actual traffic pattern closely, or to miss the peaks and valleys?

  • Add the evaluation code to the same cell, below the training loss plot:
# Evaluate on test set
gru_model.eval()
with torch.no_grad():
    test_preds = gru_model(X_test_t).cpu().numpy()
    test_actual = y_test_t.cpu().numpy()

# Inverse transform to original scale
test_preds_inv = scaler.inverse_transform(test_preds)
test_actual_inv = scaler.inverse_transform(test_actual)

# Compute RMSE
rmse = np.sqrt(np.mean((test_preds_inv - test_actual_inv) ** 2))
print(f"\nTest RMSE: {rmse:.2f} PCU")

# Plot predictions vs actual
plt.figure(figsize=(12, 5))
plt.plot(test_actual_inv, label="Actual", color="steelblue", linewidth=2)
plt.plot(test_preds_inv, label="Predicted", color="coral", linewidth=2, linestyle="--")
plt.xlabel("Test Time Step")
plt.ylabel("PCU Traffic Value")
plt.title("GRU Traffic Flow Prediction: Actual vs Predicted")
plt.legend()
plt.tight_layout()
plt.show()

What does this code do?

  • gru_model.eval() switches the model out of training mode. torch.no_grad() disables gradient tracking, which speeds up inference and saves memory.
  • scaler.inverse_transform() converts the normalized predictions back to original PCU values so the RMSE and the plot use real-world units.
  • RMSE (Root Mean Squared Error) measures the average prediction error in PCU units. The final chart overlays predicted values (coral dashed line) on actual values (blue solid line) so you can see where the model tracks well and where it drifts.
  • Run the cell again to execute the full Cell 5 (class definition, training, and evaluation together).

You should see the training loss output, then the loss curve chart, then a printed RMSE value in PCU units, and finally the prediction chart. The coral dashed line (predicted) should roughly follow the blue solid line (actual), showing the GRU has learned to anticipate traffic flow changes.

Predictions look like a flat line?

  • If the coral line is flat, the model may have collapsed to predicting the mean value. Check that your training loss actually decreased. If it did not, re-run from Cell 4 to confirm your tensors have the right shape.
  • Make sure WINDOW_SIZE is set to 10 and your sequence has enough frames. Very short sequences produce too few training samples for the GRU to learn from.

My GRU predictions are flat or don't match the actual values. Help me debug.

✔️ Awesome, I've got everything!

Great. Make sure you have saved your notebook.

ⓧ I'd like to double check the full code

Cell 5 is the final code cell in this step. Scroll up through your cell and confirm these sections are present, in order:

  • The TrafficGRU class definition with __init__ and forward methods.
  • Model instantiation (gru_model), criterion, and optimizer setup.
  • The training loop running for NUM_EPOCHS = 50 with loss printed every 10 epochs.
  • The training loss curve plot.
  • Evaluation block with torch.no_grad(), RMSE computation, and the predicted vs actual plot.

You just built a complete pipeline from raw traffic video frames all the way to a neural network that predicts future traffic intensity. Next up, there is an optional Secret Mission to push your prediction accuracy further.

Secret mission

Add Time-of-Day Features for Better Predictions

Your GRU model predicts traffic flow from a single feature: the PCU value. But real traffic is periodic. Rush hours repeat daily, and weekends look different from weekdays. Can you give the model a sense of time to improve its predictions?

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely on Google Colab's free tier, so there are no ongoing costs. Decide whether to keep your notebook and runtime active, disconnect to free up the GPU, or delete everything entirely.

Resources you used:

  • Google Colab notebook (stored in your Google Drive).
  • Active Colab GPU runtime (holds the downloaded UA-DETRAC dataset, installed packages, and YOLOv8 model weights).

Keep everything running

No action needed. Choose this if you plan to keep experimenting with the notebook.

  • Your Colab notebook stays in Google Drive and you can reopen it anytime.
  • The GPU runtime will auto-disconnect after a period of inactivity, which clears all downloaded data and installed packages from the runtime.
  • Reconnecting later will require re-running cells to reinstall packages, re-authenticate Kaggle, and re-download the dataset.

Pause - I'll come back to this later

Disconnect the runtime to free up the GPU for other users. Your notebook stays safely in Google Drive.

  • In your Colab notebook, click Runtime in the menu bar.
  • Select Disconnect and delete runtime.
  • This clears the UA-DETRAC dataset, installed packages, and model weights from the runtime storage.
  • When you come back, re-run the cells from the top to set everything up again.

Delete - I don't want to use this again

Remove everything: disconnect the runtime and delete the notebook from Google Drive.

  • In your Colab notebook, click Runtime in the menu bar, then select Disconnect and delete runtime. This clears all downloaded data, packages, and model weights.
  • Open Google Drive and find your Colab notebook file.
  • Right-click the notebook and select Move to trash.

Nice Work!

Nice Work!

You just built a complete traffic flow prediction system from real surveillance footage. From raw video frames to vehicle detection to time series forecasting, you now have hands-on experience with how intelligent transportation systems work under the hood.

You learned how to:

  • Build a YOLOv8 vehicle detection pipeline that processes real UA-DETRAC traffic frames from Kaggle and classifies vehicles into cars, motorcycles, buses, and trucks.
  • Engineer a PCU-weighted traffic flow time series that reflects the real-world road capacity impact of different vehicle types.
  • Train a GRU recurrent neural network to predict future traffic flow from historical patterns and visualize predicted vs actual values.
  • Secret Mission: Enhanced predictions by adding sinusoidal time-of-day features, demonstrating how domain-informed feature engineering improves ML models.

Ready to quiz yourself?