请写出完整的python代码,实现利用ylov5对包含任意数量香蕉的照片进行计数,并且利用pyqt设计一个界面,实现从文件夹选择图像,显示所选择的图像,点击按钮开始计数,输出计数结果和结果图像四个功能
时间: 2024-04-11 20:30:06 浏览: 10
import sys
import os
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QVBoxLayout, QFileDialog, QPushButton, QMessageBox
from PyQt5.QtGui import QPixmap
import cv2
import torch
import numpy as np
from yolov5 import detect
class BananaCounter(QMainWindow):
def __init__(self):
super().__init__()
self.setWindowTitle("Banana Counter")
self.label_image = QLabel(self)
self.label_image.setAlignment(Qt.AlignCenter)
self.btn_select_image = QPushButton("Select Image", self)
self.btn_select_image.clicked.connect(self.select_image)
self.btn_count_bananas = QPushButton("Count Bananas", self)
self.btn_count_bananas.clicked.connect(self.count_bananas)
layout = QVBoxLayout()
layout.addWidget(self.label_image)
layout.addWidget(self.btn_select_image)
layout.addWidget(self.btn_count_bananas)
widget = QWidget(self)
widget.setLayout(layout)
self.setCentralWidget(widget)
def select_image(self):
file_dialog = QFileDialog()
file_dialog.setNameFilter("Image files (*.jpg *.png)")
if file_dialog.exec_():
file_path = file_dialog.selectedFiles()[0]
self.show_image(file_path)
def show_image(self, file_path):
pixmap = QPixmap(file_path)
self.label_image.setPixmap(pixmap.scaledToWidth(self.label_image.width()))
def count_bananas(self):
file_path = self.label_image.pixmap().toImage().save("temp_image.jpg")
result = detect("temp_image.jpg")
os.remove("temp_image.jpg")
if result is not None:
num_bananas = len(result)
message = f"Number of bananas: {num_bananas}"
QMessageBox.information(self, "Banana Count Result", message)
# Display the result image with bounding boxes
image_with_boxes = draw_boxes(result, cv2.imread(file_path))
cv2.imshow("Banana Count Result", image_with_boxes)
cv2.waitKey(0)
else:
QMessageBox.warning(self, "Banana Count Result", "No bananas found in the image.")
def draw_boxes(result, image):
for box in result:
x1, y1, x2, y2, _ = box
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
return image
def detect(image_path):
weights = 'path_to_yolov5_weights' # Replace with your yolov5 weights path
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torch.hub.load('ultralytics/yolov5', 'custom', path=weights).to(device).eval()
img = cv2.imread(image_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
results = model(img)
if results.pandas().xyxy[0] is not None:
return results.pandas().xyxy[0].values.tolist()
else:
return None
if __name__ == "__main__":
app = QApplication(sys.argv)
window = BananaCounter()
window.show()
sys.exit(app.exec_())