[西安交通大学][边缘开发套件社区]铝材表面缺陷检测

This commit is contained in:
jtc
2022-09-20 16:47:28 +08:00
parent 1bd93f2cc7
commit 3f33e51c14
6 changed files with 50 additions and 150 deletions

View File

@@ -15,14 +15,13 @@
import os import os
import json import json
import stat import stat
import cv2
from StreamManagerApi import *
import time
import numpy as np
from utils import *
import glob import glob
import cv2
from StreamManagerApi import StreamManagerApi, MxDataInput
import numpy as np
from utils import xyxy2xywh
from plots import Annotator, colors from plots import box_label, colors
names = ['non_conduct', 'abrasion_mark', 'corner_leak', 'orange_peel', 'leak', 'jet_flow', 'paint_bubble', 'pit', names = ['non_conduct', 'abrasion_mark', 'corner_leak', 'orange_peel', 'leak', 'jet_flow', 'paint_bubble', 'pit',
'motley', 'dirty_spot'] 'motley', 'dirty_spot']
@@ -40,7 +39,6 @@ if __name__ == '__main__':
if ret != 0: if ret != 0:
print("Failed to init Stream manager, ret=%s" % str(ret)) print("Failed to init Stream manager, ret=%s" % str(ret))
exit() exit()
start = time.time()
# create streams by pipeline config file # create streams by pipeline config file
with open("./pipeline/AlDefectDetection.pipeline", 'rb') as f: with open("./pipeline/AlDefectDetection.pipeline", 'rb') as f:
pipelineStr = f.read() pipelineStr = f.read()
@@ -95,8 +93,6 @@ if __name__ == '__main__':
with open(ori_img_path, 'rb') as f: with open(ori_img_path, 'rb') as f:
dataInput.data = f.read() dataInput.data = f.read()
annotator = Annotator(ori_img, line_width=3, example=str(names))
# Inputs data to a specified stream based on streamName. # Inputs data to a specified stream based on streamName.
streamName = b'classification+detection' streamName = b'classification+detection'
inPluginId = 0 inPluginId = 0
@@ -135,9 +131,8 @@ if __name__ == '__main__':
f.write(('%g ' * len(line)).rstrip() % line + '\n') f.write(('%g ' * len(line)).rstrip() % line + '\n')
label = f'{classVec[0]["className"]} {classVec[0]["confidence"]:.4f}' label = f'{classVec[0]["className"]} {classVec[0]["confidence"]:.4f}'
annotator.box_label(xyxy, label, color=colors(names.index(classVec[0]["className"]), False)) save_img = box_label(ori_img, xyxy, label, color=colors[names.index(classVec[0]["className"])])
save_img = annotator.result()
cv2.imwrite(DETECT_IMG_PATH + 'result' + item, save_img) cv2.imwrite(DETECT_IMG_PATH + 'result' + item, save_img)
TESTIMGS += 1 TESTIMGS += 1
###################################################################################### ######################################################################################

View File

@@ -15,16 +15,15 @@
import os import os
import json import json
import stat import stat
import glob
import cv2 import cv2
from StreamManagerApi import * from StreamManagerApi import StreamManagerApi, MxDataInput
import time
import numpy as np import numpy as np
from utils import * from utils import preprocess, scale_coords, xyxy2xywh
from plots import Annotator, colors from plots import box_label, colors
names = ['non_conduct', 'abrasion_mark', 'corner_leak', 'orange_peel', 'leak', 'jet_flow', 'paint_bubble', 'pit', names = ['non_conduct', 'abrasion_mark', 'corner_leak', 'orange_peel', 'leak', 'jet_flow', 'paint_bubble', 'pit',
'motley', 'dirty_spot'] 'motley', 'dirty_spot']
import glob
if __name__ == '__main__': if __name__ == '__main__':
MODES = stat.S_IWUSR | stat.S_IRUSR MODES = stat.S_IWUSR | stat.S_IRUSR
@@ -81,7 +80,7 @@ if __name__ == '__main__':
r = img_size / max(h0, w0) # ratio r = img_size / max(h0, w0) # ratio
input_shape = (640, 640) input_shape = (640, 640)
pre_img = letterbox(ori_img, input_shape)[0] pre_img = preprocess(ori_img, input_shape)[0]
pre_img = np.ascontiguousarray(pre_img) pre_img = np.ascontiguousarray(pre_img)
pre_img_path = PRE_IMG_PATH + item pre_img_path = PRE_IMG_PATH + item
@@ -97,7 +96,6 @@ if __name__ == '__main__':
dataInput = MxDataInput() dataInput = MxDataInput()
with open(pre_img_path, 'rb') as f: with open(pre_img_path, 'rb') as f:
dataInput.data = f.read() dataInput.data = f.read()
annotator = Annotator(ori_img, line_width=3, example=str(names))
# Inputs data to a specified stream based on streamName. # Inputs data to a specified stream based on streamName.
streamName = b'classification+detection' streamName = b'classification+detection'
@@ -137,9 +135,8 @@ if __name__ == '__main__':
f.write(('%g ' * len(line)).rstrip() % line + '\n') f.write(('%g ' * len(line)).rstrip() % line + '\n')
label = f'{classVec[0]["className"]} {classVec[0]["confidence"]:.4f}' label = f'{classVec[0]["className"]} {classVec[0]["confidence"]:.4f}'
annotator.box_label(xyxy, label, color=colors(names.index(classVec[0]["className"]), False)) save_img = box_label(ori_img_path, xyxy, label, color=colors[names.index(classVec[0]["className"])])
save_img = annotator.result()
cv2.imwrite(DETECT_IMG_PATH + 'result' + item, save_img) cv2.imwrite(DETECT_IMG_PATH + 'result' + item, save_img)
TESTIMGS += 1 TESTIMGS += 1
###################################################################################### ######################################################################################

View File

@@ -15,10 +15,10 @@
import os import os
import json import json
import cv2 import cv2
from StreamManagerApi import * from StreamManagerApi import StreamManagerApi, MxDataInput
import numpy as np
from plots import Annotator, colors from plots import box_label, colors
from utils import * from utils import preprocess, scale_coords, xyxy2xywh
names = ['non_conduct', 'abrasion_mark', 'corner_leak', 'orange_peel', 'leak', 'jet_flow', 'paint_bubble', 'pit', names = ['non_conduct', 'abrasion_mark', 'corner_leak', 'orange_peel', 'leak', 'jet_flow', 'paint_bubble', 'pit',
'motley', 'dirty_spot'] 'motley', 'dirty_spot']
@@ -40,39 +40,38 @@ if __name__ == '__main__':
# Construct the input of the stream # Construct the input of the stream
dataInput = MxDataInput() dataInput = MxDataInput()
ori_img_path = "test.jpg" ORI_IMG_PATH = "test.jpg"
if not os.path.exists(ori_img_path): if not os.path.exists(ORI_IMG_PATH):
print("The test image does not exist.") print("The test image does not exist.")
exit() exit()
if os.path.getsize(ori_img_path) == 0: if os.path.getsize(ORI_IMG_PATH) == 0:
print("Error!The test image is empty.") print("Error!The test image is empty.")
exit() exit()
# read image # read image
ori_img = cv2.imread(ori_img_path) ori_img = cv2.imread(ORI_IMG_PATH)
h0, w0 = ori_img.shape[:2] h0, w0 = ori_img.shape[:2]
img_size = 640 r = 640 / max(h0, w0) # ratio
r = img_size / max(h0, w0) # ratio
input_shape = (640, 640) input_shape = (640, 640)
pre_img = letterbox(ori_img, input_shape)[0] pre_img = preprocess(ori_img, input_shape)[0]
pre_img = np.ascontiguousarray(pre_img) pre_img = np.ascontiguousarray(pre_img)
pre_img_path = "pre_" + ori_img_path PRE_IMG_PATH = "pre_" + ORI_IMG_PATH
cv2.imwrite(pre_img_path, pre_img) cv2.imwrite(PRE_IMG_PATH, pre_img)
with open(pre_img_path, 'rb') as f: with open(PRE_IMG_PATH, 'rb') as f:
dataInput.data = f.read() dataInput.data = f.read()
# Inputs data to a specified stream based on streamName. # Inputs data to a specified stream based on streamName.
streamName = b'classification+detection' STREAMNAME = b'classification+detection'
inPluginId = 0 INPLUGINID = 0
uniqueId = streamManagerApi.SendDataWithUniqueId(streamName, inPluginId, dataInput) uniqueId = streamManagerApi.SendDataWithUniqueId(STREAMNAME, INPLUGINID, dataInput)
if uniqueId < 0: if uniqueId < 0:
print("Failed to send data to stream.") print("Failed to send data to stream.")
exit() exit()
# Obtain the inference result by specifying streamName and uniqueId. # Obtain the inference result by specifying streamName and uniqueId.
inferResult = streamManagerApi.GetResultWithUniqueId(streamName, uniqueId, 10000) inferResult = streamManagerApi.GetResultWithUniqueId(STREAMNAME, INPLUGINID, 10000)
if inferResult.errorCode != 0: if inferResult.errorCode != 0:
print("GetResultWithUniqueId error. errorCode=%d, errorMsg=%s" % ( print("GetResultWithUniqueId error. errorCode=%d, errorMsg=%s" % (
inferResult.errorCode, inferResult.data.decode())) inferResult.errorCode, inferResult.data.decode()))
@@ -81,7 +80,6 @@ if __name__ == '__main__':
results = json.loads(inferResult.data.decode()) results = json.loads(inferResult.data.decode())
gn = np.array(ori_img.shape)[[1, 0, 1, 0]] gn = np.array(ori_img.shape)[[1, 0, 1, 0]]
annotator = Annotator(ori_img, line_width=3, example=str(names))
bboxes = [] bboxes = []
classVecs = [] classVecs = []
# draw the result and save image # draw the result and save image
@@ -93,10 +91,9 @@ if __name__ == '__main__':
xywh = (xyxy2xywh(xyxy.reshape(1, 4)) / gn).reshape(-1).tolist() # normalized xywh xywh = (xyxy2xywh(xyxy.reshape(1, 4)) / gn).reshape(-1).tolist() # normalized xywh
print(classVec) print(classVec)
label = f'{classVec[0]["className"]} {classVec[0]["confidence"]:.4f}' label = f'{classVec[0]["className"]} {classVec[0]["confidence"]:.4f}'
annotator.box_label(xyxy, label, color=colors(names.index(classVec[0]["className"]), False)) save_img = box_label(ori_img, xyxy, label, color=colors[names.index(classVec[0]["className"])])
save_img = annotator.result() cv2.imwrite('./result_' + ORI_IMG_PATH, save_img)
cv2.imwrite('./result_' + ori_img_path, save_img)
###################################################################################### ######################################################################################

View File

@@ -169,7 +169,6 @@ def compute_ap(recall, precision):
x = np.linspace(0, 1, 101) # 101-point interp (COCO) x = np.linspace(0, 1, 101) # 101-point interp (COCO)
ap = np.trapz(np.interp(x, mrec, mpre), x) # integrate ap = np.trapz(np.interp(x, mrec, mpre), x) # integrate
return ap return ap
@@ -266,10 +265,8 @@ def map_cac(opt):
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
ap = map_for_all_classes(*stats) ap = map_for_all_classes(*stats)
ap50, ap = ap[:, 0], ap.mean(1) # AP@0.5, AP@0.5:0.95 ap50, ap = ap[:, 0], ap.mean(1) # AP@0.5, AP@0.5:0.95
map50, map = ap50.mean(), ap.mean() print(f"mAP0.5: {ap50.mean()}")
print(f"mAP0.5: {map50}") print(f"mAP0.5:0.95: {ap.mean()}")
print(f"mAP0.5:0.95: {map}")
def parse_opt(): def parse_opt():

View File

@@ -15,103 +15,24 @@
import platform import platform
import os import os
from pathlib import Path from pathlib import Path
import re
from PIL import Image, ImageDraw, ImageFont from PIL import Image, ImageDraw, ImageFont
import numpy as np import numpy as np
import cv2 import cv2
# Settings colors = [(255, 56, 56), (255, 157, 151), (255, 112, 31), (255, 178, 29), (207, 210, 49), (72, 249, 10), (146, 204, 23),
RANK = int(os.getenv('RANK', -1)) (61, 219, 134), (26, 147, 52), (0, 212, 187)]
FONT = 'Arial.ttf'
import re
def is_ascii(s=''): def box_label(im, box, label='', color=(128, 128, 128), lw=3):
# Is string composed of all ASCII (no UTF) characters? (note str().isascii() introduced in python 3.7) p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
s = str(s) # convert list, tuple, None, etc. to str cv2.rectangle(im, p1, p2, color, thickness=lw, lineType=cv2.LINE_AA)
return len(s.encode().decode('ascii', 'ignore')) == len(s) if label:
tf = max(lw - 1, 1) # font thickness
w, h = cv2.getTextSize(label, 0, fontScale=lw / 3, thickness=tf)[0] # text width, height
def is_chinese(s='人工智能'): outside = p1[1] - h - 3 >= 0 # label fits outside box
# Is string composed of any Chinese characters? p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
return True if re.search('[\u4e00-\u9fff]', str(s)) else False cv2.rectangle(im, p1, p2, color, -1, cv2.LINE_AA) # filled
cv2.putText(im, label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2), 0, lw / 3, (255, 255, 255),
thickness=tf, lineType=cv2.LINE_AA)
def check_pil_font(font=FONT, size=10): return np.asarray(im)
# Return a PIL TrueType Font, downloading to CONFIG_DIR if necessary
font = Path(font)
font = font
class Colors:
# Ultralytics color palette https://ultralytics.com/
def __init__(self):
hex = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
'2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
self.palette = [self.hex2rgb('#' + c) for c in hex]
self.n = len(self.palette)
def __call__(self, i, bgr=False):
c = self.palette[int(i) % self.n]
return (c[2], c[1], c[0]) if bgr else c
@staticmethod
def hex2rgb(h): # rgb order (PIL)
return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))
colors = Colors() # create instance for 'from utils.plots import colors'
class Annotator:
if RANK in (-1, 0):
check_pil_font() # download TTF if necessary
# YOLOv5 Annotator for train/val mosaics and jpgs and detect/hub inference annotations
def __init__(self, im, line_width=None, font_size=None, font='Arial.ttf', pil=False, example='abc'):
assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to Annotator() input images.'
self.pil = pil or not is_ascii(example) or is_chinese(example)
if self.pil: # use PIL
self.im = im if isinstance(im, Image.Image) else Image.fromarray(im)
self.draw = ImageDraw.Draw(self.im)
self.font = check_pil_font(font='Arial.Unicode.ttf' if is_chinese(example) else font,
size=font_size or max(round(sum(self.im.size) / 2 * 0.035), 12))
else: # use cv2
self.im = im
self.lw = line_width or max(round(sum(im.shape) / 2 * 0.003), 2) # line width
def box_label(self, box, label='', color=(128, 128, 128), txt_color=(255, 255, 255)):
# Add one xyxy box to image with label
if self.pil or not is_ascii(label):
self.draw.rectangle(box, width=self.lw, outline=color) # box
if label:
w, h = self.font.getsize(label) # text width, height
outside = box[1] - h >= 0 # label fits outside box
self.draw.rectangle((box[0],
box[1] - h if outside else box[1],
box[0] + w + 1,
box[1] + 1 if outside else box[1] + h + 1), fill=color)
self.draw.text((box[0], box[1] - h if outside else box[1]), label, fill=txt_color, font=self.font)
else: # cv2
p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
cv2.rectangle(self.im, p1, p2, color, thickness=self.lw, lineType=cv2.LINE_AA)
if label:
tf = max(self.lw - 1, 1) # font thickness
w, h = cv2.getTextSize(label, 0, fontScale=self.lw / 3, thickness=tf)[0] # text width, height
outside = p1[1] - h - 3 >= 0 # label fits outside box
p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
cv2.rectangle(self.im, p1, p2, color, -1, cv2.LINE_AA) # filled
cv2.putText(self.im, label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2), 0, self.lw / 3, txt_color,
thickness=tf, lineType=cv2.LINE_AA)
def rectangle(self, xy, fill=None, outline=None, width=1):
# Add rectangle to image (PIL-only)
self.draw.rectangle(xy, fill, outline, width)
def text(self, xy, text, txt_color=(255, 255, 255)):
# Add text to image (PIL-only)
w, h = self.font.getsize(text) # text width, height
self.draw.text((xy[0], xy[1] - h + 1), text, fill=txt_color, font=self.font)
def result(self):
# Return annotated image as array
return np.asarray(self.im)

View File

@@ -16,7 +16,7 @@ import numpy as np
import cv2 import cv2
def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=False, scaleFill=False, scaleup=True, stride=32): def preprocess(im, new_shape, color=(114, 114, 114)):
# Resize and pad image while meeting stride-multiple constraints # Resize and pad image while meeting stride-multiple constraints
shape = im.shape[:2] # current shape [height, width] shape = im.shape[:2] # current shape [height, width]
if isinstance(new_shape, int): if isinstance(new_shape, int):
@@ -24,19 +24,12 @@ def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=False, scale
# Scale ratio (new / old) # Scale ratio (new / old)
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
if not scaleup: # only scale down, do not scale up (for better val mAP)
r = min(r, 1.0)
# Compute padding # Compute padding
ratio = r, r # width, height ratios ratio = r, r # width, height ratios
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
if auto: # minimum rectangle
dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding
elif scaleFill: # stretch
dw, dh = 0.0, 0.0
new_unpad = (new_shape[1], new_shape[0])
ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios
dw /= 2 # divide padding into 2 sides dw /= 2 # divide padding into 2 sides
dh /= 2 dh /= 2