标签：nn color list 分类 cv2 pytorch train dict AlexNet

数据集

百度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()

标签：nn,color,list,分类,cv2,pytorch,train,dict,AlexNet
来源： https://blog.csdn.net/weixin_45930948/article/details/120104737

本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享；
2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关；
3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关；
4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除；
5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。