数据集

百度AI下载的颜色块数据集，一共7种颜色，不过都是合在一起的

通过将图片转化到HSV空间，利用阈值提取不同颜色的图片，存放到不同文件夹，制作分类数据集。

根据表格数据，提取颜色代码

import numpy as np
import collections
import os
import cv2
# 定义字典存放颜色分量上下限
# 例如：{颜色: [min分量, max分量]}
# {'red': [array([160,  43,  46]), array([179, 255, 255])]}

def getColorList():
    dict = collections.defaultdict(list)

    # 黑色
    lower_black = np.array([0, 0, 0])
    upper_black = np.array([180, 255, 46])
    color_list = []
    color_list.append(lower_black)
    color_list.append(upper_black)
    dict['black'] = color_list

    #灰色
    lower_gray = np.array([0, 0, 46])
    upper_gray = np.array([180, 43, 220])
    color_list = []
    color_list.append(lower_gray)
    color_list.append(upper_gray)
    dict['gray']=color_list

    # 白色
    lower_white = np.array([0, 0, 221])
    upper_white = np.array([180, 30, 255])
    color_list = []
    color_list.append(lower_white)
    color_list.append(upper_white)
    dict['white'] = color_list

    # 红色
    lower_red = np.array([156, 43, 46])
    upper_red = np.array([180, 255, 255])
    color_list = []
    color_list.append(lower_red)
    color_list.append(upper_red)
    dict['red'] = color_list

    # 红色2
    lower_red = np.array([0, 43, 46])
    upper_red = np.array([10, 255, 255])
    color_list = []
    color_list.append(lower_red)
    color_list.append(upper_red)
    dict['red2'] = color_list

    # 橙色
    lower_orange = np.array([11, 43, 46])
    upper_orange = np.array([25, 255, 255])
    color_list = []
    color_list.append(lower_orange)
    color_list.append(upper_orange)
    dict['orange'] = color_list

    # 黄色
    lower_yellow = np.array([26, 43, 46])
    upper_yellow = np.array([34, 255, 255])
    color_list = []
    color_list.append(lower_yellow)
    color_list.append(upper_yellow)
    dict['yellow'] = color_list

    # 绿色
    lower_green = np.array([35, 43, 46])
    upper_green = np.array([77, 255, 255])
    color_list = []
    color_list.append(lower_green)
    color_list.append(upper_green)
    dict['green'] = color_list

    # 青色
    lower_cyan = np.array([78, 43, 46])
    upper_cyan = np.array([99, 255, 255])
    color_list = []
    color_list.append(lower_cyan)
    color_list.append(upper_cyan)
    dict['cyan'] = color_list

    # 蓝色
    lower_blue = np.array([100, 43, 46])
    upper_blue = np.array([124, 255, 255])
    color_list = []
    color_list.append(lower_blue)
    color_list.append(upper_blue)
    dict['blue'] = color_list

    # 紫色
    lower_purple = np.array([125, 43, 46])
    upper_purple = np.array([155, 255, 255])
    color_list = []
    color_list.append(lower_purple)
    color_list.append(upper_purple)
    dict['purple'] = color_list
    return dict


if __name__ == '__main__':
    color_dict = getColorList()
    pic_dir = 'color_train_all/'
    for images in os.listdir(pic_dir):
        img = cv2.imread(pic_dir+images)
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

        mask_white = cv2.inRange(hsv, color_dict['white'][0], color_dict['white'][1])
        # 背景颜色为黑色，加入图片有像素，则保存
        if np.max(mask_white)>0:
            res = cv2.bitwise_and(img, img, mask=mask_white)
            cv2.imwrite('train/white/'+images,res)

        mask_green = cv2.inRange(hsv, color_dict['green'][0], color_dict['green'][1])
        if np.max(mask_green)>0:
            res = cv2.bitwise_and(img, img, mask=mask_green)
            cv2.imwrite('train/green/'+images,res)

        mask_purple = cv2.inRange(hsv, color_dict['purple'][0], color_dict['purple'][1])
        if np.max(mask_purple)>0:
            res = cv2.bitwise_and(img, img, mask=mask_purple)
            cv2.imwrite('train/purple/'+images,res)

        mask_blue = cv2.inRange(hsv, color_dict['blue'][0], color_dict['blue'][1])
        if np.max(mask_blue)>0:
            res = cv2.bitwise_and(img, img, mask=mask_blue)
            cv2.imwrite('train/blue/'+images,res)

        mask_cyan = cv2.inRange(hsv, color_dict['cyan'][0], color_dict['cyan'][1])
        if np.max(mask_cyan)>0:
            res = cv2.bitwise_and(img, img, mask=mask_cyan)
            cv2.imwrite('train/cyan/'+images,res)

        mask_red = cv2.inRange(hsv, color_dict['red2'][0], color_dict['red2'][1])
        if np.max(mask_red)>0:
            res = cv2.bitwise_and(img, img, mask=mask_red)
            cv2.imwrite('train/red/'+images,res)

        mask_yellow = cv2.inRange(hsv, color_dict['yellow'][0], color_dict['yellow'][1])
        if np.max(mask_yellow)>0:
            res = cv2.bitwise_and(img, img, mask=mask_yellow)
            cv2.imwrite('train/yellow/'+images,res)

提取后文件，颜色放在单独的颜色文件下,并统一放在train文件夹下

train

• blue
  –001.jpg
  –002.jpg
  –…
• cyan
• green
• purple
• red
• white
• yellow

同理制作test文件夹

pytorch建立AlexNet模型

model.py

import torch.nn as nn
import torch

class AlexNet(nn.Module):
    def __init__(self, num_classes=1000, init_weights=False):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 48, kernel_size=11, stride=4, padding=2),  
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),                 
            nn.Conv2d(48, 128, kernel_size=5, padding=2),           
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),                 
            nn.Conv2d(128, 192, kernel_size=3, padding=1),          
            nn.ReLU(inplace=True),
            nn.Conv2d(192, 192, kernel_size=3, padding=1),          
            nn.ReLU(inplace=True),
            nn.Conv2d(192, 128, kernel_size=3, padding=1),          
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),                  
        self.classifier = nn.Sequential(
            nn.Dropout(p=0.5),
            nn.Linear(128 * 6 * 6, 2048),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(2048, 2048),
            nn.ReLU(inplace=True),
            nn.Linear(2048, num_classes),
        )
        if init_weights:
            self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, start_dim=1)
        x = self.classifier(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)

模型训练

train.py

import json
import torch
import torch.nn as nn
from torchvision import transforms, datasets, utils
import matplotlib.pyplot as plt
import numpy as np
import torch.optim as optim
from tqdm import tqdm
from model import AlexNet

def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))
    batch_size = 32

    # torchvision自带的图片预处理
    # transforms.Resize()将图片调整为指定大小
    # transforms.RandomHorizontalFlip()图片随机水平翻转
    # ......参考torchvision官网教学
    data_transform = {
        "train": transforms.Compose([transforms.Resize((224, 224)),
                                     transforms.RandomHorizontalFlip(),
                                     transforms.ToTensor(),
                                     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),
        "val": transforms.Compose([transforms.Resize((224, 224)),  # cannot 224, must (224, 224)
                                   transforms.ToTensor(),
                                   transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                                   ])}
    # torchvison中的datasets加载训练集
    train_dataset = datasets.ImageFolder(root="../train/",transform=data_transform["train"])
    train_num = len(train_dataset)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

    # {'blue': 0, 'cyan': 1, 'green': 2, 'purple': 3, 'red': 4, 'white': 5, 'yellow': 6}
    flower_list = train_dataset.class_to_idx
    # 将字典进行编码，最终生成class_indices.json文件
    cla_dict = dict((val, key) for key, val in flower_list.items())
    json_str = json.dumps(cla_dict, indent=4)
    with open('class_indices.json', 'w') as json_file:
        json_file.write(json_str)

    # torchvison中的datasets加载验证集
    validate_dataset = datasets.ImageFolder(root='../val',transform=data_transform["val"])
    val_num = len(validate_dataset)
    validate_loader = torch.utils.data.DataLoader(validate_dataset,batch_size=batch_size, shuffle=False)
    print("using {} images for training, {} images for validation.".format(train_num,val_num))

    # # 训练集图像可视化，可注释
    # test_data_iter = iter(validate_loader)
    # test_image, test_label = test_data_iter.__next__()
    # def imshow(img):
    #     img = img / 2 + 0.5  # unnormalize
    #     print(img.shape)
    #     npimg = img.numpy()
    #     plt.imshow(np.transpose(npimg, (1, 2, 0)))
    #     plt.show()
    # print(' '.join('%5s' % cla_dict[test_label[j].item()] for j in range(4)))
    # # make_grid的作用是将若干幅图像拼成一幅图像,在需要展示一批数据时很有用
    # imshow(utils.make_grid(test_image))

    # num_classes=分类个数 init_weights=初始化权重
    net = AlexNet(num_classes=7, init_weights=True)
    net.to(device)
    loss_function = nn.CrossEntropyLoss()  # 多分类常用的损失函数
    optimizer = optim.Adam(net.parameters(), lr=0.0002)  #优化器
    epochs = 5
    best_acc = 0.0   # 更新准确率最高的数值
    best_loss = 1.0  # 更新损失最低的数值
    train_steps = len(train_loader)
    for epoch in range(epochs):
        # 通过net.train()可以保证dropout/BatchNormal只在训练时候起作用
        net.train()
        running_loss = 0.0  # 统计训练过程中的损失
        train_bar = tqdm(train_loader)  # 119
        for step, data in enumerate(train_bar):
            # images: (batchsize,3,224,224)
            # labels: batchsize
            images, labels = data
            optimizer.zero_grad()
            # outputs: (batchsize,classes)
            outputs = net(images.to(device))
            loss = loss_function(outputs, labels.to(device))  # 计算预测值与真实值
            loss.backward()  # 损失反向传播
            optimizer.step()  # 更新参数
            # 计算一共多少损失
            running_loss += loss.item()
            train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch+1,epochs,loss)
            
        # 验证
        net.eval()
        acc = 0.0  # 计算精度公式 number / epoch
        with torch.no_grad():  # 进制pytorch对参数跟踪
            val_bar = tqdm(validate_loader)
            for val_data in val_bar:
                val_images, val_labels = val_data
                outputs = net(val_images.to(device))
                predict_y = torch.max(outputs, dim=1)[1]
                acc += torch.eq(predict_y, val_labels.to(device)).sum().item()

        val_accurate = acc / val_num
        print('[epoch %d] train_loss: %.3f  val_accuracy: %.3f' %(epoch + 1, running_loss/train_steps, val_accurate))
        val_loss = running_loss/train_steps
        if val_loss<best_loss:
            best_loss = val_loss
            if val_accurate > best_acc:
                best_acc = val_accurate
                torch.save(net.state_dict(), './epoch%d_train_loss_%.2fval_accuracy_%.2f.pth'
                           %(epoch + 1, running_loss/train_steps, val_accurate))
    print('Finished Training')

if __name__ == '__main__':
    main()

模型预测

import os
import cv2
import json
import torch
import numpy as np
from PIL import Image
from torchvision import transforms
from model import AlexNet

img_path = 'purple.jpg'
def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    # 和训练时图像处理一样：不过只需要resize、ToTensor、Normalize
    data_transform = transforms.Compose(
        [transforms.Resize((224, 224)),
         transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    # 读取json文件并解码
    json_path = 'class_indices.json'
    assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)
    json_file = open(json_path, "r")
    class_indict = json.load(json_file)

    # 创建模型
    model = AlexNet(num_classes=7).to(device)
    # 加载训练好的权重
    weights_path = "epoch4train_loss_0.00val_accuracy_1.00.pth"
    assert os.path.exists(weights_path), "file: '{}' dose not exist.".format(weights_path)
    model.load_state_dict(torch.load(weights_path))
    model.eval()  # 验证模式

    ''''''''''
    OpenCV转换成PIL.Image格式
    image = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    Image转换成OpenCV格式
    img = cv2.cvtColor(numpy.asarray(image), cv2.COLOR_RGB2BGR)
    '''''''''''
    image = cv2.imread(img_path)
    img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    img = data_transform(img)
    img = torch.unsqueeze(img, dim=0)  # 维度扩张，增加一个batchsize维度
    with torch.no_grad():
        output = torch.squeeze(model(img.to(device))).cpu()  # 降维
        predict = torch.softmax(output, dim=0)
        predict_cla = torch.argmax(predict).numpy()
    print_res = "class:{} confidence:{:.2}".format(class_indict[str(predict_cla)],predict[predict_cla].numpy())
    cv2.putText(image, print_res, (0, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2)
    cv2.imshow('image',image)
    cv2.waitKey(0)
if __name__ == '__main__':
    main()