基于pytorch水果图像识别系统,数据集使用Fruits 360,并将训练好的模型应用到图片识别界面中,使用GPU加速训练模型,具体代码实现
时间: 2024-05-16 16:13:36 浏览: 183
好的,我可以为你提供一个基于PyTorch实现的水果图像识别系统的代码示例。这个示例将使用Fruits 360数据集,并且将训练好的模型应用到一个基于PyQt5的图片识别界面中。同时,我们将使用GPU来加速训练模型,以便更快地获得准确的识别结果。
首先,我们需要安装PyTorch和PyQt5库,可以使用以下命令:
```
pip install torch
pip install PyQt5
```
接下来,我们将定义一个用于加载Fruits 360数据集的函数。这个函数使用torchvision库来加载数据,并将数据集分成训练集和测试集。
```python
import torch
import torchvision
import torchvision.transforms as transforms
def load_data():
transform = transforms.Compose(
[transforms.Resize((64, 64)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
trainset = torchvision.datasets.ImageFolder(root='./fruits-360/Training',
transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=32,
shuffle=True, num_workers=2)
testset = torchvision.datasets.ImageFolder(root='./fruits-360/Test',
transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=32,
shuffle=False, num_workers=2)
return trainloader, testloader
```
然后,我们将定义一个用于训练模型的函数。这个函数将使用PyTorch的GPU加速来加速训练过程。
```python
def train_model(trainloader):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = Net()
net = net.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
for epoch in range(10): # loop over the dataset multiple times
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
inputs, labels = inputs.to(device), labels.to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if i % 200 == 199: # print every 200 mini-batches
print('[%d, %5d] loss: %.3f' %
(epoch + 1, i + 1, running_loss / 200))
running_loss = 0.0
print('Finished Training')
return net
```
接着,我们将定义一个用于测试模型的函数。这个函数将使用测试集上的图像来评估模型的准确率。
```python
def test_model(net, testloader):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
correct = 0
total = 0
with torch.no_grad():
for data in testloader:
images, labels = data
images, labels = images.to(device), labels.to(device)
outputs = net(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('Accuracy of the network on the 10000 test images: %d %%' % (
100 * correct / total))
```
最后,我们将定义一个用于应用训练好的模型的函数。这个函数将加载训练好的模型,并使用PyQt5来实现一个简单的GUI界面,以便我们可以将图像加载到系统中,并使用训练好的模型来识别它们。
```python
from PyQt5.QtWidgets import *
from PyQt5.QtGui import *
from PyQt5.QtCore import *
from PIL import Image
import numpy as np
class App(QWidget):
def __init__(self):
super().__init__()
self.title = 'Fruit Recognition'
self.left = 10
self.top = 10
self.width = 640
self.height = 480
self.initUI()
def initUI(self):
self.setWindowTitle(self.title)
self.setGeometry(self.left, self.top, self.width, self.height)
# create a label
self.label = QLabel(self)
self.label.setGeometry(QRect(30, 30, 400, 400))
self.label.setAlignment(Qt.AlignCenter)
# create a button
button = QPushButton('Open', self)
button.setGeometry(QRect(500, 30, 100, 30))
button.clicked.connect(self.open_image)
self.show()
def open_image(self):
options = QFileDialog.Options()
options |= QFileDialog.DontUseNativeDialog
file_name, _ = QFileDialog.getOpenFileName(self, "Open Image", "",
"Images (*.png *.xpm *.jpg *.bmp);;All Files (*)",
options=options)
if file_name:
image = Image.open(file_name)
image = image.resize((64, 64))
image = np.array(image)
image = image.transpose((2, 0, 1))
image = image / 255
image = torch.from_numpy(image).type(torch.FloatTensor)
image = image.unsqueeze(0)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = Net()
net = net.to(device)
net.load_state_dict(torch.load('fruits_model.pth'))
outputs = net(image)
_, predicted = torch.max(outputs.data, 1)
self.label.setText('This is a ' + classes[predicted.item()] + '!')
self.label.setPixmap(QPixmap(file_name).scaled(400, 400, Qt.KeepAspectRatio))
self.label.setAlignment(Qt.AlignCenter)
if __name__ == '__main__':
classes = ('Apple Braeburn', 'Apple Golden 1', 'Apple Golden 2', 'Apple Golden 3', 'Apple Granny Smith', 'Apple Red 1', 'Apple Red 2', 'Apple Red 3', 'Apple Red Delicious', 'Apple Red Yellow 1', 'Apple Red Yellow 2', 'Apricot', 'Avocado', 'Banana', 'Beetroot', 'Blueberry', 'Cactus fruit', 'Cantaloupe 1', 'Cantaloupe 2', 'Carambula', 'Cauliflower', 'Cherry 1', 'Cherry 2', 'Cherry Rainier', 'Cherry Wax Black', 'Cherry Wax Red', 'Cherry Wax Yellow', 'Chestnut', 'Clementine', 'Cocos', 'Dates', 'Eggplant', 'Fig', 'Ginger Root', 'Granadilla', 'Grape Blue', 'Grape Pink', 'Grape White', 'Grape White 2', 'Grape White 3', 'Grape White 4', 'Grapefruit Pink', 'Grapefruit White', 'Guava', 'Hazelnut', 'Huckleberry', 'Kaki', 'Kiwi', 'Kohlrabi', 'Kumquats', 'Lemon', 'Lemon Meyer', 'Limes', 'Lychee', 'Mandarine', 'Mango', 'Mangostan', 'Maracuja', 'Melon Piel de Sapo', 'Mulberry', 'Nectarine', 'Orange', 'Papaya', 'Passion Fruit', 'Peach', 'Peach Flat', 'Pear', 'Pear Abate', 'Pear Monster', 'Pear Williams', 'Pepino', 'Pepper Green', 'Pepper Red', 'Pepper Yellow', 'Physalis', 'Physalis with Husk', 'Pineapple', 'Pineapple Mini', 'Pitahaya Red', 'Plum', 'Plum 2', 'Plum 3', 'Pomegranate', 'Pomelo Sweetie', 'Potato Red', 'Potato Red Washed', 'Potato Sweet', 'Potato White', 'Quince', 'Rambutan', 'Raspberry', 'Redcurrant', 'Salak', 'Strawberry', 'Tamarillo', 'Tangelo', 'Tomato 1', 'Tomato 2', 'Tomato 3', 'Tomato 4', 'Tomato Cherry Red', 'Tomato Maroon', 'Tomato Yellow', 'Walnut')
trainloader, testloader = load_data()
net = train_model(trainloader)
test_model(net, testloader)
torch.save(net.state_dict(), 'fruits_model.pth')
app = QApplication(sys.argv)
ex = App()
sys.exit(app.exec_())
```
注意,这个示例中我们使用了一个名为Net的神经网络模型,你可以根据需要进行替换。
至此,我们已经完成了一个基于PyTorch的水果图像识别系统的实现。你可以使用这个示例作为起点,根据需要进行修改和扩展。
阅读全文