174 lines
6.3 KiB
Python
174 lines
6.3 KiB
Python
"""
|
|
The TransCrowd paper lists ShanghaiTech dataset as from here:
|
|
|
|
https://drive.google.com/file/d/1CkYppr_IqR1s6wi53l2gKoGqm7LkJ-Lc/view
|
|
|
|
Alternatively, you could prob. download from here:
|
|
|
|
https://www.kaggle.com/datasets/tthien/shanghaitech?resource=download
|
|
|
|
It seems the directories are all wrong, though.
|
|
|
|
After downloading, execute:
|
|
|
|
$ unzip <downloaded-zip-file> -d <repo-dir>/synchronous/dataset/
|
|
|
|
To unzip the dataset correctly prior to running this script.
|
|
"""
|
|
|
|
import os
|
|
import glob
|
|
import random
|
|
|
|
import numpy as np
|
|
import cv2
|
|
import scipy.io as io
|
|
import scipy.sparse as sparse
|
|
import h5py
|
|
from tqdm.auto import tqdm
|
|
|
|
CWD = os.getcwd()
|
|
|
|
def pre_dataset_sh():
|
|
dataset_name = "ShanghaiTech"
|
|
root = CWD + "/synchronous/dataset/" + dataset_name + "/"
|
|
|
|
part_A_train = os.path.join(root, "part_A_final/train_data", "images")
|
|
part_A_test = os.path.join(root, "part_A_final/test_data", "images")
|
|
part_B_train = os.path.join(root, "part_B_final/train_data", "images")
|
|
part_B_test = os.path.join(root, "part_B_final/test_data", "images")
|
|
|
|
# Create cropped (to 1152x768) dataset directories
|
|
for base_path in part_A_train, part_A_test, part_B_train, part_B_test:
|
|
for replacement in "images_crop", "gt_density_map_crop":
|
|
if not os.path.exists(base_path.replace("images", replacement)):
|
|
os.makedirs(base_path.replace("images", replacement))
|
|
|
|
# Gather all jpg paths in part A & B, train & test
|
|
img_paths = []
|
|
for path in part_A_train, part_A_test, part_B_train, part_B_test:
|
|
for img_path in glob.glob(os.path.join(path, "*.jpg")):
|
|
img_paths.append(img_path)
|
|
|
|
# np.random.seed(0)
|
|
# random.seed(0)
|
|
for _, img_path in tqdm(img_paths, desc="Preprocessing Data"):
|
|
img_data = cv2.imread(img_path)
|
|
mat = io.loadmat(
|
|
img_path
|
|
.replace(".jpg", ".mat")
|
|
.replace("images", "ground_truth")
|
|
.replace("IMG_", "GT_IMG_")
|
|
)
|
|
gt_data = mat["image_info"][0][0][0][0][0]
|
|
|
|
# Resize to 1152x768
|
|
is_portrait = False
|
|
if img_data.shape[1] >= img_data.shape[0]: # landscape
|
|
rate_x = 1152.0 / img_data.shape[1]
|
|
rate_y = 768.0 / img_data.shape[0]
|
|
else: # portrait
|
|
rate_x = 768.0 / img_data.shape[1]
|
|
rate_y = 1152.0 / img_data.shape[0]
|
|
is_portrait = True
|
|
|
|
img_data = cv2.resize(img_data, (0, 0), fx=rate_x, fy=rate_y)
|
|
gt_data[:, 0] = gt_data[:, 0] * rate_x
|
|
gt_data[:, 1] = gt_data[:, 1] * rate_y
|
|
|
|
if is_portrait:
|
|
print("Portrait img: \'{}\' -- transposing...".format(img_path))
|
|
img_data = cv2.transpose(img_data)
|
|
gt_data = gt_data[:, ::-1]
|
|
|
|
# Compute 0/1 counts from density map
|
|
assert img_data.shape[:2] == (768, 1152)
|
|
coordinates = gt_data.round().astype(int) # To integer coords
|
|
coordinates[:, 0] = np.clip(coordinates[:, 0], a_min=0, a_max=1151)
|
|
coordinates[:, 1] = np.clip(coordinates[:, 1], a_min=0, a_max=767)
|
|
assert max(coordinates[:, 0]) < 1152
|
|
assert max(coordinates[:, 1]) < 768
|
|
sparse_mat = sparse.coo_matrix((
|
|
np.ones(coordinates.shape[0]), # data
|
|
(coordinates[:, 1], coordinates[:, 0]), # (i, j)
|
|
), # N.B. all k |- ret[i[k], j[k]] = data[k]
|
|
shape=(768, 1152),
|
|
dtype=int,
|
|
) # To same shape as image, so i, j flipped wrt. coordinates
|
|
kpoint = sparse_mat.toarray()
|
|
|
|
# Sum count as ground truth (we need to train STN, remember?)
|
|
gt_count = sparse_mat.nnz
|
|
|
|
fname = img_path.split("/")[-1]
|
|
root_path = img_path.split("IMG_")[0].replace("images", "images_crop")
|
|
|
|
# Likewise, we do not crop to patched sequences here...
|
|
# Skip directly to saving fixed-size data & kpoint.
|
|
img_path = img_path.replace("images", "images_crop")
|
|
cv2.imwrite(img_path, img_data)
|
|
with h5py.File(
|
|
img_path.replace('.jpg', '.h5').replace('images', 'gt_density_map'),
|
|
mode='w'
|
|
) as hf:
|
|
hf["kpoint"] = kpoint
|
|
hf["gt_count"] = gt_count
|
|
|
|
|
|
def make_npydata():
|
|
if not os.path.exists("./npydata"):
|
|
os.makedirs("./npydata")
|
|
|
|
try:
|
|
sh_A_train_path = CWD + '/synchronous/dataset/ShanghaiTech/part_A_final/train_data/images_crop/'
|
|
sh_A_test_path = CWD + '/synchronous/dataset/ShanghaiTech/part_A_final/test_data/images_crop/'
|
|
|
|
train_fpaths = []
|
|
for fname in os.listdir(sh_A_train_path):
|
|
if fname.split(".")[1] == "jpg":
|
|
train_fpaths.append(sh_A_train_path + fname)
|
|
train_fpaths.sort()
|
|
np.save("./npydata/ShanghaiA_train.npy", train_fpaths)
|
|
|
|
test_fpaths = []
|
|
for fname in os.listdir(sh_A_test_path):
|
|
if fname.split(".")[1] == "jpg":
|
|
test_fpaths.append(sh_A_test_path + fname)
|
|
test_fpaths.sort()
|
|
np.save("./npydata/ShanghaiA_test.npy", test_fpaths)
|
|
|
|
print("Saved ShanghaiA image list (test: {} | train: {})".format(
|
|
len(test_fpaths), len(train_fpaths)
|
|
))
|
|
except:
|
|
print("The ShanghaiA dataset path is wrong.")
|
|
|
|
try:
|
|
sh_B_train_path = CWD + '/synchronous/dataset/ShanghaiTech/part_B_final/train_data/images_crop/'
|
|
sh_B_test_path = CWD + '/synchronous/dataset/ShanghaiTech/part_B_final/test_data/images_crop/'
|
|
|
|
train_fpaths = []
|
|
for fname in os.listdir(sh_B_train_path):
|
|
if fname.split(".")[1] == "jpg":
|
|
train_fpaths.append(sh_B_train_path + fname)
|
|
train_fpaths.sort()
|
|
np.save("./npydata/ShanghaiB_train.npy", train_fpaths)
|
|
|
|
test_fpaths = []
|
|
for fname in os.listdir(sh_B_test_path):
|
|
if fname.split(".")[1] == "jpg":
|
|
test_fpaths.append(sh_B_test_path + fname)
|
|
test_fpaths.sort()
|
|
np.save("./npydata/ShanghaiB_test.npy", test_fpaths)
|
|
|
|
print("Saved ShanghaiB image list (test: {} | train: {})".format(
|
|
len(test_fpaths), len(train_fpaths)
|
|
))
|
|
except:
|
|
print("The ShanghaiB dataset path is wrong.")
|
|
|
|
if __name__ == "__main__":
|
|
# Download manually...
|
|
pre_dataset_sh() # XXX: preliminary
|
|
print("Storing dataset paths...")
|
|
make_npydata()
|