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

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 json
import stat
import cv2
from StreamManagerApi import *
import time
import numpy as np
from utils import *
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',
'motley', 'dirty_spot']
@@ -40,7 +39,6 @@ if __name__ == '__main__':
if ret != 0:
print("Failed to init Stream manager, ret=%s" % str(ret))
exit()
start = time.time()
# create streams by pipeline config file
with open("./pipeline/AlDefectDetection.pipeline", 'rb') as f:
pipelineStr = f.read()
@@ -95,8 +93,6 @@ if __name__ == '__main__':
with open(ori_img_path, 'rb') as f:
dataInput.data = f.read()
annotator = Annotator(ori_img, line_width=3, example=str(names))
# Inputs data to a specified stream based on streamName.
streamName = b'classification+detection'
inPluginId = 0
@@ -135,9 +131,8 @@ if __name__ == '__main__':
f.write(('%g ' * len(line)).rstrip() % line + '\n')
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)
TESTIMGS += 1
######################################################################################

View File

@@ -15,16 +15,15 @@
import os
import json
import stat
import glob
import cv2
from StreamManagerApi import *
import time
from StreamManagerApi import StreamManagerApi, MxDataInput
import numpy as np
from utils import *
from plots import Annotator, colors
from utils import preprocess, scale_coords, xyxy2xywh
from plots import box_label, colors
names = ['non_conduct', 'abrasion_mark', 'corner_leak', 'orange_peel', 'leak', 'jet_flow', 'paint_bubble', 'pit',
'motley', 'dirty_spot']
import glob
if __name__ == '__main__':
MODES = stat.S_IWUSR | stat.S_IRUSR
@@ -81,7 +80,7 @@ if __name__ == '__main__':
r = img_size / max(h0, w0) # ratio
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_path = PRE_IMG_PATH + item
@@ -97,7 +96,6 @@ if __name__ == '__main__':
dataInput = MxDataInput()
with open(pre_img_path, 'rb') as f:
dataInput.data = f.read()
annotator = Annotator(ori_img, line_width=3, example=str(names))
# Inputs data to a specified stream based on streamName.
streamName = b'classification+detection'
@@ -137,9 +135,8 @@ if __name__ == '__main__':
f.write(('%g ' * len(line)).rstrip() % line + '\n')
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)
TESTIMGS += 1
######################################################################################

View File

@@ -15,10 +15,10 @@
import os
import json
import cv2
from StreamManagerApi import *
from plots import Annotator, colors
from utils import *
from StreamManagerApi import StreamManagerApi, MxDataInput
import numpy as np
from plots import box_label, colors
from utils import preprocess, scale_coords, xyxy2xywh
names = ['non_conduct', 'abrasion_mark', 'corner_leak', 'orange_peel', 'leak', 'jet_flow', 'paint_bubble', 'pit',
'motley', 'dirty_spot']
@@ -40,39 +40,38 @@ if __name__ == '__main__':
# Construct the input of the stream
dataInput = MxDataInput()
ori_img_path = "test.jpg"
if not os.path.exists(ori_img_path):
ORI_IMG_PATH = "test.jpg"
if not os.path.exists(ORI_IMG_PATH):
print("The test image does not exist.")
exit()
if os.path.getsize(ori_img_path) == 0:
if os.path.getsize(ORI_IMG_PATH) == 0:
print("Error!The test image is empty.")
exit()
# read image
ori_img = cv2.imread(ori_img_path)
ori_img = cv2.imread(ORI_IMG_PATH)
h0, w0 = ori_img.shape[:2]
img_size = 640
r = img_size / max(h0, w0) # ratio
r = 640 / max(h0, w0) # ratio
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_path = "pre_" + ori_img_path
cv2.imwrite(pre_img_path, pre_img)
PRE_IMG_PATH = "pre_" + ORI_IMG_PATH
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()
# Inputs data to a specified stream based on streamName.
streamName = b'classification+detection'
inPluginId = 0
uniqueId = streamManagerApi.SendDataWithUniqueId(streamName, inPluginId, dataInput)
STREAMNAME = b'classification+detection'
INPLUGINID = 0
uniqueId = streamManagerApi.SendDataWithUniqueId(STREAMNAME, INPLUGINID, dataInput)
if uniqueId < 0:
print("Failed to send data to stream.")
exit()
# 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:
print("GetResultWithUniqueId error. errorCode=%d, errorMsg=%s" % (
inferResult.errorCode, inferResult.data.decode()))
@@ -81,7 +80,6 @@ if __name__ == '__main__':
results = json.loads(inferResult.data.decode())
gn = np.array(ori_img.shape)[[1, 0, 1, 0]]
annotator = Annotator(ori_img, line_width=3, example=str(names))
bboxes = []
classVecs = []
# 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
print(classVec)
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)
ap = np.trapz(np.interp(x, mrec, mpre), x) # integrate
return ap
@@ -266,10 +265,8 @@ def map_cac(opt):
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
ap = map_for_all_classes(*stats)
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: {map50}")
print(f"mAP0.5:0.95: {map}")
print(f"mAP0.5: {ap50.mean()}")
print(f"mAP0.5:0.95: {ap.mean()}")
def parse_opt():

View File

@@ -15,103 +15,24 @@
import platform
import os
from pathlib import Path
import re
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import cv2
# Settings
RANK = int(os.getenv('RANK', -1))
FONT = 'Arial.ttf'
import re
colors = [(255, 56, 56), (255, 157, 151), (255, 112, 31), (255, 178, 29), (207, 210, 49), (72, 249, 10), (146, 204, 23),
(61, 219, 134), (26, 147, 52), (0, 212, 187)]
def is_ascii(s=''):
# Is string composed of all ASCII (no UTF) characters? (note str().isascii() introduced in python 3.7)
s = str(s) # convert list, tuple, None, etc. to str
return len(s.encode().decode('ascii', 'ignore')) == len(s)
def is_chinese(s='人工智能'):
# Is string composed of any Chinese characters?
return True if re.search('[\u4e00-\u9fff]', str(s)) else False
def check_pil_font(font=FONT, size=10):
# 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
def box_label(im, box, label='', color=(128, 128, 128), lw=3):
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)
cv2.rectangle(im, p1, p2, color, thickness=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
tf = max(lw - 1, 1) # font thickness
w, h = cv2.getTextSize(label, 0, fontScale=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,
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 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)
return np.asarray(im)

View File

@@ -16,7 +16,7 @@ import numpy as np
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
shape = im.shape[:2] # current shape [height, width]
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)
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
ratio = r, r # width, height ratios
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
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
dh /= 2