训练代码完善

.gitignore

@@ -1,3 +1,4 @@
 train_data
 model
+test_image
 .idea

test.ipynb

@@ -0,0 +1,101 @@
+{
+ "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
+}