關(guān)鍵點是一種應(yīng)用廣泛的技術(shù)，在人臉識別、人體姿態(tài)檢測、車牌識別等方面均有較好表現(xiàn)。本文采用回歸和heatmap兩種方式回歸關(guān)鍵點，嘗試定位紅外小目標(biāo)，形成了一個適合新手學(xué)習(xí)的基礎(chǔ)項目。
筆者出于興趣搞了一個小的庫，主要是用于定位紅外小目標(biāo)。由于其具有尺度很小的特點，所以可以嘗試用點的方式代表其位置。本文主要采用了 回歸 和 heatmap 兩種方式來回歸關(guān)鍵點，是一個很簡單基礎(chǔ)的項目，代碼量很小，可供新手學(xué)習(xí)。 1 數(shù)據(jù)來源 數(shù)據(jù)集 ：數(shù)據(jù)來源自小武，經(jīng)過小武的授權(quán)使用，但不會公開。本項目只用了其中很少一部分，共108張圖片。https://github.com/pprp/landmark_annotation 部分樣例展示 上圖是數(shù)據(jù)集中的兩張圖片，紅圈代表對應(yīng)的目標(biāo)，標(biāo)注的時候只需要在其中心點一下即可得到該點對應(yīng)的橫縱坐標(biāo)。 該數(shù)據(jù)集有一個特點，每張圖只有一個目標(biāo)（不然沒法用簡單的方法回歸），多余一個目標(biāo)的圖片被剔除了。
以上是一個標(biāo)注文件的例子，1.jpg對應(yīng)1.txt 2 回歸確定關(guān)鍵點 回歸確定關(guān)鍵點比較簡單，網(wǎng)絡(luò)部分采用手工構(gòu)建的一個兩層的小網(wǎng)絡(luò)，訓(xùn)練采用的是MSELoss。 https://github.com/pprp/SimpleCVReproduction/tree/master/simple_keypoint/regression 2.1 數(shù)據(jù)加載 數(shù)據(jù)的組織比較簡單，按照以下格式組織： - data 重寫一下Dataset類，用于加載數(shù)據(jù)集。 class KeyPointDatasets (Dataset) :def __init__ (self, root_dir='./data' , transforms=None) :'images' )# self.txt_path = os.path.join(root_dir, 'labels') '*.jpg' ))'.jpg' , '.txt' ).replace('images' , 'labels' ) for item in self.img_list]if transforms is not None :def __getitem__ (self, index) :if self.transforms:with open(txt, 'r' ) as f:for i, line in enumerate(f):if i == 0 :# 第一行 else :for t in line.split()]# range from 0 to 1 return img, torch.tensor(label[0 ])def __len__ (self) :return len(self.img_list) @staticmethod def collect_fn (batch) :return torch.stack(imgs, 0 ), torch.stack(labels, 0 )
返回的結(jié)果是圖片和對應(yīng)坐標(biāo)位置。 2.2 網(wǎng)絡(luò)模型 import torchimport torch.nn as nnclass KeyPointModel (nn.Module) :def __init__ (self) :3 , 6 , 3 , 1 , 1 )6 )True )2 , 2 ))6 , 12 , 3 , 1 , 1 )12 )True )2 , 2 ))1 )12 , 2 ),def forward (self, x) :0 ], -1 )return self.classifier(x)其結(jié)構(gòu)就是卷積+pooling+卷積+pooling+global average pooling+Linear，返回長度為2的tensor。 def train (model, epoch, dataloader, optimizer, criterion) :for itr, (image, label) in enumerate(dataloader):0 ]if itr % 4 == 0 :'epoch:%2d|step:%04d|loss:%.6f' % (epoch, itr, loss.item()/bs))'train_loss' : loss.item()*100 /bs})300 10 ######################################## 360 ,480 )),0.4372 , 0.4372 , 0.4373 ],0.2479 , 0.2475 , 0.2485 ])'./data' , transforms=transforms_all)True ,3e-4 )# criterion = torch.nn.SmoothL1Loss() 30 ,0.1 )for epoch in range(total_epoch):if epoch % 10 == 0 :'weights/epoch_%d_%.3f.pt' % (epoch, loss*1000 ))
loss部分使用Smooth L1 loss或者M(jìn)SE loss均可。 2.4 測試結(jié)果 3 heatmap確定關(guān)鍵點 這部分代碼很多參考了CenterNet，不過曾經(jīng)嘗試CenterNet中的loss在這個問題上收斂效果不好，所以參考了kaggle人臉關(guān)鍵點定位的解決方法，發(fā)現(xiàn)使用簡單的MSELoss效果就很好。 這部分和CenterNet構(gòu)建heatmap的過程類似，不過半徑的確定是人工的。因為數(shù)據(jù)集中的目標(biāo)都比較小，半徑的范圍最大不超過半徑為30個像素的圓。 class KeyPointDatasets (Dataset) :def __init__ (self, root_dir='./data' , transforms=None) :1 480 // self.down_ratio360 // self.down_ratio'images' )'*.jpg' ))'.jpg' , '.txt' ).replace('images' , 'labels' ) for item in self.img_list]if transforms is not None :def __getitem__ (self, index) :if self.transforms:with open(txt, 'r' ) as f:for i, line in enumerate(f):if i == 0 :# 第一行 else :for t in line.split()]# range from 0 to 1 30 )return img, torch.tensor(heatmap).unsqueeze(0 )def __len__ (self) :return len(self.img_list) @staticmethod def collect_fn (batch) :return torch.stack(imgs, 0 ), torch.stack(labels, 0 )核心函數(shù)是draw_umich_gaussian，具體如下： def gaussian2D (shape, sigma=1 ) :1. ) / 2. for ss in shape]1 , -n:n + 1 ]2 * sigma * sigma))0 # 限制最小的值 return hdef draw_umich_gaussian (heatmap, center, radius, k=1 ) :2 * radius + 1 6 )# 一個圓對應(yīng)內(nèi)切正方形的高斯分布 0 ]), int(center[1 ])1 )1 )if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0 : # TODO debug # 將高斯分布覆蓋到heatmap上，取最大，而不是疊加 return heatmap
sigma參數(shù)直接沿用了CenterNet中的設(shè)置，沒有調(diào)節(jié)這個超參數(shù)。 3.2 網(wǎng)絡(luò)結(jié)構(gòu) 網(wǎng)絡(luò)結(jié)構(gòu)參考了知乎上一個復(fù)現(xiàn)YOLOv3中提到的模塊，Sematic Embbed Block(SEB)用于上采樣部分，將來自低分辨率的特征圖進(jìn)行上采樣，然后使用3x3卷積和1x1卷積統(tǒng)一通道個數(shù)，最后將低分辨率特征圖和高分辨率特征圖相乘得到融合結(jié)果。 class SematicEmbbedBlock (nn.Module) :def __init__ (self, high_in_plane, low_in_plane, out_plane) :3 , 1 , 1 )2 )1 )def forward (self, high_x, low_x) :return high_x * low_xclass KeyPointModel (nn.Module) :''' def __init__ (self) :3 , 6 , 3 , 1 , 1 )6 )True )2 , 2 ))6 , 12 , 3 , 1 , 1 )12 )True )2 , 2 ))12 , 20 , 3 , 1 , 1 )20 )True )2 , 2 ))20 , 40 , 3 , 1 , 1 )40 )True )40 , 20 , 20 )20 , 12 , 12 )12 , 6 , 6 )6 , 1 , 1 )def forward (self, x) :return out網(wǎng)絡(luò)模型也是自己寫的小網(wǎng)絡(luò)，用了四個卷積層，三個池化層，然后進(jìn)行了三次上采樣。最終輸出分辨率和輸入分辨率相同。 訓(xùn)練過程和基于回歸的方法幾乎一樣，代碼如下： datasets = KeyPointDatasets(root_dir='./data' , transforms=transforms_all)True ,if torch.cuda.is_available():3e-3 )# compute_loss 30 ,0.1 )for epoch in range(total_epoch):if epoch % 5 == 0 :'weights/epoch_%d_%.3f.pt' % (epoch, loss*10000 ))
用的是MSELoss進(jìn)行監(jiān)督，訓(xùn)練曲線如下： 訓(xùn)練過程中的loss曲線 3.4 測試過程 測試過程和CenterNet的推理過程一致，也用到了3x3的maxpooling來篩選極大值點。 for iter, (image, label) in enumerate(dataloader):# print(image.shape) 0 ]for i in range(bs):0 )255 * hm)0 ]# heatmap = torch.sigmoid(heatmap) # hm = cv2.cvtColor(hm, cv2.COLOR_RGB2BGR) './test_output/output_%d_%d.jpg' % (iter, i), hm)0 )以上的nms和topk代碼都在CenterNet系列最后一篇講過了。這里直接對模型輸出結(jié)果使用nms，然后進(jìn)行可視化，結(jié)果如下： 放大結(jié)果 上圖中白色的點就是目標(biāo)位置，為了更形象的查看結(jié)果，detect.py部分負(fù)責(zé)可視化。 3.5 可視化 可視化的問題經(jīng)常遇見，比如CAM、Grad CAM等可視化特征圖的時候就會碰到。以下是可視化的一個簡單的方法（參考了CSDN的一位博主的方案，具體鏈接因太過久遠(yuǎn)找不到了）。 可視化流程 def normalization (data) :return (data - np.min(data)) / _rangefor i in range(bs):480 , 360 ))0 )255 * hm)480 , 360 ))0.2 + img2 , (0 , 0 , 0 ), thickness=-1 )'./output2/%s_out.jpg' % (img_name_list[i]), superimposed_img)
注意通過處理以后的hm和原圖疊加的時候0.2只是一個參考值，這個值既不會影響原圖顯示又能將heatmap中重點關(guān)注的位置可視化出來。 可視化結(jié)果 可以看到，定位結(jié)果要比回歸更準(zhǔn)一些，圖中黑色點是獲取到最終坐標(biāo)的位置，幾乎和目標(biāo)是重疊的狀態(tài)，效果比較理想。 4 總結(jié) 筆者做這個小項目初心是想搞清楚如何用關(guān)鍵點進(jìn)行定位的。關(guān)鍵點被用在很多領(lǐng)域比如人臉關(guān)鍵點定位、車牌定位、人體姿態(tài)檢測、目標(biāo)檢測等等領(lǐng)域。當(dāng)時用小武的數(shù)據(jù)的時候，發(fā)現(xiàn)這個數(shù)據(jù)集的特點就是目標(biāo)很小，比較適合用關(guān)鍵點來做。之后又開始陸陸續(xù)續(xù)看CenterNet源碼，借鑒了其中很多代碼，這才完成了這個小項目。 由于水平有限，可能使用heatmap進(jìn)行關(guān)鍵點定位的方式有些地方并不合理，如有建議可與作者溝通。
