【Day17】自製訓練集：總結

15th鐵人賽

sunteng

2023-10-02 23:59:24

310 瀏覽

分享至

今天我們把所有前面所用過的爬蟲、讀取權重、分類全融合在一起，把我們所學全部把它並一起運用，我們可以以這方式來幫我們快速爬蟲獲取資料，並讓模型幫我們分類，我們可以讓模型做完第一次分類後再到分好資料夾，把其他不相關圖片給移除掉，之後來可以把新增圖片再讓模型訓練一次，讓模型更精確。

程式部分

我們把前面部分程式做融合，資料夾我們也盡量在前面先儲存再變數裡，方便後面調用

引入函式庫

import torch
import torch.nn as nn
from PIL import Image
import torchvision.transforms as transforms
import numpy as np
import torch.nn.functional as F
import os
import shutil
from bs4 import BeautifulSoup
from selenium import webdriver
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.chrome.service import Service
import time
import requests

資料夾位置

輸出資料夾位置要放完整的路徑，我的分類資料夾要是以下格式
C:\Users\image\train\0猴子或1人
則 output_folder = r'C:\Users\image\train'

# 資料夾路徑
input_folder = 'images'  # 輸入資料夾，包含要分類的圖片
output_folder = r'你的完整路徑'   # 輸出資料夾，用於存放分類後的圖片

# 存圖位置
local_path = 'images'

爬蟲部分

# 目標元素的class
img_class = 'rg_i'
url = 'https://www.google.com/search?sca_esv=563286194&hl=zh-TW&sxsrf=AB5stBhULplilmDPtWzjSST6XL-DTIN1Mg:1694059984176&q=%E4%BA%BA%E5%83%8F&tbm=isch&source=lnms&sa=X&ved=2ahUKEwiA29ac0ZeBAxUPH3AKHWOyCHAQ0pQJegQICxAB&biw=1920&bih=923&dpr=1'

# 設定 user agent 防止網站鎖 IP
chrome_options = webdriver.ChromeOptions()
chrome_options.add_argument("user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36")

# 指定 chromedriver.exe 的路径
chrome_path = 'C:/Users/user/Desktop/AI入門學習紀錄/新增資料夾/chromedriver.exe'
service = Service(chrome_path)

#  Chrome WebDriver 對象
chrome = webdriver.Chrome(service=service, options=chrome_options)

# 前往網站
chrome.get(url)

# 建立目錄
if not os.path.exists(local_path):
    os.makedirs(local_path)

img_url_dic = {}

# 模擬滾動視窗瀏覽更多圖片
scroll_pause_time = 1
last_height = chrome.execute_script("return document.body.scrollHeight")

while True:

    chrome.execute_script("window.scrollTo(0, document.body.scrollHeight);")
    time.sleep(scroll_pause_time)
    time.sleep(0.05)
    new_height = chrome.execute_script("return document.body.scrollHeight")


    try:
        # 找到並點擊 "顯示更多" 按鈕
        show_more_button = chrome.find_element("xpath", '//*[@id="islmp"]/div/div/div/div/div[1]/div[2]/div[2]/input')
        show_more_button.click()
    except Exception as e:
        print("無法點擊顯示更多按鈕:", e)
    if new_height == last_height:
        break

    last_height = new_height

soup = BeautifulSoup(chrome.page_source, 'html.parser')

m = 10000 # 起始圖片編號
for img in soup.find_all('img', class_=img_class):
    try:
        img_url = img.get('src')

        # 保存圖片到指定路徑
        if img_url and not img_url in img_url_dic:
            img_url_dic[img_url] = ''
            m += 1
            ext = os.path.splitext(img_url)[1]
            filename =   str(m) + ext +'.png'
            print(filename)

            # 下載圖片
            img_response = requests.get(img_url, stream=True)
            with open(os.path.join(local_path, filename), 'wb') as out_file:
                shutil.copyfileobj(img_response.raw, out_file)
            del img_response

    except Exception as e:
        print('發生錯誤:', e)

chrome.quit()

print('下載完成')

分類部分

# 檢查並創建輸出資料夾
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

class CNN(nn.Module):
    def __init__(self):
        #呼叫nn.Module裡面init的資料
        super().__init__()
        #定義神經網路
        self.conv1 = nn.Conv2d(3, 16, 3)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(16, 32, 3)
        self.fc1 = nn.Linear(32 * 30 * 30, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 2)

    def forward(self, x):
        #定義操作
        x = self.pool(F.relu(self.conv1(x)))  # 第一個卷積層和池化層
        x = self.pool(F.relu(self.conv2(x)))  # 第二個卷積層和池化層
        x = x.view(x.size(0), -1)  # 展平特徵圖  x.size(0) 來獲取批次大小，然 -1 來自動計算展平後的尺寸
        x = F.relu(self.fc1(x))  # 第一個全連接層
        x = F.relu(self.fc2(x))  # 第二個全連接層
        x = self.fc3(x)  # 最後一個全連接層（輸出層）
        return x    
            
# 載入訓練好的模型權重
model = CNN()
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

# 轉換設定
transform = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
label_encoding = {i:cnt for cnt, i in enumerate(os.listdir('images/train'))}
label_decoding = {v:k for k, v in label_encoding.items()}
# 讀取資料夾中的所有圖片並分類
for filename in os.listdir(input_folder):
    img_path = os.path.join(input_folder, filename)
    x_valid = transform(Image.open(img_path).convert("RGB"))
    img = x_valid.unsqueeze(0).to("cpu")
    output = model(img)
    _, pred = torch.max(output, dim=1)
    label = pred.tolist()[0]

    # 根據預測結果分類到不同的資料夾
    if int(label_decoding[label]) == 0:
        print(f'{filename} 猴子')
        shutil.move(img_path, os.path.join(output_folder, '0', filename))
    elif int(label_decoding[label]) == 1:
        print(f'{filename} 人')
        shutil.move(img_path, os.path.join(output_folder, '1', filename))

print("分類完成")

以這模式我們可以一開始先爬取少部分資料來訓練模型，讓模型有初步分辨能力，再來用後面總結的程式來讓爬蟲+分類一次完成，結束再反覆訓練模型讓我們模型更準確，可以減少很多人工挑選圖片時間，而分類項目可以再多一項其他類別，這樣可以把完全不相干的圖片給另外移除，今天比較沒什麼說明，主要把前面學過作總結，明天開始我們來練習讓模型辨別文字與相關運用，我們明天見!