大學生人工智慧挑戰賽—智慧零售 從目標檢測到演算法落地

語言: CN / TW / HK
時間 2022-05-14 03:16:11

theme: mk-cute

前言

這是最近的一次作業,試試本科競賽內容應該沒啥大問題吧正好水一篇部落格.

image-20220512221829274

首先下載資料集,提取碼 wwsj

檢視資料集給出的是json格式，訓練集和測試集有標註(共110張)，其餘還有無標註的需要自己手工標註。但是既然只是作業又不是去參加比賽，那就直接當小資料量樣本訓練.

image-20220512222354898

image-20220512222407851

構思

目前資料量較小,而且很明顯是一個目標檢測任務,並且涉及到演算法落地的問題,所以開始之前一定要理清思路,想清楚每一步應該怎麼做.

  • 找到一個合適的目標檢測模型,基於這個小樣本資料集進行訓練,得到一個效果較好的模型
  • 將python訓練得到的模型進行轉換,轉為onnx以及tensorRT等形式,方便後續演算法落地
  • 有了轉換後的模型,進行c++改寫模型載入以及檢測部分程式碼

既然是快速實現一次作業,那必然要“站在巨人的肩膀上”,所以使用的大部分都是網上的開原始碼.

開始動手

1.資料集準備

從網盤下載資料,因為資料量很小,所以train和test全部拿來訓練,一共110張圖片.然後把資料集轉為voc格式

```python

將所給的資料轉為voc資料集格式

import os import numpy as np import codecs import json from glob import glob import cv2 import shutil from sklearn.model_selection import train_test_split

1.標籤路徑

labelme_path = "./2022 年（第 15 屆）中國大學生計算機設計大賽人工智慧挑戰賽-智慧零售賽項資料集/TrainingDataset/" # 使用labelme打的標籤（包含每一張照片和對應json格式的標籤） saved_path = "./VOCdevkit/VOC2007/" # 儲存路徑

2.voc格式的資料夾，如果沒有，就建立一個

if not os.path.exists(saved_path + "Annotations"): os.makedirs(saved_path + "Annotations") if not os.path.exists(saved_path + "JPEGImages/"): os.makedirs(saved_path + "JPEGImages/") if not os.path.exists(saved_path + "ImageSets/Main/"): os.makedirs(saved_path + "ImageSets/Main/")

3.獲取json檔案

files = glob(labelme_path + "*.json") files = [i.split("/")[-1].split(".json")[0] for i in files] # 獲取每一個json檔名 print(len(files))

4.讀取每一張照片和對應標籤，生成xml

for json_file_ in files: json_filename = labelme_path + json_file_ + ".json" json_file = json.load(open(json_filename, "r", encoding="utf-8")) height, width, channels = cv2.imread(labelme_path + json_file_ + ".png").shape with codecs.open(saved_path + "Annotations/" + json_file_ + ".xml", "w", "utf-8") as xml: xml.write('\n') xml.write('\t' + 'UAV_data' + '\n') xml.write('\t' + json_file_ + ".png" + '\n') xml.write('\t\n') xml.write('\t\tThe UAV autolanding\n') xml.write('\t\tUAV AutoLanding\n') xml.write('\t\tflickr\n') xml.write('\t\tNULL\n') xml.write('\t\n') xml.write('\t\n') xml.write('\t\tNULL\n') xml.write('\t\tChaojieZhu\n') xml.write('\t\n') xml.write('\t\n') xml.write('\t\t' + str(width) + '\n') xml.write('\t\t' + str(height) + '\n') xml.write('\t\t' + str(channels) + '\n') xml.write('\t\n') xml.write('\t\t0\n') for multi in json_file["labels"]: #print(len(multi)) xmin=multi['x1'] xmax=multi['x2'] ymin=multi['y1'] ymax=multi['y2'] label = multi["name"]

        xml.write('\t<object>\n')
        xml.write('\t\t<name>' + label + '</name>\n')
        xml.write('\t\t<pose>Unspecified</pose>\n')
        xml.write('\t\t<truncated>1</truncated>\n')
        xml.write('\t\t<difficult>0</difficult>\n')
        xml.write('\t\t<bndbox>\n')
        xml.write('\t\t\t<xmin>' + str(xmin) + '</xmin>\n')
        xml.write('\t\t\t<ymin>' + str(ymin) + '</ymin>\n')
        xml.write('\t\t\t<xmax>' + str(xmax) + '</xmax>\n')
        xml.write('\t\t\t<ymax>' + str(ymax) + '</ymax>\n')
        xml.write('\t\t</bndbox>\n')
        xml.write('\t</object>\n')
        print(json_filename, xmin, ymin, xmax, ymax, label)
    xml.write('</annotation>')

5.複製圖片到 VOC2007/JPEGImages/下

image_files = glob(labelme_path + "*.png") print("copy image files to VOC007/JPEGImages/") for image in image_files: shutil.copy(image, saved_path + "JPEGImages/")

6.劃分train，test，val格式資料集

txtsavepath = saved_path + "ImageSets/Main/" ftrainval = open(txtsavepath + '/trainval.txt', 'w') ftest = open(txtsavepath + '/test.txt', 'w') ftrain = open(txtsavepath + '/train.txt', 'w') fval = open(txtsavepath + '/val.txt', 'w') total_files = glob("./VOCdevkit/VOC2007/Annotations/*.xml") total_files = [i.split("/")[-1].split(".xml")[0] for i in total_files]

test_filepath = "/Users/ysj/Desktop/2022 年（第 15 屆）中國大學生計算機設計大賽人工智慧挑戰賽-智慧零售賽項資料集/TestDataset/"

for file in total_files: ftrainval.write(file + "\n")

test

for file in os.listdir(test_filepath):

ftest.write(file.split(".png")[0] + "\n")

split，根據test_size這個引數來確定test的數量

train_files, val_files = train_test_split(total_files, test_size=0.001, random_state=42)

train

for file in train_files: ftrain.write(file + "\n") #ftest.write(file + "\n")

val

for file in val_files: fval.write(file + "\n")

ftrainval.close() ftrain.close() fval.close()

ftest.close()

```

得到的資料集如下

image-20220513112741030

2. 訓練模型

準備好了資料集,接著就得找一個好的模型進行訓練.為了後面的部署方便,我這裡選擇的是YOLOX.其他大多數模型在後面轉ONNX格式的時候會運算元不相容或者其他問題無法轉換.為了簡單起見所以直接選擇YOLOX而且程式碼中就自帶有轉ONNX和TRT部分的程式碼.

把YOLOX克隆之後改一下里面對應的類別數,類別名稱,把剛才準備好的資料複製到datasets裡面.下載一個yolox_s的預訓練模型,然後開始train(為了節約,直接半精度訓練) 預設訓練最多300epoch,想更改可以去yolox_base.py裡面改max_epoch.訓練耗時並不久,很快就得到了一個訓練集上表現良好的模型.

image-20220513113958230

然後驗證一下模型效果

python tools/eval.py -f ../exps/example/yolox_voc/yolox_voc_s.py -c ../YOLOX_outputs/yolox_voc_s/best_ckpt.pth -b 8 -d 0 --conf 0.001 --fp16

image-20220513114131157

image-20220513114159823

使用模型預測一下圖片

python tools/demo.py image -n yolox-s -c ./YOLOX_outputs/yolox_voc_s/best_ckpt.pth --path assets/good_640.png --conf 0.5 --nms 0.45 --tsize 640 --save_result --device gpu

opencv不支援中文顯示,一般都需要引入其他字元庫或者改寫PutText,當然還可以嘗試用PIL的ImageDraw來繪製圖片,也相當於改寫繪製函式.這裡我直接全部寫成拼音圖簡單

image-20220513125343680

3. 轉換模型格式

我們目前得到的是pytorch生成的pth,我們目標是onnx和trt.使用export_onnx.py我們可以得到onnx檔案.因為有onnxsim,所以轉換後的模型是優化過的,大小會比pth小很多.

trt.py可以得到trt的.engine檔案,但是如果想要trt檔案.這個時候使用tensorRT的trtexec可以將onnx轉為trt檔案trtexec --onnx='xxx.onnx' --saveEngine='xxx.trt' --workspace=xxx --fp16

4. 使用tensorRT改寫

這部分可以參考yolox中demo/TensorRT下的cpp進行仿寫,也可以根據TensorRT自帶的一些example來改寫,還有一些網上開源的程式碼也可以參考.如果想快速實現,可以參考一下這個http://gitee.com/xiaoyuerCV/tensorrt-yolox/tree/master

裡面的CMakeLists根據自己的路徑引入庫和連結,然後它的程式碼裡有一個小地方需要自己加上,這個應該是最近TensorRT更新過所以繼承的時候要新增,如下圖

image-20220513130649991

其他地方基本不用動,改改自己的類別以及一些引數就行.然後一些功能根據自己需要新增,比如我想得到每張圖裡商品的名稱,置信度以及總價格

image-20220513131317829

效果

image-20220513131540405

image-20220513131841247

目前還是輸入圖片路徑進行檢測,後期可以改寫成用Capture捕獲攝像頭進行檢測,當然也可以用python搭建簡單的api直接tensorrt呼叫模型作預測.總之只要模型有了,後面的可玩性還是很強的.而且不得不說曠視確實牛皮,yolox訓練快效果好,之前也試過其他模型對於這批資料效果並不理想.

最後

配置一系列環境比如TensorRT、torch2trt、apex、opencv……應該都比較順利，如果有問題可以在評論區一起討論下，如果問題多我就再水一篇配置環境的部落格。

「」
「」