
目錄
- 概述
- 算法原理
- 核心邏輯
- 效果演示
- 使用方式
- 參考文獻
參考文獻:需要本文的詳細復現過程的項目源碼、數據和預訓練好的模型可從該地址處獲取完整版:
地址
概述
本文基于論文 Multi-Label Classification using Deep Convolutional Neural Network[1] 實現圖像自動標注程序。
計算機技術的進步和互聯網產業的不斷發展導致了網絡圖像數量的爆炸式增長,如何管理種類繁多的海量圖像成為了一個重要問題。自動圖像標注(Automatic Image Tagging)作為一項重要的圖像管理技術,可以利用計算機自動為每張圖像打上與其內容有關的標簽,從而幫助用戶更好地搜索和訪問圖像。
圖1:圖像自動標注任務
近年來,隨著深度學習技術的發展,深度神經網絡能夠捕捉到更多且更加復雜的圖像特征,這使得圖像標注算法的性能也隨之受益。圖像標注與圖像多標簽分類有著天然的緊密連系,后者會根據內容將一張圖像歸納到多個類別中。綜上,本文基于目前先進的深度神經網絡 VGG-Net[2] 和大規模圖像多標簽分類數據集 MS-COCO-2017[3] 訓練自動圖像標注模型。
算法原理
VGG-Net 是一種經典的卷積神經網絡 (Convolutional Neural Network) 架構,其核心思想是通過更深的網絡結構以及使用較小的卷積核來提取更豐富的圖像特征。VGG-Net 通過堆疊多個卷積層來加深網絡,且卷積層全部采用大小為 3×3 的小卷積核,步長為 1,填充為 1。這種設計通過堆疊多個小卷積核來增加網絡的非線性表達能力,且相比使用較大的卷積核,能減少參數數量。在若干卷積層后,VGG-Net 使用 2×2 的最大池化層,步長為 2。池化層用于減少特征圖的尺寸,并保留主要的特征。在最后的卷積層之后,VGG-Net 通過三個全連接層對特征進行進一步處理,最后輸出分類結果。在每個卷積層和全連接層之后,VGG-Net 使用 ReLU (Rectified Linear Unit) 激活函數,以增加網絡的非線性。
本文使用一個線性層和 Sigmoid 函數構建模型的分類器,并利用二元交叉熵損失(Binary Cross-Entropy, BCE)進行訓練。
Sigmoid(x)=11+e?x
核心邏輯
程序的核心代碼如下所示:
# transformtransform = v2.Compose([v2.Resize(256),v2.CenterCrop(224),v2.RandomHorizontalFlip(),v2.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),v2.ToImage(),v2.ToDtype(torch.float32, scale=True),v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),])# devicedevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')# datasettrain_dataset = COCO_Dataset(configs['train_annotations_path'], configs['train_images_dir'], transform)train_dataloader = DataLoader(train_dataset, batch_size = configs['batch_size'], shuffle = True, num_workers=8, pin_memory=True)test_dataset = COCO_Dataset(configs['test_annotations_path'], configs['test_images_dir'], transform)test_dataloader = DataLoader(test_dataset, batch_size = configs['batch_size'], shuffle = False, num_workers=8, pin_memory=True)# modelmodel = ImageTaggingModel().to(device)optimizer = optim.Adam(model.parameters(), lr=configs['learning_rate'], weight_decay=configs['weight_decay'])lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=configs['lr_decay_step'], gamma=configs['lr_decay_rate'])# logloss_epoch = []precise_epoch = []recall_epoch = []f1_epoch = []# train & testfor epoch_id in range(configs['epochs']):current_loss = 0# trainmodel.train()for batch in tqdm(train_dataloader, desc='Training(Epoch %d)' % epoch_id, ascii=' 123456789#'):optimizer.zero_grad()images = batch['images'].to(device)labels = batch['labels'].to(device)logits = model(images)loss = F.binary_cross_entropy_with_logits(logits, labels)current_loss += loss.item()loss.backward()optimizer.step()lr_scheduler.step()current_loss /= len(train_dataloader)print('Current Average Loss:', current_loss)loss_epoch.append(current_loss)plt.plot(loss_epoch)plt.xlabel('Epoch')plt.ylabel('Loss')plt.title('Loss-Epoch')plt.savefig(os.path.join(configs['logs_dir'], "Loss.png"), dpi=300)plt.clf()# testmodel.eval()TT_num = 0FT_num = 0FF_num = 0with torch.no_grad():for batch in tqdm(test_dataloader, desc='Testing(Epoch %d)' % epoch_id, ascii=' 123456789#'):images = batch['images'].to(device)labels = batch['labels'].to(device)logits = model(images)probs = F.sigmoid(logits)predictions = (probs > configs['threshold']).to(labels.dtype)TT_num += torch.sum(predictions * labels).item()FT_num += torch.sum(predictions * (1 - labels)).item()FF_num += torch.sum((1 - predictions) * labels).item()precise = TT_num / (TT_num + FT_num)recall = TT_num / (TT_num + FF_num)f1_score = 2 * precise * recall / (precise + recall)precise_epoch.append(precise)recall_epoch.append(recall)f1_epoch.append(f1_score)print("Precise = %.2f, Recall = %.2f, F1-score = %.2f" % (precise, recall, f1_score))plt.plot(precise_epoch, label='Precise')plt.plot(recall_epoch, label='Recall')plt.plot(f1_epoch, label='F1-score')plt.xlabel('Epoch')plt.ylabel('Value')plt.title('Result')plt.legend()plt.savefig(os.path.join(configs['logs_dir'], "Result.png"), dpi=300)plt.clf()# save modeltorch.save(model.state_dict(), configs['checkpoint'])
以上代碼僅作展示,更詳細的代碼文件請參見附件。
效果演示
配置環境并運行 main.py腳本,效果如圖4所示。
此外,網站還提供了在線體驗功能。用戶只需要輸入一張大小不超過 1MB 的 JPG 圖像,網站就會自動為圖像打上標記并展示詞云,如圖5所示。
使用方式
解壓附件壓縮包并進入工作目錄。如果是Linux系統,請使用如下命令:
unzip ImageCaptioning.zip
cd ImageCaptioning
代碼的運行環境可通過如下命令進行配置:
pip install -r requirements.txt
如果在本地測試自動圖像標注程序,請運行如下命令:
python main.py
如果希望在線部署,請運行如下命令:
python main-flask.py
參考文獻
-
[1] Lydia A A, Francis F S. Multi-label classification using deep convolutional neural network[C]//2020 international conference on innovative trends in information technology (ICITIIT). IEEE, 2020: 1-6.
-
[2] Simonyan K, Zisserman A. Very deep convolutional networks for large-scale image recognition[J]. arXiv preprint arXiv:1409.1556, 2014.
-
[3] Lin T Y, Maire M, Belongie S, et al. Microsoft coco: Common objects in context[C]//Computer Vision–ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13. Springer International Publishing, 2014: 740-755.
參考文獻:需要本文的詳細復現過程的項目源碼、數據和預訓練好的模型可從該地址處獲取完整版:
地址