109 lines
3.6 KiB
Python
109 lines
3.6 KiB
Python
|
import time
|
||
|
|
from pathlib import Path
|
||
|
|
import datetime
|
||
|
|
|
||
|
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
import torch.backends.cudnn as cudnn
|
||
|
|
import os
|
||
|
|
|
||
|
|
from read_data import LoadImages, LoadStreams
|
||
|
|
from tools.draw_chinese import cv2ImgAddText
|
||
|
|
|
||
|
|
class PCBDetection():
|
||
|
|
|
||
|
|
|
||
|
|
def __init__(self, video_path=None):
|
||
|
|
self.model = torch.hub.load((os.getcwd()) + "/algorithm/yolov5", 'custom', source='local', path='./weight/pcb.pt', force_reload=True)
|
||
|
|
self.classes = self.model.names
|
||
|
|
self.imgsz = 640
|
||
|
|
self.stride = self.model.stride
|
||
|
|
# print( self.stride)
|
||
|
|
self.frame = [None]
|
||
|
|
|
||
|
|
if video_path is not None:
|
||
|
|
self.video_name = video_path
|
||
|
|
else:
|
||
|
|
self.video_name = 'vid2.mp4' # A default video file
|
||
|
|
|
||
|
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
|
|
self.dataset = LoadImages(self.video_name, img_size=self.imgsz, stride = self.stride)
|
||
|
|
|
||
|
|
def use_webcam(self, source):
|
||
|
|
# self.dataset.release() # Release any existing video capture
|
||
|
|
# self.cap = cv2.VideoCapture(0) # Open default webcam
|
||
|
|
# print('use_webcam')
|
||
|
|
self.source = source
|
||
|
|
cudnn.benchmark = True
|
||
|
|
# self.dataset = LoadStreams(source, img_size=self.imgsz)
|
||
|
|
self.dataset = LoadStreams(source)
|
||
|
|
|
||
|
|
def class_to_label(self, x):
|
||
|
|
return self.classes[int(x)]
|
||
|
|
|
||
|
|
def get_frame(self):
|
||
|
|
|
||
|
|
chinese_name = ['漏孔', '鼠牙洞', '开路', '短路', '毛刺', '杂铜']
|
||
|
|
i = 0
|
||
|
|
|
||
|
|
for im0s in self.dataset:
|
||
|
|
# print(self.dataset.mode)
|
||
|
|
# print(self.dataset)
|
||
|
|
if self.dataset.mode == 'stream':
|
||
|
|
img = im0s[0].copy()
|
||
|
|
else:
|
||
|
|
img = im0s.copy()
|
||
|
|
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
||
|
|
results = self.model(img, size=640)
|
||
|
|
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
||
|
|
|
||
|
|
# Loop through each detected object and count the people
|
||
|
|
accuracy = 0
|
||
|
|
num_problem = len(results.xyxy[0])
|
||
|
|
bgr = (0, 255, 0)
|
||
|
|
|
||
|
|
|
||
|
|
for obj in results.xyxy[0]:
|
||
|
|
|
||
|
|
xmin, ymin, xmax, ymax = map(int, obj[:4])
|
||
|
|
|
||
|
|
accuracy = obj[4]
|
||
|
|
|
||
|
|
c = int(obj[-1])
|
||
|
|
|
||
|
|
|
||
|
|
if self.classes[c] == 'missing_hole':
|
||
|
|
color = (255, 200, 90)
|
||
|
|
elif self.classes[c] == 'mouse_bite':
|
||
|
|
color = (0, 0, 255)
|
||
|
|
elif self.classes[c] == 'caopen_circuitr':
|
||
|
|
color = (0, 255, 0)
|
||
|
|
elif self.classes[c] == 'short':
|
||
|
|
color = (50, 50, 50)
|
||
|
|
elif self.classes[c] == 'spur':
|
||
|
|
color = (255, 0, 0)
|
||
|
|
elif self.classes[c] == 'spurious_copper':
|
||
|
|
color = (0, 0, 0)
|
||
|
|
|
||
|
|
cv2.rectangle(img, (xmin, ymin), (xmax, ymax), color, 2)
|
||
|
|
|
||
|
|
|
||
|
|
img = cv2ImgAddText(img,
|
||
|
|
f'{chinese_name[c]}',
|
||
|
|
xmax + 2,
|
||
|
|
ymin - 1,
|
||
|
|
(0, 250, 0),
|
||
|
|
20,)
|
||
|
|
# cv2.putText(img, f"{self.classes[c]}, {round(float(accuracy), 2)}", (xmin, ymin),
|
||
|
|
# cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
|
||
|
|
|
||
|
|
ret, jpeg = cv2.imencode(".jpg", img)
|
||
|
|
|
||
|
|
|
||
|
|
resText=f'PCB检测到{num_problem}个缺陷'
|
||
|
|
# print(num_people)
|
||
|
|
i = i+1
|
||
|
|
return jpeg.tobytes(), resText
|
||
|
|
|