YOLOv8で学習したデータをncnnで利用する

by

YOLOv8で学習させたデータを、モバイル向けフレームワークncnn(https://github.com/Tencent/ncnn)で使えるように変換し、推定を実行する。スマートフォンで実行する前提で変換しているが、今回はpython(Google Colaboratory)のncnnで実行する。

ncnnに変換

ncnnで学習したデータを使用するには、ncnn用のフォーマットに変換する必要がある。YOLOv8では変換の方法が公式で紹介されているので参考にする(https://docs.ultralytics.com/ja/modes/export/)。公式を参考に以下のプログラムを実行した。学習させたモデルのパスは変換したいファイルに合わせて変更する。このプログラムを実行すると、/content/drive/MyDrive/yolov8/runs/detect/train2/weightsの中にbest_ncnn_modelというフォルダが作成され、その中に幾つかのファイルが作成される。作成されたファイルの中でmodel.ncnn.paramとmodel.ncnn.binがncnnで推定を実行するときに必要なファイルとなる。

Python
from ultralytics import YOLO

model = YOLO("/content/drive/MyDrive/yolov8/runs/detect/train2/weights/best.pt")
model.export(format='ncnn')

ncnnのインストール

AndroidやiOSでncnnを使うときは、ncnn公式のHow to buildを参考にして行う(https://github.com/Tencent/ncnn/wiki/how-to-build)。今回はGoogle Colaboratoryを使ってpythonで実行するのでpipでインストールする。

Python
!pip install ncnn

ncnnでの推定

ncnnを使って推定を実行する。まずは抵抗の推定を行う。ncnnを使った推定を行う手順は以下のようになる。

  1. モデルの読み込み
  2. 画像の読み込み
  3. ncnn入力用に画像を変換(画像サイズを640に、色の形式をBGRからRGBにする。RGBを0~255から0~1にする)
  4. 推定の実行(画像をin0に入力すると、結果がout0から得られるので、変数に格納する)
  5. 得られた結果から有効な結果を選定する。(得られる結果はスコアの低いものが含まれるため、閾値以上のものを有効とする。またnms処理を行う)
  6. 結果の表示

ここで特に苦労したのは、ncnnに入力する形式と出力される形式が初めはわからなかったことである。試行錯誤の結果、入力する画像をサイズ640のRGBの正規化されたデータに変換することが必要であることがわかった。そして、得られる結果の形式は

[x,x,x,…,x],[y,y,y,…,y],[w,w,w,…,w],[h,h,h,…,h],[cls1のスコア,cls1のスコア,cls1のスコア,…,cls1のスコア],[cls2のスコア,cls2のスコア,cls2のスコア,…,cls2のスコア],…

のようにリストごとにx,y,w,h,cls1,cls2,…の結果が格納される。ここでcls1,cls2はアノテーションを行った時のクラスの数得られる。つまり今回のように抵抗だけの１クラスであればcls1まで得られ、カラーコードのように11クラスで学習させていれば、cls1~cls11までの11クラスのスコアの結果が得られる。１クラスであれば、そのクラスのスコアのみを判定すればいいが、複数クラスの場合、最もスコアの高いクラスを推定結果とする。また、このリストの形式では扱いにくいので、転置の処理を行い

[x,y,w,h,cls1,cls2,…],[x,y,w,h,cls1,cls2,…],[x,y,w,h,cls1,cls2,…],…

とすることで、結果1つずつに対して有効かどうかの判定を行いやすくなる。

さらに、YOLOv8をそのまま使って推定を行ったときには、閾値処理やnms処理をやってくれたが、ncnnでは自分でやる必要がある。そのため、今回は閾値以上の結果を格納し、その結果をnms処理した。
nms処理ではclsとそのスコア順にソートし、同じクラスの結果を比較した。このとき重なりが大きくスコアの低い結果を削除し、残ったものを最終的な結果とした。

プログラム全体と実行結果の画像を以下に示す。一部を除いてほとんどの抵抗で推定が成功しているのがわかる。これは閾値処理や学習のパラメータを調整することで、改善できる。

Python
import cv2
import ncnn
import numpy as np
from google.colab.patches import cv2_imshow
import traceback

IMG_PATH="/content/drive/MyDrive/yolov8/data/originals/DSC_0178.JPG"
SAVE_PATH="/content/drive/MyDrive/yolov8/ncnn_result.JPG"
NCNN_PARAM = "/content/drive/MyDrive/yolov8/runs/detect/train2/weights/best_ncnn_model/model.ncnn.param"
NCNN_BIN = "/content/drive/MyDrive/yolov8/runs/detect/train2/weights/best_ncnn_model/model.ncnn.bin"
IMGSZ = 640
CLS_TH = 0.5
IOU_TH = 0.5

# nms処理をする関数
# bboxesの構造は [x1,y1,x2,y2,cls,score]
def nms(bboxes, iou_th):
  # 処理しやすいようにソート
  bboxes.sort(key=lambda x: (x[4], x[5]))
  result_list = []

  while(len(bboxes)>0):
    # 比較対象1
    base = bboxes.pop(-1)
    result_list.append(base)
    base_size = (base[2]-base[0])*(base[3]-base[1])
    i = len(bboxes)-1
    while (i>=0):
      # 比較対象2
      target = bboxes[i]
      if (target[4] != base[4]):
        # クラスが違う場合はスキップ
        break
      # IoU計算のための処理
      target_size = (target[2]-target[0])*(target[3]-target[1])

      ox1 = max(base[0],target[0])
      oy1 = max(base[1],target[1])
      ox2 = min(base[2],target[2])
      oy2 = min(base[3],target[3])
      w = max(0, (ox2-ox1))
      h = max(0,(oy2-oy1))
      overlap_size = w*h
      IoU = overlap_size / (base_size+target_size-overlap_size)
      # IoUが閾値以上なら削除
      if (IoU>iou_th):
        del bboxes[i]
      i-=1

  return result_list


# netの初期化
net = ncnn.Net()
# モデルの読み込み
net.load_param(NCNN_PARAM)
net.load_model(NCNN_BIN)
# extractorの生成
ex = net.create_extractor()

# 画像読み込み
img = cv2.imread(IMG_PATH)

# 読み込んだ画像をモデルに入力する型に変換(サイズと形式を変換する)
img_height, img_width = img.shape[:2]

mat_in = ncnn.Mat.from_pixels_resize(
    img,
    ncnn.Mat.PixelType.PIXEL_BGR2RGB,
    img_width,
    img_height,
    IMGSZ,IMGSZ
)

# 正規化
mean_vals = []
norm_vals = [1 / 255.0, 1 / 255.0, 1 / 255.0]
mat_in.substract_mean_normalize(mean_vals, norm_vals)

# 推定の実行
ex.input("in0", mat_in)
ret, mat_out = ex.extract("out0")
out = np.array(mat_out)

# 結果を扱いやすいように転置する
output_list = out.T
result_list = []
for output_data in output_list:
  # 結果はxywh形式で得られるのでxyxy形式に変換
  x1 = int((output_data[0]-output_data[2]/2.0)*img_width/640)
  y1 = int((output_data[1]-output_data[3]/2.0)*img_height/640)
  x2 = int((output_data[0]+output_data[2]/2.0)*img_width/640)
  y2 = int((output_data[1]+output_data[3]/2.0)*img_height/640)

  # 最もスコアが高いものを調べ、閾値以上なら有効とする
  cls = np.argmax(output_data[4:])
  if(output_data[4+cls]>CLS_TH):
    result_list.append([x1,y1,x2,y2,cls, output_data[4+cls]])

# nms処理
result_list = nms(result_list, IOU_TH)

# 結果を表示
for result_data in result_list:
  cv2.rectangle(img,(result_data[0],result_data[1]),(result_data[2],result_data[3]),(0,255,0),2)

cv2.imwrite(SAVE_PATH, img)
cv2_imshow(img)

ncnnでの抵抗値の推定

ncnnを使用して抵抗値の推定まで行うプログラムを作成した。主な内容はYOLOv8で学習データから推定を行うと同じである。
ncnnを使用して推定を行う場合には、切り取った画像を一度保存し、imreadで再び読み込んだ画像を使用しなければうまく推定ができなかった。また、切り取った画像ごとにループの処理を行なっているが、ループ内で毎回ncnnのcreate_extractorでextractorを生成しなければ、推定の結果が前回のループの結果と同じとなってしまっていた。これらの問題を解決することによって、ncnnでも推定を行うことができた。そのプログラムを以下に示す。このプログラムを実行した結果をプログラムの下に示す。結果を見ると抵抗値が推定できていることがわかる。ncnnで抵抗値の推定ができたのでモバイル端末での応用が可能となる。

Python
import cv2
import ncnn
import numpy as np
from google.colab.patches import cv2_imshow
import traceback

import time

IMG_PATH="/content/drive/MyDrive/yolov8/data/originals/DSC_0178.JPG"
SAVE_PATH="/content/drive/MyDrive/yolov8/ncnn_result.JPG"
RESISTANCE_PARAM = "/content/drive/MyDrive/yolov8/runs/detect/train2/weights/best_ncnn_model/model.ncnn.param"
RESISTANCE_BIN = "/content/drive/MyDrive/yolov8/runs/detect/train2/weights/best_ncnn_model/model.ncnn.bin"
COLOR_CODE_PARAM = "/content/drive/MyDrive/yolov8/runs/detect/train4/weights/best_ncnn_model/model.ncnn.param"
COLOR_CODE_BIN = "/content/drive/MyDrive/yolov8/runs/detect/train4/weights/best_ncnn_model/model.ncnn.bin"

IMGSZ = 640
CLS_TH = 0.5
IOU_TH = 0.5

SI_UNIT_LIST = [[1, ''],   #0
                [10, ''],   #1
                [0.1, 'k'], #2
                [1, 'k'],   #3
                [10, 'k'],  #4
                [0.1, 'M'], #5
                [1, 'M'],   #6
                [10, 'M'],  #7
                [0.1, 'G'], #8
                [1, 'G'],   #9
                [0.1, '']]  #10

# nms処理をする関数
# bboxesの構造は [x1,y1,x2,y2,cls,score]
def nms(bboxes, iou_th):
  # 処理しやすいようにソート
  bboxes.sort(key=lambda x: (x[4], x[5]))
  result_list = []

  while(len(bboxes)>0):
    # 比較対象1
    base = bboxes.pop(-1)
    result_list.append(base)
    base_size = (base[2]-base[0])*(base[3]-base[1])
    i = len(bboxes)-1
    while (i>=0):
      # 比較対象2
      target = bboxes[i]
      if (target[4] != base[4]):
        # クラスが違う場合はスキップ
        break
      # IoU計算のための処理
      target_size = (target[2]-target[0])*(target[3]-target[1])

      ox1 = max(base[0],target[0])
      oy1 = max(base[1],target[1])
      ox2 = min(base[2],target[2])
      oy2 = min(base[3],target[3])
      w = max(0, (ox2-ox1))
      h = max(0,(oy2-oy1))
      overlap_size = w*h
      IoU = overlap_size / (base_size+target_size-overlap_size)
      # IoUが閾値以上なら削除
      if (IoU>iou_th):
        del bboxes[i]
      i-=1

  return result_list

def predict(img, extractor, imgsz, cls_th, iou_th):
  # 読み込んだ画像をモデルに入力する型に変換(サイズと形式を変換する)
  img_height, img_width = img.shape[:2]
  mat_in = ncnn.Mat.from_pixels_resize(
      img,
      ncnn.Mat.PixelType.PIXEL_BGR2RGB,
      img_width,
      img_height,
      imgsz,imgsz
  )

  # 正規化
  mean_vals = []
  norm_vals = [1 / 255.0, 1 / 255.0, 1 / 255.0]
  mat_in.substract_mean_normalize(mean_vals, norm_vals)

  # 推定の実行
  extractor.input("in0", mat_in)
  ret, mat_out = extractor.extract("out0")
  out = np.array(mat_out)

  # 結果を扱いやすいように転置する
  output_list = out.T
  result_list = []
  for output_data in output_list:
    # 結果はxywh形式で得られるのでxyxy形式に変換
    x1 = int((output_data[0]-output_data[2]/2.0)*img_width/640)
    y1 = int((output_data[1]-output_data[3]/2.0)*img_height/640)
    x2 = int((output_data[0]+output_data[2]/2.0)*img_width/640)
    y2 = int((output_data[1]+output_data[3]/2.0)*img_height/640)

    # 最もスコアが高いものを調べ、閾値以上なら有効とする
    cls = np.argmax(output_data[4:])
    if(output_data[4+cls]>cls_th):
      result_list.append([x1,y1,x2,y2,cls, output_data[4+cls]])
  # nms処理
  return nms(result_list, iou_th)

def find_index(data, column, value):
  for i, row in enumerate(data):
    if row[column] == value:
      return i
  return None

# netの初期化
resistance_net = ncnn.Net()
# モデルの読み込み
resistance_net.load_param(RESISTANCE_PARAM)
resistance_net.load_model(RESISTANCE_BIN)
# extractorの生成
resistance_ex = resistance_net.create_extractor()

# 画像読み込み
img = cv2.imread(IMG_PATH)

# 抵抗の推定を実行
result_list = predict(img,resistance_ex, IMGSZ, CLS_TH, IOU_TH)

# カラーコードのnet初期化
color_code_net = ncnn.Net()
# モデルの読み込み
color_code_net.load_param(COLOR_CODE_PARAM)
color_code_net.load_model(COLOR_CODE_BIN)


# 抵抗画像の切り取りと抵抗値推定のループ
for resistance_result in result_list:
  left,top,right,bottom=resistance_result[0:4]

  # 切り取り
  roi_img = img[top:bottom,left:right]
  # 一時保存した画像でないとうまく推定できない
  cv2.imwrite("/content/drive/MyDrive/yolov8/tmp.JPG", roi_img)
  roi_img = cv2.imread("/content/drive/MyDrive/yolov8/tmp.JPG")
  
  # extractorの生成(ループの中で推定前に生成する必要がある)
  color_code_ex = color_code_net.create_extractor()
  # カラーコード推定実行
  color_code_results = predict(roi_img,color_code_ex,IMGSZ,CLS_TH,IOU_TH)

  error = True
  resistance_value = ""
  
  if len(color_code_results)==4: #4つのカラーコードを検出
    # 許容差(金のカラーコード)を含む場合そのインデックスを取得
    tolerance_index = find_index(color_code_results,4,10)
    if(tolerance_index!=None):
      # 許容差のカラーコードの中心座標を計算
      tolerance_center = np.array([(color_code_results[tolerance_index][0]+color_code_results[tolerance_index][2])/2,(color_code_results[tolerance_index][1]+color_code_results[tolerance_index][3])/2])
      #　許容差のカラーコードからの距離を計算
      distance_list = []
      for j, (c_left,c_top,c_right,c_bottom, cls, score) in enumerate(color_code_results):
        color_code_center = np.array([(c_left+c_right)/2, (c_top+c_bottom)/2])
        distance_list.append([np.linalg.norm(color_code_center-tolerance_center), cls])
      # 距離順にソート
      distance_list.sort(reverse=True)

      # 許容差以外に金のカラーコードを含む場合対策
      if(distance_list[1][1] == 10):
        distance_list[1], distance_list[2] = distance_list[2], distance_list[1]
      #抵抗値計算
      resistance_value_num = round((distance_list[0][1]*10+distance_list[1][1])*SI_UNIT_LIST[distance_list[2][1]][0],2)
      resistance_value = str(resistance_value_num) + SI_UNIT_LIST[distance_list[2][1]][1]
      error = False

  # 結果の描画
  if(error):
    cv2.rectangle(img,(left,top),(right,bottom),(0,0,255),2)
  else:
    cv2.rectangle(img,(left,top),(right,bottom),(0,255,0),2)
    cv2.putText(img,resistance_value,(left,top-30),cv2.FONT_HERSHEY_SIMPLEX,5,(0,255,0),10,cv2.LINE_AA)

cv2.imwrite(SAVE_PATH, img)
cv2_imshow(img)

コメントする