本專案以 YOLOV8 加上 QT5 視窗程式,即時偵測照片中的物件。請先下載 ui_mainwindow.ui,置於專案下的 ui 目錄。
ModelThread
載入模型可能需要一段時間,所以使用新執行緒載入。
from PyQt5.QtCore import QThread, pyqtSignal from ultralytics import YOLO class ModelThread(QThread): callback = pyqtSignal(object) def __init__(self, parent=None): super().__init__(parent) def run(self): model=YOLO("./yolov8n.pt") self.callback.emit(model)
PictureThread
此段是載入檔案中的圖片,如果檔案很多,載入時間也會很久,所以也交由新執行緒載入。
import os from PyQt5.QtCore import QThread, pyqtSignal from PyQt5.QtGui import QPixmap class PictureThread(QThread): callback=pyqtSignal(object) def __init__(self, path, parent=None): super().__init__(parent) self.path=path self.runFlag=True def run(self): ls = os.listdir(self.path) files = [] for l in ls: ll = l.lower() if ll.endswith('.jpg') or ll.endswith('.png'): files.append(os.path.join(self.path, l)) index=0
total=len(files) while index < total and self.runFlag: pix = QPixmap(files[index]) pix = pix.scaled(400, 300) pix.tag=files[index] self.callback.emit(pix) index+=1 QThread.msleep(10)#要停一下, 否則 UI 無法立即顯示
DetectThread
此段為辨識圖片代碼。
import platform from PIL import Image, ImageFont, ImageDraw from PyQt5.QtCore import QThread, pyqtSignal import numpy as np import cv2 class DetectThread(QThread): callback = pyqtSignal(object) def __init__(self, model, file, parent=None): super().__init__(parent) self.model=model self.file=file def text(self, img, text, xy=(0, 0), color=(0, 0, 0), size=12): pil = Image.fromarray(img) s = platform.system() if s == "Linux": font = ImageFont.truetype( '/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc', size) elif s == "Darwin": font = ImageFont.truetype('....', size) else: font = ImageFont.truetype('simsun.ttc', size) ImageDraw.Draw(pil).text(xy, text, font=font, fill=color) return np.asarray(pil) def run(self): img = cv2.imdecode( np.fromfile(self.file, dtype=np.uint8), cv2.IMREAD_COLOR)[:,:,::-1].copy() results = self.model.predict(img)[0] names = [results.names[int(i.cpu().numpy())] for i in results.boxes.cls] boxes = results.boxes.xyxy for box, name in zip(boxes, names): box = box.cpu().numpy() x1 = int(box[0]) y1 = int(box[1]) x2 = int(box[2]) y2 = int(box[3]) img = cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 3) img = self.text(img, name, (x1, y1 - 20), color=(0, 0, 255), size=100) self.callback.emit(img)
主程式
主程式代碼如下
import sys import time from PyQt5.QtCore import QSize from PyQt5.QtGui import QIcon, QPixmap, QImage from PyQt5.QtWidgets import QMainWindow, QApplication, QPushButton, QListWidgetItem, QFileDialog from DetectThread import DetectThread from LoadModelThread import LoadModelThread from PictureThread import PictureThread from ui.ui_mainwindow import Ui_MainWindow class YoloV8App(QMainWindow, Ui_MainWindow): def __init__(self, parent=None): super().__init__(parent) self.setupUi(self) self.resize(1920,1080) self.btnPath.clicked.connect(self.btnPath_click) self.lblPath.setText("D:/pictures/primitive/2020/20200126_Philippines_day5_Boracay") self.path=self.lblPath.text() self.lblStatus.setText("載入模型中, 請稍後.....") self.modelThread=ModelThread() self.modelThread.callback.connect(self.ModelThreadCallback) self.modelThread.start() def btnPath_click(self): path=QFileDialog.getExistingDirectory() if path!='': self.path=path.replace("\\","/") self.lblPath.setText(self.path) self.listWidget.clear() self.pictureThread.callback=False self.pictureThread=PictureThread(self.path) self.pictureThread.callback.connect(self.pictureThreadCallback) self.pictureThread.start() def ModelThreadCallback(self, model): self.model=model self.lblStatus.setText("") self.pictureThread=PictureThread(self.path) self.pictureThread.callback.connect(self.showPicture) self.pictureThread.start() def pictureThreadCallback(self, pix): btn=QPushButton() btn.setIcon(QIcon(pix)) btn.setIconSize(QSize(400,300)) btn.tag=pix.tag btn.clicked.connect(self.btn_click) item=QListWidgetItem() item.setSizeHint(QSize(400,300)) self.listWidget.addItem(item) self.listWidget.setItemWidget(item, btn) def btn_click(self): btn=self.sender() self.lblStatus.setText('yolov8 辨識中.....') self.t1=time.time() self.detectThread=DetectThread(self.model, btn.tag) self.detectThread.callback.connect(self.detectThreadCallback) self.detectThread.start() def detectThreadCallback(self, img): self.t2=time.time() self.lblStatus.setText(f'辨識時間 {self.t2-self.t1:.7f}秒') pix=QPixmap( QImage(img, img.shape[1], img.shape[0], img.shape[1]*3, QImage.Format_RGB888 ) ) pr=pix.width()/pix.height() lr=self.lblImg.width()/self.lblImg.height() if pr>lr: pix=pix.scaled(self.lblImg.width(), self.lblImg.width()/pr) else: pix=pix.scaled(self.lblImg.height()*lr, self.lblImg.height()) self.lblImg.setPixmap(pix) def closeEvent(self, evnet): if self.pictureThread is not None: self.pictureThread.runFlag=False app=QApplication(sys.argv) mainWindow=YoloV8App() mainWindow.showMaximized() app.exec()