🧙‍♂️古希腊掌管《用 AI 识别乐高》的神

事情的起因是这样的，每年的结婚纪念日我们都有拼乐高的传统。

但是拼的时候经常都找不到对应的积木兼职，太痛苦了。😖

于是心生一计，想通过手机识别到我要找的积木，然后直接用框给我标出来，省时省力不费眼，岂不美哉。😎

恰巧之前写过浏览器上运行识别狗的一个功能 "🚫 为了防止狗上沙发，写了一个浏览器实时识别目标功能 📷"，想着拿来改造一下应该就行了。但是 coco-ssd 只能识别出日常的 80 多种物体。所以需要自己训练一个，或者找一个训练好的"识别积木模型"。 🤖

解决思路 📚：

1. 📷 调用手机摄像头，获取摄像头中的画面用 canvas 绘制

这里我使用的是p5.js 一个流行的 JavaScript 库，简化了视觉和交互体验的创建，提供了易于使用的 API 用于绘图、处理用户输入以及处理如视频等媒体。

其中 setup() 和 draw() 是内置函数 😊 不需要调用

js 复制代码

/**
 * 初始化函数，设置画布大小并配置视频捕获属性。
 * 该函数不接受参数，也不返回任何值。
 */
function setup() {
  // 创建画布并设置其尺寸
  canvas = createCanvas(640, 480);

  // 计算水平和垂直缩放因子，以保持捕获的视频与画布尺寸一致
  scaleX = 640 / +canvas.canvas.width || 640;
  scaleY = 480 / +canvas.canvas.height || 480;

  // 定义视频捕获的约束，指定使用后置摄像头
  let constraints = {
    video: {
      facingMode: "environment",
    },
  };
  // 创建视频元素并配置其大小，注册视频准备就绪的回调函数
  video = createCapture(constraints, videoReady);
  video.size(640, 480);
  video.hide(); // 隐藏视频元素，仅使用其捕获的视频数据
}

2. 🔍 加载对应的识别乐高的模型

原本想要使用ml5.js 但是发现需要自己再训练乐高的模型且训练速度很慢，限制很多，作罢 😕。

目前使用的是 roboflow.js 同样是基于 tensorFlow.js 但是社区中有很多的训练好可直接使用的模型。

这里模型配置可信值我降低到了 0.15 ,因为发现高可信值的模型识别率太低了 😏。

js 复制代码

/**
 * 异步函数 videoReady，初始化视频处理模型并准备就绪。
 */
async function videoReady() {
  console.log("videoReady");
  // 等待模型加载
  model = await getModel();

  // 配置模型的阈值
  model.configure({ threshold: 0.15 });
  // 更新UI，表示模型已准备好
  loadText.innerHTML = "modelReady";
  console.log("modelReady");
  // 选择要识别的目标
  processSelect();
  // 开始检测
  detect();
}

/**
 * 异步函数 getModel，从roboflow服务加载指定的模型。
 */
async function getModel() {
  return await roboflow
    .auth({
      publishable_key: "rf_lOpB8GQvE5Uwp6BAu66QfHyjPA13", // 使用API密钥进行授权
    })
    .load({
      model: "hex-lego", // 指定要加载的模型名称
      version: 3, // 指定要加载的模型版本
    });
}

3. 🧱 选择要识别的"目标积木"类型

绑定要选择的"目标积木"

js 复制代码

function processSelect() {
  const { classes } = model.getMetadata();
  console.log("classes", classes);
  classes.forEach((className) => {
    const option = document.createElement("option");
    option.value = className;
    option.text = className;
    selectRef.appendChild(option);
  });
}

4. 🔎 让该模型识别图像并返回出识别对象的信息

调用模型的 API 进行识别，以便于后续的绘制

js 复制代码

const detect = async () => {
  if (!play || !model) {
    console.log("model is not available");
    timer = setTimeout(() => {
      requestAnimationFrame(detect);
      clearTimeout(timer);
    }, 2000);
    return;
  }
  detections = await model.detect(canvas.canvas);
  console.log("detections", detections);
  requestAnimationFrame(detect);
};

5. 🎨 通过识别出的对象信息在 canvas 上进行绘制

获取到模型返回的信息保存并将识别到的信息都用 canvas 绘制出来

js 复制代码

function draw() {
  image(video, 0, 0);

  for (let i = 0; i < detections.length; i += 1) {
    const object = detections[i];
    let { x, y, width, height } = object.bbox;

    width *= scaleX;
    height *= scaleY;
    x = x * scaleX - width / 2;
    y = y * scaleY - height / 2;

    stroke(0, 0, 255);
    if (object.class.includes(selectVal)) stroke(0, 255, 0);
    strokeWeight(4);
    noFill();
    rect(Math.floor(x), Math.floor(y), Math.floor(width), Math.floor(height));
    noStroke();
    fill(255);
    textSize(24 * scaleX);
    text(object.class, Math.floor(x) + 10, Math.floor(y) + 24);
  }
}

6. 📲 部署在手机上

安装 termux
安装 python3
运行 python3 -m http.server 8000

最终代码

html 复制代码

<!DOCTYPE html>
<html lang="en">
  <head>
    <meta charset="UTF-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <title>find Lego</title>
    <script src="https://cdnjs.cloudflare.com/ajax/libs/p5.js/1.1.9/p5.min.js"></script>
    <script src="https://cdn.roboflow.com/0.2.26/roboflow.js"></script>
    <style></style>
  </head>
  <body>
    <div id="loadText">model is loading... please wait.</div>
    <select name="classify" id="selectRef"></select>
    <button id="start">start</button>
    <button id="stop">stop</button>
  </body>

  <script>
    let model;
    let video;
    let canvas;
    let play = true;
    let timer;
    let detections = [
      {
        bbox: {
          x: 0,
          y: 0,
          width: 100,
          height: 100,
        },
        class: "testing...",
      },
    ];

    let scaleX = 1;
    let scaleY = 1;

    let selectVal;
    const selectRef = document.getElementById("selectRef");
    selectRef.addEventListener("change", (e) => {
      selectVal = e.target.value;
      console.log("selectVal", e.target.value);
    });
    function setup() {
      canvas = createCanvas(640, 480);

      scaleX = 640 / +canvas.canvas.width || 640;
      scaleY = 480 / +canvas.canvas.height || 480;

      let constraints = {
        video: {
          facingMode: "environment",
        },
      };
      video = createCapture(constraints, videoReady);
      video.size(640, 480);
      video.hide();
    }
    function draw() {
      image(video, 0, 0);

      for (let i = 0; i < detections.length; i += 1) {
        const object = detections[i];
        let { x, y, width, height } = object.bbox;

        width *= scaleX;
        height *= scaleY;
        x = x * scaleX - width / 2;
        y = y * scaleY - height / 2;

        stroke(0, 0, 255);
        if (object.class.includes(selectVal)) stroke(0, 255, 0);
        strokeWeight(4);
        noFill();
        rect(
          Math.floor(x),
          Math.floor(y),
          Math.floor(width),
          Math.floor(height)
        );
        noStroke();
        fill(255);
        textSize(24 * scaleX);
        text(object.class, Math.floor(x) + 10, Math.floor(y) + 24);
      }
    }

    const loadText = document.getElementById("loadText");

    async function videoReady() {
      console.log("videoReady");
      model = await getModel();

      model.configure({ threshold: 0.15 });
      loadText.innerHTML = "modelReady";
      console.log("modelReady");
      processSelect();
      detect();
    }

    function processSelect() {
      const { classes } = model.getMetadata();
      console.log("classes", classes);
      classes.forEach((className) => {
        const option = document.createElement("option");
        option.value = className;
        option.text = className;
        selectRef.appendChild(option);
      });
    }

    async function getModel() {
      return await roboflow
        .auth({
          publishable_key: "rf_lOpB8GQvE5Uwp6BAu66QfHyjPA13",
        })
        .load({
          model: "hex-lego",
          version: 3, // <--- YOUR VERSION NUMBER
        });
    }

    const detect = async () => {
      if (!play || !model) {
        console.log("model is not available");
        timer = setTimeout(() => {
          requestAnimationFrame(detect);
          clearTimeout(timer);
        }, 2000);
        return;
      }
      detections = await model.detect(canvas.canvas);
      console.log("detections", detections);
      requestAnimationFrame(detect);
    };

    const stopBtn = document.getElementById("stop");
    stopBtn.addEventListener("click", () => {
      play = false;
      video.pause();
      // TODO
    });

    const startBtn = document.getElementById("start");
    start.addEventListener("click", () => {
      play = true;
      video.play();
    });
  </script>
</html>

总结 🎓

技术选型：

p5.js：作为基础库，简化了在网页上实现图形、视频处理和用户交互的过程。
roboflow.js：替代 ml5.js，提供了更丰富的预训练模型，包括特定于乐高的模型，降低了自行训练模型的门槛。

核心功能实现：

摄像头画面获取与显示：通过 createCapture 获取后置摄像头画面，并利用 createCanvas 创建画布实时显示视频流。
模型加载与配置：异步加载 roboflow 提供的乐高积木识别模型，并根据需求调整识别阈值以提高识别率。
目标选择与识别：动态生成下拉菜单供用户选择要识别的积木类型，根据用户选择调用模型进行实时识别。
结果绘制：将模型识别到的积木位置信息在 canvas 上以矩形框和文字形式标注出来，直观指示积木位置。
控制逻辑：添加了开始和停止按钮，允许用户控制视频流的播放与暂停，便于实际操作。

部署

介绍了如何在 Android 设备上的 Termux 应用中部署此项目，通过 Python 简易服务器让项目在本地网络中可访问，便于在手机上测试和使用。

改进方向 🚀

性能优化：针对移动设备的性能限制，可以考虑在模型推理阶段加入轻量化处理，比如降低视频帧率或使用更轻量级的模型版本，以减少计算资源消耗。
用户体验：增加用户引导和反馈机制，比如识别开始前的提示、识别过程中的加载动画，以及识别不到积木时的友好提示，提升整体用户体验。
离线支持：探索将模型文件本地化的方法，使应用在无网络环境下也能使用，但这可能需要解决模型文件较大和跨平台兼容性的问题。
模型精度提升：虽然降低识别阈值可以增加识别数量，但可能会引入更多误识别。可以通过收集自己的乐高数据集，对现有模型进行微调，以提高在特定场景下的识别精度。

写在最后 😅

大家如果还有什么好方法的话可以一起分享一下 😊

还没等摸鱼的时候写好功能，老婆已经拼完了。。。