图像分类优化

src/driverless_car/main.py

@@ -5,40 +5,38 @@
 from torchvision import datasets, transforms
 import matplotlib.pyplot as plt
 import numpy as np
-import time
-
-
-# 修复1：解决中文字体问题（改用英文显示，避免字体依赖）
-# 若需中文，可注释以下行并使用系统自带中文字体，如：
-# plt.rcParams["font.family"] = ["Microsoft YaHei", "SimSun", "Arial"]
-# plt.rcParams["axes.unicode_minus"] = False
-# --------------------------
-# 直接使用英文显示类别，避免中文字体问题
-classes = ('Airplane', 'Car', 'Bird', 'Cat', 'Deer', 'Dog', 'Frog', 'Horse', 'Ship', 'Truck')
+import os
+from PIL import Image
+import random
 
 # --------------------------
-# 修复2：设置Matplotlib后端为TkAgg，解决tostring_rgb报错
+# 1. 基础配置（解决可视化和字体问题）
 # --------------------------
 import matplotlib
 matplotlib.use('TkAgg')  # 更换后端，兼容PyCharm的可视化
 
+# 类别名称（对应CIFAR-10的10类）
+classes = ('Airplane', 'Car', 'Bird', 'Cat', 'Deer', 'Dog', 'Frog', 'Horse', 'Ship', 'Truck')
+
 # --------------------------
-# 1. 数据预处理与加载（模拟无人机采集的图像数据）
+# 2. 数据预处理（简化，减少计算量）
 # --------------------------
-transform = transforms.Compose([
+# 简化变换：移除数据增强（加快训练，牺牲一点泛化能力）
+basic_transform = transforms.Compose([
     transforms.Resize((32, 32)),
     transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
 ])
 
-train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
-test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
-
-train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
-test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
+# 加载CIFAR-10数据集（使用简化变换）
+train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=basic_transform)
+test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=basic_transform)
+# 增大批次大小，加快训练（根据内存调整，默认128）
+train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
+test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)
 
 # --------------------------
-# 2. 搭建轻量化CNN模型（无人机端深度学习模型）
+# 3. 搭建轻量化CNN模型（原模型，参数少，速度快）
 # --------------------------
 class DroneCNN(nn.Module):
     def __init__(self):
@@ -50,7 +48,7 @@ def __init__(self):
         self.fc1 = nn.Linear(128 * 4 * 4, 512)
         self.fc2 = nn.Linear(512, 10)
         self.relu = nn.ReLU()
-        self.dropout = nn.Dropout(0.5)
+        self.dropout = nn.Dropout(0.3)  # 降低dropout比例，加快计算
 
     def forward(self, x):
         x = self.pool(self.relu(self.conv1(x)))
@@ -62,148 +60,203 @@ def forward(self, x):
         x = self.fc2(x)
         return x
 
+# 初始化模型、损失函数、优化器
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = DroneCNN().to(device)
 criterion = nn.CrossEntropyLoss()
+# 使用SGD优化器（比Adam稍快，或保留Adam，影响不大）
 optimizer = optim.Adam(model.parameters(), lr=0.001)
 
 # --------------------------
-# 3. 可视化函数（修复非阻塞显示问题）
+# 4. 训练函数（大幅优化速度）
 # --------------------------
-def show_dataset_samples():
-    """显示数据集的样本图像（模拟无人机采集的图像）"""
-    data_iter = iter(train_loader)
-    images, labels = next(data_iter)
-    images = images / 2 + 0.5  # 反归一化
-
-    plt.figure(figsize=(10, 6))
-    for i in range(12):
-        plt.subplot(3, 4, i+1)
-        plt.imshow(np.transpose(images[i].numpy(), (1, 2, 0)))
-        plt.title(classes[labels[i]])
-        plt.axis('off')
-    plt.suptitle('Drone Collected Image Samples (CIFAR-10 Simulation)', fontsize=14)
-    plt.tight_layout()
-    plt.show(block=False)  # 保留非阻塞
-    plt.pause(0.1)  # 修复：添加pause，解决后端渲染问题
-    # 移除time.sleep，改用plt.pause更稳定
+# 全局开启交互模式，用于实时绘图
+plt.ion()
+fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))  # 训练曲线窗口
 
 def plot_training_curve(train_losses, train_accs, test_accs):
-    """绘制训练损失和准确率曲线"""
-    plt.figure(figsize=(12, 4))
+    """更新训练曲线窗口，减少绘图开销"""
+    ax1.clear()
+    ax2.clear()
     # 损失曲线
-    plt.subplot(1, 2, 1)
-    plt.plot(train_losses, label='Training Loss')
-    plt.xlabel('Iteration (Batch)')
-    plt.ylabel('Loss')
-    plt.title('Training Loss Change')
-    plt.legend()
-    plt.grid(True)
+    ax1.plot(train_losses, label='Training Loss', color='blue')
+    ax1.set_xlabel('Iteration (Batch)')
+    ax1.set_ylabel('Loss')
+    ax1.set_title('Training Loss Change')
+    ax1.legend()
+    ax1.grid(True)
     # 准确率曲线
-    plt.subplot(1, 2, 2)
-    plt.plot(train_accs, label='Training Accuracy')
-    plt.plot(test_accs, label='Test Accuracy')
-    plt.xlabel('Iteration (Batch)')
-    plt.ylabel('Accuracy (%)')
-    plt.title('Training/Test Accuracy Change')
-    plt.legend()
-    plt.grid(True)
-    plt.tight_layout()
-    plt.show(block=False)
-    plt.pause(0.1)  # 修复：添加pause
+    ax2.plot(train_accs, label='Training Accuracy', color='green')
+    ax2.plot(test_accs, label='Test Accuracy', color='red')
+    ax2.set_xlabel('Iteration (Batch)')
+    ax2.set_ylabel('Accuracy (%)')
+    ax2.set_title('Training/Test Accuracy Change')
+    ax2.legend()
+    ax2.grid(True)
+    plt.draw()
+    plt.pause(0.01)  # 减少暂停时间，加快绘图
 
-def show_predictions():
-    """显示模型的预测结果"""
+def calculate_test_acc_fast(test_loader, sample_batches=5):
+    """快速计算测试集准确率：只抽取少量批次，不遍历整个测试集"""
+    test_correct = 0
+    test_total = 0
     model.eval()
-    data_iter = iter(test_loader)
-    images, labels = next(data_iter)
-    images = images.to(device)
-    labels = labels.to(device)
-
-    outputs = model(images)
-    _, predicted = torch.max(outputs, 1)
-    images = images / 2 + 0.5  # 反归一化
-
-    plt.figure(figsize=(12, 8))
-    for i in range(16):
-        plt.subplot(4, 4, i+1)
-        plt.imshow(np.transpose(images[i].cpu().numpy(), (1, 2, 0)))
-        true_label = classes[labels[i]]
-        pred_label = classes[predicted[i]]
-        color = 'green' if true_label == pred_label else 'red'
-        plt.title(f'True: {true_label}\nPred: {pred_label}', color=color)
-        plt.axis('off')
-    plt.suptitle('Drone Image Classification Model Predictions', fontsize=14)
-    plt.tight_layout()
-    plt.show(block=False)
-    plt.pause(0.1)  # 修复：添加pause
+    with torch.no_grad():
+        # 只取前sample_batches个批次，大幅减少耗时
+        for i, (test_inputs, test_labels) in enumerate(test_loader):
+            if i >= sample_batches:
+                break
+            test_inputs, test_labels = test_inputs.to(device), test_labels.to(device)
+            test_outputs = model(test_inputs)
+            _, test_predicted = torch.max(test_outputs.data, 1)
+            test_total += test_labels.size(0)
+            test_correct += (test_predicted == test_labels).sum().item()
+    model.train()
+    if test_total == 0:
+        return 0.0
+    return 100 * test_correct / test_total
 
-# --------------------------
-# 4. 训练模型并实时可视化
-# --------------------------
-def train_model(epochs=2):
+def train_model(epochs=1):  # 减少训练轮数，默认1轮
     train_losses = []
     train_accs = []
     test_accs = []
     model.train()
 
-    # 显示数据集样本
-    show_dataset_samples()
-
     for epoch in range(epochs):
         running_loss = 0.0
         correct = 0
         total = 0
+        # 每200个批次更新一次曲线（原100，减少更新频率）
+        update_interval = 200
         for i, (inputs, labels) in enumerate(train_loader):
             inputs, labels = inputs.to(device), labels.to(device)
 
+            # 前向传播+反向传播+优化
             optimizer.zero_grad()
             outputs = model(inputs)
             loss = criterion(outputs, labels)
             loss.backward()
             optimizer.step()
 
+            # 统计指标
             running_loss += loss.item()
             _, predicted = torch.max(outputs.data, 1)
             total += labels.size(0)
             correct += (predicted == labels).sum().item()
 
-            # 每100个批次记录并可视化
-            if i % 100 == 99:
-                train_loss = running_loss / 100
+            # 降低更新频率，减少计算和绘图开销
+            if i % update_interval == update_interval - 1:
+                train_loss = running_loss / update_interval
                 train_acc = 100 * correct / total
                 train_losses.append(train_loss)
                 train_accs.append(train_acc)
 
-                # 计算测试集准确率
-                test_correct = 0
-                test_total = 0
-                with torch.no_grad():
-                    for test_inputs, test_labels in test_loader:
-                        test_inputs, test_labels = test_inputs.to(device), test_labels.to(device)
-                        test_outputs = model(test_inputs)
-                        _, test_predicted = torch.max(test_outputs.data, 1)
-                        test_total += test_labels.size(0)
-                        test_correct += (test_predicted == test_labels).sum().item()
-                test_acc = 100 * test_correct / test_total
+                # 快速计算测试集准确率（只取5个批次）
+                test_acc = calculate_test_acc_fast(test_loader, sample_batches=5)
                 test_accs.append(test_acc)
 
                 print(f'Epoch {epoch+1}, Batch {i+1} | Loss: {train_loss:.3f} | Train Acc: {train_acc:.2f}% | Test Acc: {test_acc:.2f}%')
                 running_loss = 0.0
                 correct = 0
                 total = 0
 
-                # 实时绘制曲线
                 plot_training_curve(train_losses, train_accs, test_accs)
 
-    # 显示预测结果
-    show_predictions()
-    # 修复：最后添加plt.show(block=True)，防止窗口闪退
+    # 保存模型
+    torch.save(model.state_dict(), 'drone_model.pth')
+    print('模型已保存为drone_model.pth')
+    model.eval()
+    plt.ioff()  # 关闭交互模式
+    plt.show(block=False)
+
+# --------------------------
+# 5. 模拟无人机实时图像输入（优化推理速度）
+# --------------------------
+def load_drone_images(folder_path):
+    """读取本地文件夹中的图像，模拟无人机采集的图像流"""
+    img_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.gif']
+    img_paths = []
+    for file in os.listdir(folder_path):
+        if os.path.splitext(file)[1].lower() in img_extensions:
+            img_paths.append(os.path.join(folder_path, file))
+    if not img_paths:
+        raise ValueError(f'文件夹{folder_path}中未找到任何图像文件！')
+    return img_paths
+
+def preprocess_image(img_path):
+    """预处理单张图像，简化操作"""
+    # 读取图像（RGB模式）
+    img = Image.open(img_path).convert('RGB')
+    original_img = img.copy()
+    # 预处理
+    img = basic_transform(img)
+    # 添加batch维度
+    img = torch.unsqueeze(img, 0)
+    return original_img, img.to(device)
+
+def drone_real_time_inference(folder_path, delay=0.1):  # 降低延迟，默认0.1秒
+    """模拟无人机实时图像输入，优化推理速度"""
+    print(f'\n开始模拟无人机实时图像流（读取文件夹：{folder_path}），每{delay}秒处理一张图像...')
+    img_paths = load_drone_images(folder_path)
+    # 创建单个显示窗口，减少窗口创建开销
+    fig, ax = plt.subplots(figsize=(6, 4))  # 缩小窗口，加快绘图
+    plt.ion()
+
+    for img_path in img_paths:
+        try:
+            # 预处理图像
+            original_img, input_tensor = preprocess_image(img_path)
+
+            # 模型预测（简化，减少计算）
+            with torch.no_grad():
+                outputs = model(input_tensor)
+                probabilities = torch.softmax(outputs, dim=1)
+                pred_idx = torch.argmax(probabilities, dim=1).item()
+                pred_class = classes[pred_idx]
+                pred_conf = probabilities[0][pred_idx].item() * 100
+
+            # 可视化优化：只更新图像和文本，不重建窗口
+            ax.clear()
+            ax.imshow(original_img)
+            ax.axis('off')
+            # 简化文本显示，减少渲染开销
+            text = f'{pred_class} ({pred_conf:.1f}%)'
+            ax.text(5, 5, text, fontsize=10, color='red',
+                    bbox=dict(facecolor='white', alpha=0.7, edgecolor='none'))
+            ax.set_title('Drone Real-Time View', fontsize=12)
+            plt.draw()
+            plt.pause(delay)
+
+            # 简化控制台输出
+            print(f'图像：{os.path.basename(img_path)} → {pred_class} ({pred_conf:.1f}%)')
+
+        except Exception as e:
+            print(f'处理图像{img_path}时出错：{e}')
+            continue
+
+    # 结束后保持窗口
+    plt.ioff()
+    ax.text(0.5, 0.5, 'Done!', fontsize=14, ha='center', va='center',
+            transform=ax.transAxes, bbox=dict(facecolor='red', alpha=0.8))
+    plt.draw()
     plt.show(block=True)
-    print('Training Finished!')
 
 # --------------------------
-# 运行主程序
+# 主程序运行
 # --------------------------
 if __name__ == '__main__':
-    train_model(epochs=2)
+    # 第一步：训练模型（优化后速度大幅提升）
+    train_model(epochs=1)  # 可改为2轮，仍比原来快很多
+
+    # 第二步：加载模型（可选）
+    # model.load_state_dict(torch.load('drone_model.pth', map_location=device))
+    # model.eval()
+    # print('模型已加载')
+
+    # 第三步：模拟无人机实时图像输入
+    drone_image_folder = r"C:\Users\hyq52\Desktop\P1\potoh"
+    if not os.path.exists(drone_image_folder):
+        os.makedirs(drone_image_folder)
+        print(f'已创建文件夹：{drone_image_folder}，请放入测试图片后重新运行！')
+    else:
+        drone_real_time_inference(drone_image_folder, delay=0.1)  # 低延迟，快速播放