TensorFlow.js (4) YOLOv5 Live demo

YOLOv5로 실시간으로 detection하는 방법을 공유해보도록 하겠습니다. YOLOv5는 ultralytics회사에서 주도적으로 개발하는 object detection model입니다. detection 성능은 현재 2022년까지 최상의 성능을 내고있습니다. 오늘은 YOLOv5모델 중 YOLOv5n를 사용해 live demo를 진행해보도록 하겠습니다.

해당 블로그에서 실시간으로 demo가 가능하므로 카메라를 요청할수 있어요! 카메라로 어떤 정보나 해킹은 없으니 안심하고 사용하세요! ㅠㅠ

1. YOLOv5n TensorFlow.js 변환

YOLOv5에서는 다음 명령어로 TensorFlow.js모델로 변환하는 코드를 간단하게 제공하고 있습니다.

git clone https://github.com/ultralytics/yolov5.git
python export.py --weights yolov5n.pt --include tfjs

변환이 잘되었다면 다음과 같이 파일들이 생성되었을 것입니다.

yolov5n fjs model — yolov5n tensorflow.js

.bin파일들은 모델의 weight를 저장하는 것이며 model.json은 모델의 구조를 담고 있는 파일입니다. 

변환 과정에 대해 조금 설명하겠습니다. 이후에 tensorflow.js코드에 필요한 정보라서요! export.py 코드에서는 yolov5n을 다음과 같은 변환을 거칩니다.

PyTorch -> TensorFlow -> TensorFlow.js

이때 TensorFlow로 변환될 때 agnosticNMS를 추가해주었기 때문에 output의 형태가 [boxes, scores, classes, valid_detections]가 됨을 아셔야하는데요. 그 이유는 TensorFlow.js로 변환하고 난 뒤 tensorflow.js 코드에서 yolov5n모델의 결과(output)를 웹상에 그려줘야하기 때문입니다! (output형태가 실제로 어떻게 변화하는 지 코드로 보고싶으면 여기로 -> TensorFlow변환시 output변화 )

output의 의미는 다음과 같습니다.

boxes: detect된 물체의 bounding box position (x1, y1, x2, y2)
- x1: bounding box의 왼쪽 x좌표
- y1: bounding box의 위 y좌표
- x2: bounding box의 오른쪽 x 좌표
- y2: bounding box의 아래 y좌표
scores: detect된 물체의 confidence score
classes: detect된 물체의 class index
valid_detections: nms를 통한 최종 detect된 물체 총 개수

2. YOLOv5n Live demo

이전글들에서 live demo를 해보았기때문에 해당 글에서는 YOLOv5모델을 돌리기 위한 중요한 코드만 설명드릴게요. 티스토리블로그에서 돌릴 수 있는 전체 코드는 다음 github tfjs_tutorial 에서 찾아 보시면 됩니다.

2.1 HTML skeleton

<body>
  <div id="main">
    <div class="container">
      <div class="canvas-wrapper">
        <canvas id="output"></canvas>
        <video id="video" playsinline style="
          -webkit-transform: scaleX(-1);
          transform: scaleX(-1);
          visibility: hidden;
          width: auto;
          height: auto;
          ">
        </video>
      </div>
    </div>
  </div>
</body>

<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@latest"></script>

html코드에 대한 내용입니다. video element를 통해 실시간으로 camera를 읽어오는 것이고 canvas element를 통해 model detection한 결과를 camera화면 위에 그릴 것입니다. 그리고 아래의 script element를 보면 tfjs를 cdn을 통해 import하여 tfjs의 모듈들을 사용함을 정의하였습니다.

2.2 javascript 기본구조

javascript 큰 구조는 다음과 같습니다. 아래의 함수들을 통해 live로 yolov5 detection을 진행할 것입니다.

async function app() {
  camera = await Camera.setupCamera(); //camera setup
  detector = await createDetector(); //load yolov5n model
  renderPrediction(); //draw detection result into canvas
};
app();

2.3 Load Model

이제부터 javascript코드를 작성할것입니다.

const yolov5n_weight = "https://raw.githubusercontent.com/da22so/tfjs_models/main/yolov5n_web_model/model.json"

async function createDetector() {
  return tf.loadGraphModel(yolov5n_weight);
}

위에서 만들어진 yolovn5의 model.json과 .bin파일들은 위의 github url에 올려놓은 것이며 tf.loadGraphModel을 통해 model을 load하게 됩니다.

2.4 Model input shape에 맞게 stream image 변환

실시간으로 들어오는 image stream을 yolov5n모델의 input shape에 맞춰 주기위해 다음과 같은 코드를 추가하였습니다.

  let [modelWidth, modelHeight] = detector.inputs[0].shape.slice(1, 3); //get model's input shape
  const input = tf.tidy(() => {
    return tf.image.resizeBilinear(tf.browser.fromPixels(camera.video), [modelWidth, modelHeight])
      .div(255.0).expandDims(0);
      // 실시간으로 들어오는 camera.video를 model input shape에 맞게 변환
      // normalize를 위해 255로 나눠줌
      // 3차원을 4차원으로 변환
  });

2.5 object detection from yolov5n

아래의 코드로 detection하게 됩니다.

      detect_res = await detector.executeAsync(input,); //detection!!

2.6 detection result를 canvas에 그리기

위의 detect_res의 값이 아래 res와 동일한데요. 위에서 tensorflow로 변환시에 output shape이 [boxes, scores, classes, valid_detections]이었음을 기억하면 아래의 코드가 이해되실거예여!

const [boxes, scores, classes, valid_detections] = res;
    const boxes_data = boxes.dataSync();
    const scores_data = scores.dataSync();
    const classes_data = classes.dataSync();
    const valid_detections_data = valid_detections.dataSync()[0];
    tf.dispose(res);
    var i;
    for (i = 0; i < valid_detections_data; ++i) { // valid_detections수만큼 물체 인식
      let [x1, y1, x2, y2] = boxes_data.slice(i * 4, (i + 1) * 4); //slicing을 통한 한 물체의 bounding box좌표 가져오기
	  ...// 생략
      const width = x2 - x1; 
      const height = y2 - y1;
      const klass = coco_names[classes_data[i]]; // class index를 coco class이름으로 매칭
      const score = scores_data[i].toFixed(2); 

      // bounding box 그리기
      this.ctx.strokeStyle = "#00FFFF";
      this.ctx.lineWidth = 4;
      this.ctx.strokeRect(x1, y1, width, height);

      // label과 confidence score 그리기
      this.ctx.fillStyle = "#00FFFF";
      const textWidth = this.ctx.measureText(klass + ":" + score).width;
      const textHeight = parseInt(font, 10); // base 10
      this.ctx.fillRect(x1, y1, textWidth + 4, textHeight + 4);
    }
	...//생략
    }

3. YOLOv5 Live demo 결과

webcam사용을 승낙하셧다면 아래에 detection결과가 짜짠!!!!!

저작자표시

'AI Engineering > TensorFlow' 카테고리의 다른 글

TFLite 뽀개기 (4) XNNPACK 이해 및 성능 비교 (0)	2022.08.10
TensorFlow.js (3) TensorFlow.js 변환 (0)	2022.04.03
TensorFlow.js (2) - WebGL 기반 hand pose detection (2)	2022.03.23
TensorFlow.js (1) - TensorFlow.js 이해 및 detection 예제 (1)	2022.03.17
Mediapipe (2) - custom segmentation model with mediapipe (0)	2022.03.14