概要
本記事では、Crowd Counting（群集カウント、群衆カウント）という、機械学習を利用した群衆画像内の人数推定モデルの実装について説明していきます。実装するモデルは、この道では（おそらく）わりと有名なCSRNetというCNNベースの密度推定モデルを使用します。すべてのコードをじっくり解説、というよりはCrowd Countingに興味はあるけど何から手をつけて良いかさっぱり。。という初級者、中級者の学習や研究のはじめの一歩として見ていただければと思います。
なお、本記事はこちらのページの実装部分の内容を日本語で説明しているものになります。英語が読める方は元の記事にも目を通してみてください。
環境の準備
jupyter notebook形式（.ipynb）のファイルを扱うので、jupyter notebookが使用できるように環境の準備をしてください。いろんなライブラリを使用するのでAnacondaをインストールし、仮想環境上での実装を推奨します。
実装にはPyTorchとGPUを使用するのであらかじめGPU環境の準備（CUDAなど）をお願い致します。
自分が実装に使用した環境はPython3.8.0, pytorch1.7.1ですが、必ずしも全く同じバージョンである必要はないと思います。
ただ、Python3を使用する前提でコードの修正などの説明をするので、Python2環境を使用する場合などはその限りではないことをご承知おきください。
CSRNetについて
CSRNetは2018年に発表されたCrowd CountingのためのCNNベースの密度推定モデルです。密度推定モデルとは、群集の画像から密集度をヒートマップとして表現した群集密度マップ（density map）を作成し、ピクセル単位で数の計測を行うモデルです。
CSRNetは特徴量抽出に特化したフロントエンドのCNNと、プーリング演算の役割を担うバックエンドのCNNを組み合わせたネットワーク構造が特徴であり、それまでの手法による精度を大幅に更新しています。
CSRNetを何もない状態からすべて実装するのはかなり大変なので、今回はGitHubのコードを使用して実装を行います。GitHub上のコードを使用する場合はgit cloneコマンドが便利です。ご使用のコマンドプロンプト（ターミナル）で以下を入力します。
$ git clone https://github.com/leeyeehoo/CSRNet-pytorch.git
git cloneコマンドの詳細はここでは省略しますが、第1引数にクローンしたいリポジトリを、第2引数にクローン先のディレクトリを指定することで、指定したディレクトリにGitHub上のリポジトリ（コードなど）を複製します。なお、第2引数は省略可です。
（Command 'git' not found...などと返ってきたら、gitライブラリをインストールする必要があります。conda（またはpip） install gitなどでパッケージをインストールしてください。）
以下、このディレクトリを作成済みという前提で説明が進みます。
データセットの準備
実装するモデルは教師あり学習なので、ラベル（アノテーション、正解）付きのデータセットが必要です。今回はShanghaiTech Datasetというデータセットを使用します。このデータセットは2016年に上海科技大学のYingying Zhangらによって発表されたCrowd Counting用データセットで、Part_AとPart_Bの2つに分けられています。本記事投稿時点では、こちらのkaggleのページからダウンロードすることができます。（kaggleのアカウント登録とサインインが必要です。もちろん無料。）
Ground-truth density mapを生成する
まずは密度マップの正解データ（Ground-truth）を生成します。make_dataset.ipynbを使用しますが、Python2で書かれているため、いくつか訂正箇所があります。全ての修正箇所に対して詳細に解説はしきれないので、各自下の修正済みコードを見て変更箇所を探して都度修正するか、下のコードを写経するか、コピペしてください。
まず、1番上のセルを実行してください。command not found系のエラーが出たら、各自で必要に応じてcondaやpipでライブラリのインストールをしてください。
# ライブラリのインポート
import h5py
import scipy.io as io
import PIL.Image as Image
import numpy as np
import os
import glob
from matplotlib import pyplot as plt
from scipy.ndimage.filters import gaussian_filter
import scipy
import scipy.spatial
import json
from matplotlib import cm as CM
from image import *
from model import CSRNet
import torch
from tqdm import tqdm
%matplotlib inline
1番目のセルが実行できたら、2番目のセルを実行してください。おそらくSyntaxErrorが出たと思います。printの中身を（）で囲む必要があります。この修正はこの先いくつかの箇所で必要なので、その都度修正してください。
# 画像から密度マップを作成するための関数
def gaussian_filter_density(gt):
print (gt.shape)
density = np.zeros(gt.shape, dtype=np.float32)
gt_count = np.count_nonzero(gt)
if gt_count == 0:
return density
pts = np.array(list(zip(np.nonzero(gt)[1], np.nonzero(gt)[0])))
leafsize = 2048
# build kdtree
tree = scipy.spatial.KDTree(pts.copy(), leafsize=leafsize)
# query kdtree
distances, locations = tree.query(pts, k=4)
print ('generate density...')
for i, pt in enumerate(pts):
pt2d = np.zeros(gt.shape, dtype=np.float32)
pt2d[pt[1],pt[0]] = 1.
if gt_count > 1:
sigma = (distances[i][1]+distances[i][2]+distances[i][3])*0.1
else:
sigma = np.average(np.array(gt.shape))/2./2. #case: 1 point
density += scipy.ndimage.filters.gaussian_filter(pt2d, sigma, mode='constant')
print ('done.')
return density
3番目のセルでは、rootに自分がデータセットをダウンロードした場所のパスを代入します。このrootは次のセルで使用します。ここはコピペせずに各自ShanghaiTechを保存してあるフォルダの絶対パスを入力してください。
# データセットのパスを入力
root = '/home/leeyh/Downloads/CSRNet-pytorch/archive/ShanghaiTech/'
4番目のセルではpart_Aとpart_Bの学習データ・テストデータそれぞれのパスに修正が必要です。（上記のkaggleからデータセットをインストールする方法以外でデータセットを入手した場合は、特にパスに誤りがないか確認してください）
part_A_train = os.path.join(root,'part_A/train_data','images')
part_A_test = os.path.join(root,'part_A/test_data','images')
part_B_train = os.path.join(root,'part_B/train_data','images')
part_B_test = os.path.join(root,'part_B/test_data','images')
path_sets = [part_A_train,part_A_test]
5番目のセル。ここは修正なし
img_paths = []
for path in path_sets:
for img_path in glob.glob(os.path.join(path, '*.jpg')):
img_paths.append(img_path)
6番目のセル。エラーがなければ、このコードでPart_Aの密度マップの正解データ（Ground-truth density map）を生成します。実行完了までしばらくかかるので、コーヒー片手にリラックスしときましょう。
for img_path in img_paths:
print (img_path)
mat = io.loadmat(img_path.replace('.jpg','.mat').replace('images','ground-truth').replace('IMG_','GT_IMG_'))
img= plt.imread(img_path)
k = np.zeros((img.shape[0],img.shape[1]))
gt = mat["image_info"][0,0][0,0][0]
for i in range(0,len(gt)):
if int(gt[i][1])<img.shape[0] and int(gt[i][0])<img.shape[1]:
k[int(gt[i][1]),int(gt[i][0])]=1
k = gaussian_filter_density(k)
with h5py.File(img_path.replace('.jpg','.h5').replace('images','ground-truth'), 'w') as hf:
hf['density'] = k
7~9番目のセルではサンプルとしてPart_Aの画像・密度マップ・正解の人数が表示されます。
plt.imshow(Image.open(img_paths[0]))
gt_file = h5py.File(img_paths[0].replace('.jpg','.h5').replace('images','ground-truth'),'r')
groundtruth = np.asarray(gt_file['density'])
plt.imshow(groundtruth,cmap=CM.jet)
np.sum(groundtruth)
10~12番目のセル。Part_Bも同様の手順で密度マップの正解データ（Ground-truth density map）の生成を行います。生成にはPart_Aのときと同様にけっこうな時間がかかります。実行完了したらGround-truth density mapの生成は完了です。
# パスをSHT_Bのものに変更
path_sets = [part_B_train,part_B_test]
img_paths = []
for path in path_sets:
for img_path in glob.glob(os.path.join(path, '*.jpg')):
img_paths.append(img_path)
# SHT_Bの画像から密度マップを生成
for img_path in img_paths:
print (img_path)
mat = io.loadmat(img_path.replace('.jpg','.mat').replace('images','ground-truth').replace('IMG_','GT_IMG_'))
img= plt.imread(img_path)
k = np.zeros((img.shape[0],img.shape[1]))
gt = mat["image_info"][0,0][0,0][0]
for i in range(0,len(gt)):
if int(gt[i][1])<img.shape[0] and int(gt[i][0])<img.shape[1]:
k[int(gt[i][1]),int(gt[i][0])]=1
k = gaussian_filter_density(k)
with h5py.File(img_path.replace('.jpg','.h5').replace('images','ground-truth'), 'w') as hf:
hf['density'] = k
モデルの学習（訓練）
train.pyを使用してモデルの学習を行います。モデルが生成した密度マップと、さっきまで生成していた密度マップの正解データとの誤差を最小化することで学習を実行しています。
少しだけmodel.pyとimage.pyの内容を修正します。model.pyは18~19行目を、image.pyは次のように変更してください。
コードを見る
for i in xrange(len(self.frontend.state_dict().items())):
self.frontend.state_dict().items()[i][1].data[:] = mod.state_dict().items()[i][1].data[:]
↓
for i in range(len(self.frontend.state_dict().items())):
list(self.frontend.state_dict().items())[i][1].data[:] = list(mod.state_dict().items())[i][1].data[:]
gt_path = img_path.replace('.jpg','.h5').replace('images','ground_truth')
↓
gt_path = img_path.replace('.jpg','.h5').replace('images','ground-truth')
target = cv2.resize(target,(target.shape[1]/8,target.shape[0]/8),interpolation = cv2.INTER_CUBIC)*64
↓
target = cv2.resize(target,(target.shape[1]//8,target.shape[0]//8),interpolation = cv2.INTER_CUBIC)*64
git clone先のディレクトリ上で次のコマンドを実行し、モデルを学習させます。（初めて学習を実行する際には、同ディレクトリ内のjsonファイルを書き換える必要があります。コマンドを実行する前に下の「jsonファイルの書き換えに関して」を読んで下さい。）
$ cd CSRNet-pytorch
$ python train.py part_A_train.json part_A_val.json 0 0
shanghaitech_dataset part_Bを学習に使用したい場合は、上記のpart_Aの部分をpart_Bに置き換えれば良いと思います。3番目以降の変数はGPU番号などですが、よほど特殊な状況でない限りはそのままで大丈夫なはずです。
.jsonファイルの書き換えに関して
上記のgit cloneによってディレクトリを作成していれば、ディレクトリ内にshanghaitech_dataset用のjsonファイルが複数あると思います。
jsonファイルには各画像データの絶対パスがリストで書かれています。しかしデフォルトの状態では絶対パスが自分の環境のものと異なるため実行できません。適宜自分の環境に適した絶対パスの書きかえを行ってください。最近のテキストエディタであれば、同じ文字列の一括編集（同時編集）が出来ると思うので上手く活用すると良いでしょう。
（書き換え完了後に上のコードを実行してJSONDecodeErrorが出力されたら、書き換えの際に余計な文字などが入っていてjsonファイルに問題がある可能性が高いです。）
モデルの評価・検証
モデルの学習が完了していると使用しているディレクトリ上にmodel_best.pth.tarみたいな名前のファイルが出来ているかと思います。その中に学習したモデルのパラメータ（重み）が保存されています。val.ipynbを使用してこのモデルの評価をしてみましょう。val.ipynbもいくつか修正箇所があるので、make_dataset.ipynbのときと同様に各自で変更をしてください。
最初の2つのセルはそのままで大丈夫です。
#ライブラリのインポート
import h5py
import scipy.io as io
import PIL.Image as Image
import numpy as np
import os
import glob
from matplotlib import pyplot as plt
from scipy.ndimage.filters import gaussian_filter
import scipy
import json
import torchvision.transforms.functional as F
from matplotlib import cm as CM
from image import *
from model import CSRNet
import torch
%matplotlib inline
from torchvision import datasets, transforms
transform=transforms.Compose([
transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
次の2つはデータセットのパスに合わせて、各自変更をお願いします。path_setsの中身には評価したいデータセットを入力します。
# データセットのパスを入力
root = '/home/leeyh/Downloads/CSRNet-pytorch/archive/ShanghaiTech/'
part_A_train = os.path.join(root,'part_A/train_data','images')
part_A_test = os.path.join(root,'part_A/test_data','images')
part_B_train = os.path.join(root,'part_B/train_data','images')
part_B_test = os.path.join(root,'part_B/test_data','images')
path_sets = [part_A_test]
次も特に変更なしです。
img_paths = []
for path in path_sets:
for img_path in glob.glob(os.path.join(path, '*.jpg')):
img_paths.append(img_path)
model = CSRNet()
model = model.cuda()
次のcheckpointの中には、自分のディレクトリ上にある学習済みモデルのパスを入力してください。
# 学習したパラメータ（重み）を代入
checkpoint = torch.load('part_A/0model_best.pth.tar')
model.load_state_dict(checkpoint['state_dict'])
次のコードで指定したデータセットのMAE（絶対平均誤差）を算出します。正しく学習が完了していれば、Part_Aのテストデータなら70前後、Part_Bのテストデータなら10前後になるかと思います。
mae = 0
for i in tqdm(range(len(img_paths))):
img = transform(Image.open(img_paths[i]).convert('RGB')).cuda()
gt_file = h5py.File(img_paths[i].replace('.jpg','.h5').replace('images','ground-truth'),'r')
groundtruth = np.asarray(gt_file['density'])
output = model(img.unsqueeze(0))
mae += abs(output.detach().cpu().sum().numpy()-np.sum(groundtruth))
print (mae/len(img_paths))
次のコードを実行すると、指定したテストデータ画像のモデルによる推定人数、モデルが作成した密度マップ、正解の人数、密度マップの正解データ、元のテストデータ画像が出力されます。2,9,15行目のパスは適宜変更してください。
from matplotlib import cm as c
img = transform(Image.open('part_A/test_data/images/IMG_100.jpg').convert('RGB')).cuda()
output = model(img.unsqueeze(0))
print("Predicted Count : ",int(output.detach().cpu().sum().numpy()))
temp = np.asarray(output.detach().cpu().reshape(output.detach().cpu().shape[2],output.detach().cpu().shape[3]))
plt.imshow(temp,cmap = c.jet)
plt.show()
temp = h5py.File('part_A/test_data/ground-truth/IMG_100.h5', 'r')
temp_1 = np.asarray(temp['density'])
plt.imshow(temp_1,cmap = c.jet)
print("Original Count : ",int(np.sum(temp_1)) + 1)
plt.show()
print("Original Image")
plt.imshow(plt.imread('part_A/test_data/images/IMG_100.jpg'))
plt.show()
次のコードは上のコードを参考に作成した、未知の画像に対する学習済みモデルの推定人数、モデルが作成した密度マップ、元の画像を出力するコードです。人数を推定したい画像のパスを'org_data.jpg'の部分に入力してください。
from matplotlib import cm as c
img = transform(Image.open('org_data.jpg').convert('RGB')).cuda()
output = model(img.unsqueeze(0))
print("Predicted Count : ",int(output.detach().cpu().sum().numpy()))
temp = np.asarray(output.detach().cpu().reshape(output.detach().cpu().shape[2],output.detach().cpu().shape[3]))
plt.imshow(temp,cmap = c.jet)
plt.show()
print("Original Image")
plt.imshow(plt.imread('org_data.jpg'))
plt.show()
最後に
ここまで読んでいただきありがとうございます。自分なりに説明したつもりではありますが、たくさんのコードを使用しているのでどこかしらで予期せぬエラーが発生してしまったかもしれません。それでもめげずに試行錯誤する過程にこそ、学びがあると思います。自分の環境ではなぜかtqdmが上手く動作しなかったので、tqdmを全てのコードから外しました。こんな感じで臨機応変に対応していただければと思います。