1、專案說明
在該專案中,主要向大家介紹如何使用目標檢測和語意分割來實作對指標型表計讀數,
在電力能源廠區需要定期監測表計讀數,以保證設備正常運行及廠區安全,但廠區分布分散,人工巡檢耗時長,無法實時監測表計,且部分作業環境危險導致人工巡檢無法觸達,針對上述問題,希望通過攝像頭拍照->智能讀數的方式高效地完成此任務,

為實作智能讀數,我們采取目標檢測->語意分割->讀數后處理的方案:
第一步,使用目標檢測模型定位出影像中的表計;
第二步,使用語意分割模型將各表計的指標和刻度分割出來;
第三步,根據指標的相對位置和預知的量程計算出各表計的讀數,

2、資料準備
本案例開放了表計檢測資料集、指標和刻度分割資料集、表計測驗圖片(只有圖片無真值標注),使用這些圖片可以完成目標檢測模型、語意分割模型的訓練、模型預測,點擊下表中的鏈接可下載資料集,提前下載資料集不是必須的,因為在接下來的模型訓練部分中提供的訓練腳本會自動下載資料集,
表計測驗圖片 表計檢測資料集 指標和刻度分割資料集
meter_test meter_det meter_seg
解壓后的表計檢測資料集的檔案夾內容如下:
訓練集有725張圖片,測驗集有58張圖片,
meter_det/
|-- annotations/ # 標注檔案所在檔案夾
| |-- instance_train.json # 訓練集標注檔案
| |-- instance_test.json # 測驗集標注檔案
|-- test/ # 測驗圖片所在檔案夾
| |-- 20190822_105.jpg # 測驗集圖片
| |-- … …
|-- train/ # 訓練圖片所在檔案夾
| |-- 20190822_101.jpg # 訓練集圖片
| |-- … …
解壓后的指標和刻度分割資料集的檔案夾內容如下:
訓練集有374張圖片,測驗集有40張圖片,
meter_seg/
|-- annotations/ # 標注檔案所在檔案夾
| |-- train # 訓練集標注圖片所在檔案夾
| | |-- 105.png
| | |-- … …
| |-- val # 驗證集合標注圖片所在檔案夾
| | |-- 110.png
| | |-- … …
|-- images/ # 圖片所在檔案夾
| |-- train # 訓練集圖片
| | |-- 105.jpg
| | |-- … …
| |-- val # 驗證集圖片
| | |-- 110.jpg
| | |-- … …
|-- labels.txt # 類別名串列
|-- train.txt # 訓練集圖片串列
|-- val.txt # 驗證集圖片串列
解壓后的表計測驗圖片的檔案夾內容如下:
一共有58張測驗圖片,
meter_test/
|-- 20190822_105.jpg
|-- 20190822_142.jpg
|-- … …
3、模型選擇
PaddleX提供了豐富的視覺模型,在目標檢測中提供了RCNN和YOLO系列模型,在語意分割中提供了DeepLabV3P和BiSeNetV2等模型,
因最終部署場景是本地化的服務器GPU端,算力相對充足,因此在本專案中采用精度和預測性能皆優的PPYOLOV2進行表計檢測,
考慮到指標和刻度均為細小區域,我們采用精度更優的DeepLabV3P進行指標和刻度的分割,
4、表計檢測模型訓練
本專案中采用精度和預測性能的PPYOLOV2進行表計檢測,
訓練結束后,最優模型精度bbox_mmap達到100%,
5、 指標和刻度分割模型訓練
本專案中采用精度更優的DeepLabV3P進行指標和刻度的分割,
訓練結束后,最優模型精度miou達84.09,
準備階段
第一步:創建Notebook模型任務
step1:進入BML主頁,點擊立即使用
🔗:https://ai.baidu.com/bml/

step2:點擊Notebook,創建“通用任務”

step3:填寫任務資訊

第二步:下載任務操作模板
下載鏈接:https://aistudio.baidu.com/aistudio/datasetdetail/120387

目標檢測模型訓練
第一步:配置Notebook
1.找到昨天創建的Notebook任務,點擊配置

2.配置選擇
- 開發語言:Python3.7
- AI框架:PaddlePaddle2.0.0
- 資源規格:GPU V100

3.打開Notebook

4.上傳本次Notebook操作模型
若沒來得及下載,請點擊鏈接下載:https://aistudio.baidu.com/aistudio/datasetdetail/120387


第二步:環境準備
1.安裝filelock
!pip install filelock


2.安裝PaddleX
!pip install paddlex==2.0.0
注意:安裝paddlex的時候需要制定版本,


3.升級paddlepaddle-gpu為2.1.3版本
!pip install paddlepaddle-gpu==2.1.3.post101 -f https://www.paddlepaddle.org.cn/whl/linux/mkl/avx/stable.html


第三步:目標檢測模型訓練
訓練程序說明:
定義資料預處理 -> 定義資料集路徑 -> 初始化模型 -> 模型訓練

1.呼叫PaddleX
import paddlex as pdxfrom paddlex import transforms as T

2.定義訓練和驗證時的transforms
API詳細說明:https://github.com/PaddlePaddle/PaddleX/blob/release/2.0-rc/paddlex/cv/transforms/operators.py
train_transforms = T.Compose([ T.MixupImage(mixup_epoch=250), T.RandomDistort(), T.RandomExpand(im_padding_value=[123.675, 116.28, 103.53]), T.RandomCrop(), T.RandomHorizontalFlip(), T.BatchRandomResize( target_sizes=[320, 352, 384, 416, 448, 480, 512, 544, 576, 608], interp='RANDOM'), T.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
eval_transforms = T.Compose([ T.Resize( 608, interp='CUBIC'), T.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])

3.下載用于目標檢測訓練的表計讀數資料集
meter_det_dataset = 'https://bj.bcebos.com/paddlex/examples/meter_reader/datasets/meter_det.tar.gz'pdx.utils.download_and_decompress(meter_det_dataset, path='./')

可在左側檔案夾區域查看資料集

4.設定訓練引數
詳細API說明:https://github.com/PaddlePaddle/PaddleX/blob/develop/paddlex/cv/datasets/coco.py#L26
train_dataset = pdx.datasets.CocoDetection( data_dir='meter_det/train/', ann_file='meter_det/annotations/instance_train.json', transforms=train_transforms, shuffle=True)eval_dataset = pdx.datasets.CocoDetection( data_dir='meter_det/test/', ann_file='meter_det/annotations/instance_test.json', transforms=eval_transforms)



5.訓練結束后查看bestmodel


第四步:保存Notebook并關閉、停止運行


語意分割模型訓練
第一步:重新安裝環境
1.啟動Notebook并打開

2.重新執行安裝的三條命令



第二步:指標和刻度分割模型訓練
1.呼叫paddlex
import paddlex as pdxfrom paddlex import transforms as T

2.定義訓練和驗證時的transforms
詳細API說明參考:https://github.com/PaddlePaddle/PaddleX/blob/release/2.0-rc/paddlex/cv/transforms/operators.py
train_transforms = T.Compose([ T.Resize(target_size=512), T.RandomHorizontalFlip(), T.Normalize( mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),])
eval_transforms = T.Compose([ T.Resize(target_size=512), T.Normalize( mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),])

3.下載和解壓指標刻度分割資料集
meter_seg_dataset = 'https://bj.bcebos.com/paddlex/examples/meter_reader/datasets/meter_seg.tar.gz'pdx.utils.download_and_decompress(meter_seg_dataset, path='./')

4.定義訓練和驗證所用的資料集,配置相應路徑
詳細API說明參考:https://github.com/PaddlePaddle/PaddleX/blob/release/2.0-rc/paddlex/cv/datasets/seg_dataset.py#L22
train_dataset = pdx.datasets.SegDataset( data_dir='meter_seg', file_list='meter_seg/train.txt', label_list='meter_seg/labels.txt', transforms=train_transforms, shuffle=True)
eval_dataset = pdx.datasets.SegDataset( data_dir='meter_seg', file_list='meter_seg/val.txt', label_list='meter_seg/labels.txt', transforms=eval_transforms, shuffle=False)

5.選擇PaddleX內置的DeepLabV3P模型進行訓練
API說明:https://github.com/PaddlePaddle/PaddleX/blob/release/2.0-rc/paddlex/cv/models/segmenter.py#L150
num_classes = len(train_dataset.labels)model = pdx.seg.DeepLabV3P( num_classes=num_classes, backbone='ResNet50_vd', use_mixed_loss=True)

6.設定訓練時的引數
各引數介紹與調整說明:https://paddlex.readthedocs.io/zh_CN/develop/appendix/parameters.html
model.train( num_epochs=2, train_dataset=train_dataset, train_batch_size=4, eval_dataset=eval_dataset, pretrain_weights='IMAGENET', learning_rate=0.1, save_dir='output/deeplabv3p_r50vd')



7.訓練結束后查看bestmodel


第三步:保存Notebook并關閉、停止運行


提示:Notebook一旦運行即會開始計費,如果不用請及時停止!以免浪費免費額度
模型預測
第一步:重新安裝環境
1.啟動Notebook并打開

2.重新執行安裝的三條命令



第二步:上傳預測的py檔案
1.點擊下方鏈接將py檔案下載至本地
https://aistudio.baidu.com/aistudio/datasetdetail/120795

2.上傳至Notebook中

第三步:模型預測

1.上傳reader_infer.py檔案后,執行一下命令進行模型預測
!python work/meter_reader/reader_infer.py --det_model_dir output/ppyolov2_r50vd_dcn/best_model --seg_model_dir output/deeplabv3p_r50vd/best_model/ --image meter_det/test/20190822_105.jpg


2.打開output/result查看預測結果

第四步:保存Notebook并關閉、停止運行


推理代碼:
# coding: utf8
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import os.path as osp
import numpy as np
import math
import cv2
import argparse
from paddlex import transforms as T
import paddlex as pdx
# 讀數后處理中有把圓形表盤轉成矩形的操作,矩形的寬即為圓形的外周長
# 因此要求表盤影像大小為固定大小,這里設定為[512, 512]
METER_SHAPE = [512, 512] # 高x寬
# 圓形表盤的中心點
CIRCLE_CENTER = [256, 256] # 高x寬
# 圓形表盤的半徑
CIRCLE_RADIUS = 250
# 圓周率
PI = 3.1415926536
# 在把圓形表盤轉成矩形后矩形的高
# 當前設定值約為半徑的一半,原因是:圓形表盤的中心區域除了指標根部就是背景了
# 我們只需要把外圍的刻度、指標的尖部保存下來就可以定位出指標指向的刻度
RECTANGLE_HEIGHT = 120
# 矩形表盤的寬,即圓形表盤的外周長
RECTANGLE_WIDTH = 1570
# 當前案例中只使用了兩種型別的表盤,第一種表盤的刻度根數為50
# 第二種表盤的刻度根數為32,因此,我們通過預測的刻度根數來判斷表盤型別
# 刻度根數超過閾值的即為第一種,否則是第二種
TYPE_THRESHOLD = 40
# 兩種表盤的配置資訊,包含每根刻度的值,量程,單位
METER_CONFIG = [{
'scale_interval_value': 25.0 / 50.0,
'range': 25.0,
'unit': "(MPa)"
}, {
'scale_interval_value': 1.6 / 32.0,
'range': 1.6,
'unit': "(MPa)"
}]
# 分割模型預測類別id與類別名的對應關系
SEG_CNAME2CLSID = {'background': 0, 'pointer': 1, 'scale': 2}
def parse_args():
parser = argparse.ArgumentParser(description='Meter Reader Infering')
parser.add_argument(
'--det_model_dir',
dest='det_model_dir',
help='The directory of the detection model',
type=str)
parser.add_argument(
'--seg_model_dir',
dest='seg_model_dir',
help='The directory of the segmentation model',
type=str)
parser.add_argument(
'--image_dir',
dest='image_dir',
help='The directory of images to be inferred',
type=str,
default=None)
parser.add_argument(
'--image',
dest='image',
help='The image to be inferred',
type=str,
default=None)
parser.add_argument(
'--use_erode',
dest='use_erode',
help='Whether erode the lable map predicted from a segmentation model',
action='store_true')
parser.add_argument(
'--erode_kernel',
dest='erode_kernel',
help='Erode kernel size',
type=int,
default=4)
parser.add_argument(
'--save_dir',
dest='save_dir',
help='The directory for saving the predicted results',
type=str,
default='./output/result')
parser.add_argument(
'--score_threshold',
dest='score_threshold',
help="Predicted bounding boxes whose scores are lower than this threshlod are filtered",
type=float,
default=0.5)
parser.add_argument(
'--seg_batch_size',
dest='seg_batch_size',
help="The number of images fed into the segmentation model during one forward propagation",
type=int,
default=2)
return parser.parse_args()
def is_pic(img_name):
"""判斷是否是圖片
引數:
img_name (str): 圖片路徑
回傳:
flag (bool): 判斷值,
"""
valid_suffix = ['JPEG', 'jpeg', 'JPG', 'jpg', 'BMP', 'bmp', 'PNG', 'png']
suffix = img_name.split('.')[-1]
flag = True
if suffix not in valid_suffix:
flag = False
return flag
class MeterReader:
"""檢測表盤的位置,分割各表盤內刻度和指標的位置,并根據分割結果計算出各表盤的讀數
引數:
det_model_dir (str): 用于定位表盤的檢測模型所在路徑,
seg_model_dir (str): 用于分割刻度和指標的分割模型所在路徑,
"""
def __init__(self, det_model_dir, seg_model_dir):
if not osp.exists(det_model_dir):
raise Exception("Model path {} does not exist".format(
det_model_dir))
if not osp.exists(seg_model_dir):
raise Exception("Model path {} does not exist".format(
seg_model_dir))
self.detector = pdx.load_model(det_model_dir)
self.segmenter = pdx.load_model(seg_model_dir)
def decode(self, img_file):
"""影像解碼
引數:
img_file (str|np.array): 影像路徑,或者是已解碼的BGR影像陣列,
回傳:
img (np.array): BGR影像陣列,
"""
if isinstance(img_file, str):
img = cv2.imread(img_file).astype('float32')
else:
img = img_file.copy()
return img
def filter_bboxes(self, det_results, score_threshold):
"""過濾置信度低于閾值的檢測框
引數:
det_results (list[dict]): 檢測模型預測介面的回傳值,
score_threshold (float):置信度閾值,
回傳:
filtered_results (list[dict]): 過濾后的檢測狂,
"""
filtered_results = list()
for res in det_results:
if res['score'] > score_threshold:
filtered_results.append(res)
return filtered_results
def roi_crop(self, img, det_results):
"""摳取影像上各檢測框的影像區域
引數:
img (np.array):BRG影像陣列,
det_results (list[dict]): 檢測模型預測介面的回傳值,
回傳:
sub_imgs (list[np.array]): 各檢測框的影像區域,
"""
sub_imgs = []
for res in det_results:
# Crop the bbox area
xmin, ymin, w, h = res['bbox']
xmin = max(0, int(xmin))
ymin = max(0, int(ymin))
xmax = min(img.shape[1], int(xmin + w - 1))
ymax = min(img.shape[0], int(ymin + h - 1))
sub_img = img[ymin:(ymax + 1), xmin:(xmax + 1), :]
sub_imgs.append(sub_img)
return sub_imgs
def resize(self, imgs, target_size, interp=cv2.INTER_LINEAR):
"""影像縮放至固定大小
引數:
imgs (list[np.array]):批量BGR影像陣列,
target_size (list|tuple):縮放后的影像大小,格式為[高, 寬],
interp (int):影像差值方法,默認值為cv2.INTER_LINEAR,
回傳:
resized_imgs (list[np.array]):縮放后的批量BGR影像陣列,
"""
resized_imgs = list()
for img in imgs:
img_shape = img.shape
scale_x = float(target_size[1]) / float(img_shape[1])
scale_y = float(target_size[0]) / float(img_shape[0])
resize_img = cv2.resize(
img, None, None, fx=scale_x, fy=scale_y, interpolation=interp)
resized_imgs.append(resize_img)
return resized_imgs
def seg_predict(self, segmenter, imgs, batch_size):
"""分割模型完成預測
引數:
segmenter (pdx.seg.model):加載后的分割模型,
imgs (list[np.array]):待預測的輸入BGR影像陣列,
batch_size (int): 分割模型前向預測一次時輸入影像的批量大小,
回傳:
seg_results (list[dict]): 輸入影像的預測結果,
"""
seg_results = list()
num_imgs = len(imgs)
for i in range(0, num_imgs, batch_size):
batch = imgs[i:min(num_imgs, i + batch_size)]
result = segmenter.predict(batch)
seg_results.extend(result)
return seg_results
def erode(self, seg_results, erode_kernel):
"""對分割模型預測結果中label_map做影像腐蝕操作
引數:
seg_results (list[dict]):分割模型的預測結果,
erode_kernel (int): 影像腐蝕的卷積核的大小,
回傳:
eroded_results (list[dict]):對label_map進行腐蝕后的分割模型預測結果,
"""
kernel = np.ones((erode_kernel, erode_kernel), np.uint8)
eroded_results = seg_results
for i in range(len(seg_results)):
test_resulte = seg_results[i]['label_map']
# print('***********************************',type(test_resulte))
# eroded_results[i]['label_map'] = cv2.erode(
# seg_results[i]['label_map'], kernel)
eroded_results[i]['label_map'] = cv2.erode(
test_resulte.astype('uint8'), kernel)
return eroded_results
def circle_to_rectangle(self, seg_results):
"""將圓形表盤的預測結果label_map轉換成矩形
圓形到矩形的計算方法:
因本案例中兩種表盤的刻度起始值都在左下方,故以圓形的中心點為坐標原點,
從-y軸開始逆時針計算極坐標到x-y坐標的對應關系:
x = r + r * cos(theta)
y = r - r * sin(theta)
注意:
1. 因為是從-y軸開始逆時針計算,所以r * sin(theta)前有負號,
2. 還是因為從-y軸開始逆時針計算,所以矩形從上往下對應圓形從外到內,
可以想象把圓形從-y軸切開再往左右拉平時,圓形的外圍是上面,內圍在下面,
引數:
seg_results (list[dict]):分割模型的預測結果,
回傳值:
rectangle_meters (list[np.array]):矩形表盤的預測結果label_map,
"""
rectangle_meters = list()
for i, seg_result in enumerate(seg_results):
label_map = seg_result['label_map']
# rectangle_meter的大小已經由預先設定的全域變數RECTANGLE_HEIGHT, RECTANGLE_WIDTH決定
rectangle_meter = np.zeros(
(RECTANGLE_HEIGHT, RECTANGLE_WIDTH), dtype=np.uint8)
for row in range(RECTANGLE_HEIGHT):
for col in range(RECTANGLE_WIDTH):
theta = PI * 2 * (col + 1) / RECTANGLE_WIDTH
# 矩形從上往下對應圓形從外到內
rho = CIRCLE_RADIUS - row - 1
y = int(CIRCLE_CENTER[0] + rho * math.cos(theta) + 0.5)
x = int(CIRCLE_CENTER[1] - rho * math.sin(theta) + 0.5)
rectangle_meter[row, col] = label_map[y, x]
rectangle_meters.append(rectangle_meter)
return rectangle_meters
def rectangle_to_line(self, rectangle_meters):
"""從矩形表盤的預測結果中提取指標和刻度預測結果并沿高度方向壓縮成線狀格式,
引數:
rectangle_meters (list[np.array]):矩形表盤的預測結果label_map,
回傳:
line_scales (list[np.array]):刻度的線狀預測結果,
line_pointers (list[np.array]):指標的線狀預測結果,
"""
line_scales = list()
line_pointers = list()
for rectangle_meter in rectangle_meters:
height, width = rectangle_meter.shape[0:2]
line_scale = np.zeros((width), dtype=np.uint8)
line_pointer = np.zeros((width), dtype=np.uint8)
for col in range(width):
for row in range(height):
if rectangle_meter[row, col] == SEG_CNAME2CLSID['pointer']:
line_pointer[col] += 1
elif rectangle_meter[row, col] == SEG_CNAME2CLSID['scale']:
line_scale[col] += 1
line_scales.append(line_scale)
line_pointers.append(line_pointer)
return line_scales, line_pointers
def mean_binarization(self, data_list):
"""對影像進行均值二值化操作
引數:
data_list (list[np.array]):待二值化的批量陣列,
回傳:
binaried_data_list (list[np.array]):二值化后的批量陣列,
"""
batch_size = len(data_list)
binaried_data_list = data_list
for i in range(batch_size):
mean_data = np.mean(data_list[i])
width = data_list[i].shape[0]
for col in range(width):
if data_list[i][col] < mean_data:
binaried_data_list[i][col] = 0
else:
binaried_data_list[i][col] = 1
return binaried_data_list
def locate_scale(self, line_scales):
"""在線狀預測結果中找到每根刻度的中心位置
引數:
line_scales (list[np.array]):批量的二值化后的刻度線狀預測結果,
回傳:
scale_locations (list[list]):各影像中每根刻度的中心位置,
"""
batch_size = len(line_scales)
scale_locations = list()
for i in range(batch_size):
line_scale = line_scales[i]
width = line_scale.shape[0]
find_start = False
one_scale_start = 0
one_scale_end = 0
locations = list()
for j in range(width - 1):
if line_scale[j] > 0 and line_scale[j + 1] > 0:
if find_start == False:
one_scale_start = j
find_start = True
if find_start:
if line_scale[j] == 0 and line_scale[j + 1] == 0:
one_scale_end = j - 1
one_scale_location = (
one_scale_start + one_scale_end) / 2
locations.append(one_scale_location)
one_scale_start = 0
one_scale_end = 0
find_start = False
scale_locations.append(locations)
return scale_locations
def locate_pointer(self, line_pointers):
"""在線狀預測結果中找到指標的中心位置
引數:
line_scales (list[np.array]):批量的指標線狀預測結果,
回傳:
scale_locations (list[list]):各影像中指標的中心位置,
"""
batch_size = len(line_pointers)
pointer_locations = list()
for i in range(batch_size):
line_pointer = line_pointers[i]
find_start = False
pointer_start = 0
pointer_end = 0
location = 0
width = line_pointer.shape[0]
for j in range(width - 1):
if line_pointer[j] > 0 and line_pointer[j + 1] > 0:
if find_start == False:
pointer_start = j
find_start = True
if find_start:
if line_pointer[j] == 0 and line_pointer[j + 1] == 0:
pointer_end = j - 1
location = (pointer_start + pointer_end) / 2
find_start = False
break
pointer_locations.append(location)
return pointer_locations
def get_relative_location(self, scale_locations, pointer_locations):
"""找到指標指向了第幾根刻度
引數:
scale_locations (list[list]):批量的每根刻度的中心點位置,
pointer_locations (list[list]):批量的指標的中心點位置,
回傳:
pointed_scales (list[dict]):每個表的結果組成的list,每個表的結果由字典表示,
字典有兩個關鍵詞:'num_scales'、'pointed_scale',分別表示預測的刻度根數、
預測的指標指向了第幾根刻度,
"""
pointed_scales = list()
for scale_location, pointer_location in zip(scale_locations,
pointer_locations):
num_scales = len(scale_location)
pointed_scale = -1
if num_scales > 0:
for i in range(num_scales - 1):
if scale_location[
i] <= pointer_location and pointer_location < scale_location[
i + 1]:
pointed_scale = i + (
pointer_location - scale_location[i]
) / (scale_location[i + 1] - scale_location[i] + 1e-05
) + 1
result = {'num_scales': num_scales, 'pointed_scale': pointed_scale}
pointed_scales.append(result)
return pointed_scales
def calculate_reading(self, pointed_scales):
"""根據刻度的間隔值和指標指向的刻度根數計算表盤的讀數
"""
readings = list()
batch_size = len(pointed_scales)
for i in range(batch_size):
pointed_scale = pointed_scales[i]
# 刻度根數大于閾值的為第一種表盤
if pointed_scale['num_scales'] > TYPE_THRESHOLD:
reading = pointed_scale['pointed_scale'] * METER_CONFIG[0][
'scale_interval_value']
else:
reading = pointed_scale['pointed_scale'] * METER_CONFIG[1][
'scale_interval_value']
readings.append(reading)
return readings
def get_meter_reading(self, seg_results):
"""對分割結果進行讀數后處理得到各表盤的讀數
引數:
seg_results (list[dict]): 分割模型的預測結果,
回傳:
meter_readings (list[dcit]): 各表盤的讀數,
"""
rectangle_meters = self.circle_to_rectangle(seg_results)
line_scales, line_pointers = self.rectangle_to_line(rectangle_meters)
binaried_scales = self.mean_binarization(line_scales)
binaried_pointers = self.mean_binarization(line_pointers)
scale_locations = self.locate_scale(binaried_scales)
pointer_locations = self.locate_pointer(binaried_pointers)
pointed_scales = self.get_relative_location(scale_locations,
pointer_locations)
meter_readings = self.calculate_reading(pointed_scales)
return meter_readings
def print_meter_readings(self, meter_readings):
"""列印各表盤的讀數
引數:
meter_readings (list[dict]):各表盤的讀數
"""
for i in range(len(meter_readings)):
print("Meter {}: {}".format(i + 1, meter_readings[i]))
def visualize(self, img, det_results, meter_readings, save_dir="./"):
"""可視化影像中各表盤的位置和讀數
引數:
img (str|np.array): 影像路徑,或者是已解碼的BGR影像陣列,
det_results (dict): 檢測模型的預測結果,
meter_readings (list): 各表盤的讀數,
save_dir (str):可視化后的圖片保存路徑,
"""
vis_results = list()
for i, res in enumerate(det_results):
# 將檢測結果中的關鍵詞`score`替換成讀數,就可以呼叫pdx.det.visualize畫圖了
res['score'] = meter_readings[i]
vis_results.append(res)
# 檢測結果可視化時會濾除score低于threshold的框,這里讀數都是>=-1的,所以設定thresh=-1
pdx.det.visualize(img, vis_results, threshold=-1, save_dir=save_dir)
def predict(self,
img_file,
save_dir='./',
use_erode=True,
erode_kernel=4,
score_threshold=0.5,
seg_batch_size=2):
"""檢測影像中的表盤,而后分割出各表盤中的指標和刻度,對分割結果進行讀數后處理后得到各表盤的讀數,
引數:
img_file (str):待預測的圖片路徑,
save_dir (str): 可視化結果的保存路徑,
use_erode (bool, optional): 是否對分割預測結果做影像腐蝕,默認值:True,
erode_kernel (int, optional): 影像腐蝕的卷積核大小,默認值: 4,
score_threshold (float, optional): 用于濾除檢測框的置信度閾值,默認值:0.5,
seg_batch_size (int, optional):分割模型前向推理一次時輸入表盤影像的批量大小,默認值為:2,
"""
img = self.decode(img_file)
det_results = self.detector.predict(img)
filtered_results = self.filter_bboxes(det_results, score_threshold)
sub_imgs = self.roi_crop(img, filtered_results)
sub_imgs = self.resize(sub_imgs, METER_SHAPE)
seg_results = self.seg_predict(self.segmenter, sub_imgs,
seg_batch_size)
seg_results = self.erode(seg_results, erode_kernel)
meter_readings = self.get_meter_reading(seg_results)
self.print_meter_readings(meter_readings)
self.visualize(img, filtered_results, meter_readings, save_dir)
def infer(args):
image_lists = list()
if args.image is not None:
if not osp.exists(args.image):
raise Exception("Image {} does not exist.".format(args.image))
if not is_pic(args.image):
raise Exception("{} is not a picture.".format(args.image))
image_lists.append(args.image)
elif args.image_dir is not None:
if not osp.exists(args.image_dir):
raise Exception("Directory {} does not exist.".format(
args.image_dir))
for im_file in os.listdir(args.image_dir):
if not is_pic(im_file):
continue
im_file = osp.join(args.image_dir, im_file)
image_lists.append(im_file)
meter_reader = MeterReader(args.det_model_dir, args.seg_model_dir)
if len(image_lists) > 0:
for image in image_lists:
meter_reader.predict(image, args.save_dir, args.use_erode,
args.erode_kernel, args.score_threshold,
args.seg_batch_size)
if __name__ == '__main__':
args = parse_args()
infer(args)
轉載請註明出處,本文鏈接:https://www.uj5u.com/qita/382136.html
標籤:其他
