1. 程式人生 > 其它 >深度學習實戰之垃圾分類

深度學習實戰之垃圾分類

垃圾分類,指按一定規定或標準將垃圾分類儲存、分類投放和分類搬運,從而轉變成公共資源的一系列活動的總稱。分類的目的是提高垃圾的資源價值和經濟價值,力爭物盡其用;然而我們在日常生活中認為對垃圾分類還是有些不知所措的,對幹垃圾、溼垃圾……分的不是很清楚,由此我們就想到了使用深度學習的方法進行分類。

簡介

本篇博文主要會帶領大家進行資料的預處理、網路搭建、模型訓練、模型測試

1. 獲取資料集
這裡筆者已經為大家提供了一個比較完整的資料集,所以大家不必再自己去收集資料了
資料集連結:https://pan.baidu.com/s/1OhA7yQt9evqNIP5CIPjdgw​
提取碼:Z5A1
如下為資料集中的部分資料展示

這裡就不過多的展示了,因為有些圖片比較的不雅
ok,不說題外話了,我們繼續

2.預處理資料集

import torch,visdom,time
import os,csv,random,glob
from PIL import Image
from torchvision import transforms
from torch.utils.data import DataLoader,Dataset

class Data(Dataset):
    def __init__(self,root,resize,model):
        super(Data, self).__init__()
        self.root=root
        self.resize=resize
        
        # dividing data
        if model=='train':
            self.images=self.images[:int(0.6*len(self.images))]
            self.labels=self.labels[:int(0.6*len(self.labels))]
        if model=='verify':
            self.images=self.images[int(0.6*len(self.images)):int(0.8*len(self.images))]
            self.labels=self.labels[int(0.6*len(self.labels)):int(0.8*len(self.labels))]
        else:
            self.images=self.images[int(0.8*len(self.images)):]
            self.labels=self.labels[int(0.8*len(self.labels)):]
def main():
    # show data
    viz = visdom.Visdom()
    db = Data('data', 64, 'train')

    x,y = next(iter(db))
    print('sample:', x.shape, y.shape, y)

    viz.image(db.normalize(x), win='sample_x', opts=dict(title='sample_x'))

if __name__ == '__main__':
    main()

這裡需要對資料進行增強,增強的最終結果如下圖所示:

  1. 搭建網路
    接下來進行搭建網路了,筆者這裡使用的是resnet的網路結構
import  torch
from    torch import  nn
from    torch.nn import functional as F
class ResBlk(nn.Module):
  
    def __init__(self, ch_in, ch_out, stride=1):
     
        super(ResBlk, self).__init__()

        self.conv1 = nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=stride, padding=1)
        self.bn1 = nn.BatchNorm2d(ch_out)
        self.conv2 = nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(ch_out)

        self.extra = nn.Sequential()
        if ch_out != ch_in:
            # [b, ch_in, h, w] => [b, ch_out, h, w]
            self.extra = nn.Sequential(
                nn.Conv2d(ch_in, ch_out, kernel_size=1, stride=stride),
                nn.BatchNorm2d(ch_out)
            )
class ResNet18(nn.Module):

    def __init__(self, num_class):
        super(ResNet18, self).__init__()

        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, stride=3, padding=0),
            nn.BatchNorm2d(16)
        )
        # followed 4 blocks
        # [b, 16, h, w] => [b, 32, h ,w]
        self.blk1 = ResBlk(16, 32, stride=3)
        # [b, 32, h, w] => [b, 64, h, w]
        self.blk2 = ResBlk(32, 64, stride=3)
        # # [b, 64, h, w] => [b, 128, h, w]
        self.blk3 = ResBlk(64, 128, stride=2)
        # # [b, 128, h, w] => [b, 256, h, w]
        self.blk4 = ResBlk(128, 256, stride=2)

        self.outlayer = nn.Linear(256*3*3, num_class)

    def forward(self, x):
      
        return x

def main():
   
if __name__ == '__main__':
    main()

  1. 訓練
import torch
from torch import optim, nn
import visdom
import torchvision
from torch.utils.data import DataLoader

from Data_Pre import Data
# from    resnet import ResNet18
from torchvision.models import resnet18

from utils import Flatten

batchsz = 32
lr = 1e-3
epochs = 20

device = torch.device('cuda')
torch.manual_seed(1234)

train_db = Data('data', 224, model='train')
val_db = Data('data', 224, model='verify')
test_db = Data('data', 224, model='test')
train_loader = DataLoader(train_db, batch_size=batchsz, shuffle=True,
                          num_workers=4)
val_loader = DataLoader(val_db, batch_size=batchsz, num_workers=2)
test_loader = DataLoader(test_db, batch_size=batchsz, num_workers=2)

viz = visdom.Visdom()

def evalute(model, loader):
    model.eval()

    correct = 0
    total = len(loader.dataset)

    for x, y in loader:
        x, y = x.to(device), y.to(device)
        with torch.no_grad():
            logits = model(x)
            pred = logits.argmax(dim=1)
        correct += torch.eq(pred, y).sum().float().item()
    return correct / total

視覺化訓練結果



可以看到經過次的迭代後,模型的效果還是不錯的,哈哈哈哈哈。。。。。。。
5. 測試模型

專案完整程式碼:https://github.com/taotao/garbage
有什麼問題歡迎同筆者討論:1017190168
最後希望大家能夠動手實踐實踐

最新版本:

自筆者於2019年8月發出了深度學習實戰之垃圾分類博文以來,已被訪問上萬次,幾百次的收藏。同時也有很多的人詢問筆者關於這個專案的問題。在這裡非常的感謝大家的厚愛。其實很早就想再迭代一個垃圾分類專案出來,但是由於之前比較的忙,一直沒有時間將其落實。最近沒什麼事情,筆者花了一些時間又重新做了一個垃圾分類系統。希望對大家學習研究有所幫助。

新版本連結:

深度學習之垃圾分類系統(識別率近99%)
垃圾分類識別系統(目標檢測版)

近期專案:

新冠肺炎(Covid-19)病毒檢測
歡迎大家閱讀、獲取。

申明近期筆者接計算機視覺方面的畢業設計、比賽,有需要的歡迎聯絡!!!