Предсказание траектории летящего объекта

В этой статье мы обсудим решение задачи предсказания координат летящего объекта. Представим, что вы хотите сделать ПВО против комаров. Зная координаты комара на нескольких кадрах видео, надо сказать, где он окажется на следующем кадре.
Или, скажем, вы пишите AI для браузерной игрушки и надо предсказывать, где игрок будет через секунду, чтобы стрелять с реалистичным упреждением.

Можно построить сложную модель учитывающую ветер, инерцию и всю физику объекта, а можно просто покидать данные в нейросетку и получить вполне сносный результат, который, оказывается, одинаково хорошо работает и для отслеживания комаров, дронов, птиц, самолётов и других активно маневрирующих объектов. Так вот, эта статья про моделирование полёта через нейросети для ленивых.

Коротко:

Предполагаем, что на достаточно коротком отрезке траектория полёта укладывается в кривую Безьe, по нескольким точкам догадываемся о характеристиках кривой и используем её для предсказаний.

1. Один процесс собирает датасет с координатами и обучает нейросеть
2. Второй процесс имея предыдущие координаты, собственно, делает предсказание, выдаёт координаты в будущем.

Для простоты считаем, что у нас есть 5 координат в прошлом: T-4, T-3, T-2, T-1, и T-0 и расстояние по времени между ними одинаковое. Если говорить строго, то в реальной жизни так не бывает, потому что кадры с видео приходят с разной задержкой, не на всех кадрах можно распознать объект и т.д., потому кроме координат по-хорошему надо обрабатывать и временную метку. Выносим это за пределы статьи и будем оперировать только с плоскими координатами.

Итак, считаем что у нас есть координаты пикселей для каждой из T.

Датасет в CSV выглядит вот так:

filename, x_t+1, y_t+1, x_t-0, y_t-0, x_t-1, y_t-1, x_t-2, y_t-2, x_t-3, y_t-3, x_t-4, y_t-4 

kxs2ut1j, -0.003441, -0.006477,  0.000000,  0.000000, -0.076425,  0.012893, -0.232717,  0.032203, -0.468875,  0.057928, -0.784900,  0.090070

Это координаты относительные к точке T-0. То есть мы считаем что T-0 всегда находится ровно в середине кадра с координатами 0,0; при этом -0.5, -0.5 это левый-верхний угол кадра, +0.5, +0.5 это правый-нижний угол кадра.

В качестве инструментов используем JavaScript и библиотеку brain.js

Объявим зависимости

const brain = require("brain.js");
const fs = require("fs");

Проинициализируем примитивную fully-connected нейросеть в один слой и три нейрона. Если есть пред-тренированные веса, загружаем их

const config = {
  binaryThresh: 0.5,
  hiddenLayers: [3], // array of ints for the sizes of the hidden layers in the network
  activation: "sigmoid", //supported activation types: ["sigmoid", "relu", "leaky-relu", "tanh"],
  learningRate: 0.000003, // scales with delta to effect training rate --> number between 0 and 1
  momentum: 0.01,
  log: true,
  logPeriod: 1,
  iterations: 1000, //number of iterations per epoch
  errorThresh: 0.0000005,
};

const net = new brain.NeuralNetwork(config);

//if there are weights saved from the previous run, fetch them
if (fs.existsSync("brain")) { 
  let json = JSON.parse(fs.readFileSync("brain"));
  net.fromJSON(json);
  console.log("Reused the weights from the brain trained last time");
}

Прочитаем наш датасет из csv файла

const data = fs.readFileSync("./dataset/data.csv", "UTF-8");

// split the contents by new line
let lines = data.split(/\r?\n/);

Создадим два массива с датасетом: в input пойдут координаты T-4… T-0, в output пойдут координаты T+1

let dataset = [];

lines.forEach((line) => {
  let raw = line.split(",");
  let rec = {
    input: [],
    output: []
  };

  for (let i = 1; i < raw.length; i++)
    raw[i] = (parseFloat(raw[i]) +1) / 2;

  rec.input.push(raw[3]); //present coordinates T-0
  rec.input.push(raw[4]);
  rec.input.push(raw[5]); //past coordinates T-1
  rec.input.push(raw[6]);
  rec.input.push(raw[7]); //past coordinates T-2
  rec.input.push(raw[8]);
  rec.input.push(raw[9]); //past coordinates T-3
  rec.input.push(raw[10]);
  rec.input.push(raw[11]); //past coordinates T-4
  rec.input.push(raw[12]);

  rec.output.push(raw[1]); //future coordinates T+1
  rec.output.push(raw[2]);

  dataset.push(rec);
});

Запускаем обучение. На моём лаптопе обучение в тысячу итераций с сотней датапоинтов занимает всего 38ms без использования GPU.
Для объектов летящих по баллистической траектории достаточно данных со ста кадров. Для активно маневрирующих объектов — чуть больше.

console.time("training");
net.train(dataset);
console.timeEnd("training");
fs.writeFileSync("brain", JSON.stringify(net.toJSON())); //save the weights

Результат налицо:

Тут зелёным показаны прошлые точки, жёлтым реальная координата в будущем, красным — предсказанная координата.

То есть за всего 38 несчастных миллисекунд и с помощью данных со ста кадров мы научили нейросетку предсказывать координату со вполне себе достаточной точностью. Если вы можете позволить себе потратить ещё 38мс то можно прогнать ещё цикл, предсказание станет ещё точнее.

Целью статьи было показать, что с помощью примитивнейшей нейросетки можно легко заменить тяжелое моделирование физики и получить приличную точность предсказаний будущих координат летящего объекта.

Весь код и синтетический датасет для иллюстрации, конечно, на гитхабе.