102 lines
2.7 KiB
Plaintext
102 lines
2.7 KiB
Plaintext
{
|
||
"cells": [
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 9,
|
||
"metadata": {
|
||
"collapsed": true
|
||
},
|
||
"outputs": [],
|
||
"source": [
|
||
"import torch\n",
|
||
"import torch.nn as nn\n",
|
||
"import torchvision.transforms as transforms\n",
|
||
"import torchvision.models as models\n",
|
||
"from PIL import Image\n",
|
||
"\n",
|
||
"# 设置设备\n",
|
||
"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
|
||
"\n",
|
||
"# 数据预处理和标准化\n",
|
||
"transform = transforms.Compose([\n",
|
||
" transforms.Resize((224, 224)), # 调整图像大小为模型要求的大小\n",
|
||
" transforms.ToTensor(),\n",
|
||
" transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
|
||
"])\n",
|
||
"\n",
|
||
"# 加载已训练的MobileNetV2模型\n",
|
||
"model = models.mobilenet_v2(pretrained=False) # 创建一个新的MobileNetV2模型\n",
|
||
"num_classes = 12 # 假设您有12个类别，根据您的实际数据集进行调整\n",
|
||
"model.classifier[1] = nn.Linear(in_features=1280, out_features=num_classes)\n",
|
||
"model.load_state_dict(torch.load(\"./model/1/epochs10 96.13.pt\"))\n",
|
||
"model = model.to(device)\n",
|
||
"model.eval()\n",
|
||
"\n",
|
||
"class_names = [\"battery\",\"brick\",\"bottle\",\"butt\",\"cans\",\"carrot_piece\",\"fruits\",\"leaf\",\"nothing\",\"paper\",\"potato\",\"vegetable\"]\n",
|
||
"\n",
|
||
"def classify_image(image_path):\n",
|
||
" image = Image.open(image_path)\n",
|
||
" image = transform(image).unsqueeze(0) # 加载图像并进行预处理，添加一个批次维度\n",
|
||
" image = image.to(device)\n",
|
||
"\n",
|
||
" with torch.no_grad():\n",
|
||
" outputs = model(image)\n",
|
||
" _, predicted = torch.max(outputs.data, 1)\n",
|
||
"\n",
|
||
" class_idx = predicted.item()\n",
|
||
" class_name = class_names[class_idx]\n",
|
||
" return class_name\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 10,
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"bottle\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"print(classify_image(\"./test_image/1690813083599.jpg\"))"
|
||
],
|
||
"metadata": {
|
||
"collapsed": false
|
||
}
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 3,
|
||
"outputs": [],
|
||
"source": [],
|
||
"metadata": {
|
||
"collapsed": false
|
||
}
|
||
}
|
||
],
|
||
"metadata": {
|
||
"kernelspec": {
|
||
"display_name": "Python 3",
|
||
"language": "python",
|
||
"name": "python3"
|
||
},
|
||
"language_info": {
|
||
"codemirror_mode": {
|
||
"name": "ipython",
|
||
"version": 2
|
||
},
|
||
"file_extension": ".py",
|
||
"mimetype": "text/x-python",
|
||
"name": "python",
|
||
"nbconvert_exporter": "python",
|
||
"pygments_lexer": "ipython2",
|
||
"version": "2.7.6"
|
||
}
|
||
},
|
||
"nbformat": 4,
|
||
"nbformat_minor": 0
|
||
}
|