ReclassifyAnimeCG - 來用深度學習整理CG圖庫吧

會想寫這個工具的主要原因是, 我有追蹤一些繪師並且用下載器去追蹤下載那些繪師的作品, 不過基本上繪師畫的圖通常就是幾天放個一張, 或是一次推出的CG圖庫都是複合的動漫畫主題, 然後下載器下載又不可能自動幫我分類這些圖片, 導致這些圖片就越來越雜亂, 每個資料夾都混雜複數個主題&角色, 每個資料夾的名稱又沒有特定的主題性, 讓我自己事後很難找圖不說, 連看都懶得看了...。 後來就想到現在深度學習分類圖片已經很強了, 那我乾脆來寫個工具輔助我整理圖庫吧!!

做好的成果如下:

https://github.com/zmcx16/ReclassifyAnimeCG

使用的方法大致如下:

1. 設置config.yaml

設定想整理的圖庫位置(predict_data_path), 以及先手動整理好一小部分的資料當訓練資料(train_data_path), 還有預定之後模型分類完存放圖庫的位置(output_data_path), 還有設定想訓練/使用的深度學習模型。

2. 執行preprocess.py

preprocess.py會生成 train.txt, test.txt, index.txt 到 label_path 位置。train.txt儲存訓練資料的每張圖片的位置, 以及對應的資料夾名稱(類別); test.txt則是根據config.yaml裡的train_ratio來決定要從training_data裡拿多少資料做為testing data, 如果今天是想測試模型好壞或調參數會設定到, 如果是單純想使用來整理圖片不調參數, 就直接設定成1.0; index.txt則是記錄每個類別index對應的資料夾名稱。

3. 執行train.py

train.py會開始訓練模型, 並把訓練好的模型存在train_model_path位置

4. 執行predict.py

如果config.yaml設定use_test_txt = true, predict.py會使用之前preprocess.py生成的testing data來做分類預測, 並計算這次分類的準確率, 可以根據分類好壞來調整模型跟參數改善模型; 如果use_test_txt = false, predict.py會遍歷predict_data_path位置的所有圖檔, 並針對這些圖檔做分類預測, 以及根據預測結果儲存到output_data_path。

5. 如果是調整或訓練模型到步驟4就結束了, 而如果是要整體圖庫的話, 手動確認分類結果, 把分類錯誤的部分砍掉, 之後把剩餘正確(部分)的分類結果複製到training_data內。

6. 跳回步驟2, 反覆做到個人滿意分類結果為止。

介紹跟使用上差不多是這樣, 再來是一些開發過程心得, 內容繁瑣有想看在看~

會做這個工具主要是輔助整理圖庫, 所以最主要的兩個重點:

1. 模型準確率:

模型的準確率會大大影響上述步驟4耗費的人力, 只要準確率越高, 需要人工處理介入的部分就越少, 這塊目前只實作好直接套用PyTorch內建模型, 如果真的要提高準確率, 目前想法是結合動漫角色臉部識別做前處理, 之後在丟給CNN視覺模型, 這樣應該比起單純丟原始圖片更能大幅提高準確率, 這塊就等有空在來研究怎麼做了。

2. 讓程式提供更多功能減少人工工作:

這程式最終想提供的還是幫忙整理圖庫, 所以能讓程式做的就希望盡量交給程式處理, 另外還有要怎麼整理才能滿足瀏覽圖片的需求, 個人是有幾個想法:

a. 不去動到原始檔案的資料, 而是建立一個資料庫, 對每一個圖檔都建立tag標籤, 之後可以用tag的方式搜尋資料庫, 並瀏覽符合搜尋結果的圖片。

這算是最終想做到的方案, 現在網路上也有不少圖庫網站是提供這種方式讓使用者找圖跟瀏覽圖片, 不過這等於要把整套系統都做出來, 包含資料庫, 檔案瀏覽器, 看圖軟體三大功能都實作才行, 要是之後檔案有動到也有資料庫遷移的問題, 個人目前還沒有那麼偉大的夢想想去自己做這一套系統出來, 所以就先不考慮了。

b. 對想分類的圖片一樣用資料夾管理, 不過一樣不去動到原始檔案的資料而且也不複製圖片, 圖片檔案一律用symbolic link去連結, 事後如果原始檔案有遷移需求的話, 簡單寫個程式去替換symbolic link就好, 這樣的好處是硬碟不用因為整理圖庫多花空間, 想怎麼分類就怎麼分類, 而且也可以直接用市面現成的看圖軟體瀏覽圖片。

c. 對想分類的圖片一樣用資料夾管理, 原始檔案就直接搬移或複製。

目前的ReclassifyAnimeCG提供了上述 (b.) & (c.) 的功能, 以我來說我主要就是用(b.)的功能, 除此之外, 為了提升程式效能&減少重複的運算, 也實作了重複檔案的檢查(crc32), 藉以刪除重複檔案以及不預測已經在training data的圖片等等, 這些功能都可以透過config.yaml控制, 不過也因為這樣config.yaml也越來越複雜了就是...。

最後來推薦一下看圖軟體FastStone Image Viewer, 上述的(b.)會把分類好的圖片建立symbolic link, 可是不給力的Windows系統竟然File Explorer不能預覽symbolic link連結的圖片, 而且如果用相片檢視器瀏覽symbolic link連結的圖片, 他會直接開啟到原始檔案的路徑, 所以瀏覽下一張圖也會在原始檔案的資料夾, 而不是symbolic link圖檔在的資料夾...。 這問題用上面的看圖軟體就可以解決了, 所以基本上問題都解決了, 以後整理好的圖庫我應該都會用FastStone Image Viewer觀看~。 

最後附上我目前在整理的某繪師圖庫, 圖檔數量超過2W張, 目前我用上述的流程跑到Round2, 大概已經分類好2000多張圖片, 我第一Round準備了100多張, 20幾個自己分類的訓練資料, 程式跑完就從分類器幫分類好的那20幾個分類資料把部分正確分類的複製到training data, 並且在手動擴充更多自己想分類的資料樣本到training data, 所以第一Round跑完我的訓練資料就擴充到2000多張圖片, 57個分類; 之後就繼續跑Round2, 反覆幾次後應該就能把我想分類的圖片都分類完, 而且後期隨著training_data的樣本數變多, 分類錯誤的情況就會大幅改善, 我也不用勞心勞力把某類分類完才能繼續做下一步。

以我來說, 這個工具就是幫助我大幅減少人工整理圖庫的時間, 還有藉由可以多次rerun改善模型的方式, 不用一定要把某類一次分類完, 大幅降低精神污染的危險, 個人還算蠻滿意這次的成果, 再來要做的就是盡量提高模型的準確率, 減少需要人工過濾分類好的圖庫的時間~。