Improved preprocess script performance

This commit is contained in:
Zhengyi Chen 2024-03-03 19:00:37 +00:00
parent 12aabb0d3f
commit a9dd8dee04

View file

@ -23,6 +23,7 @@ import random
import numpy as np import numpy as np
import cv2 import cv2
import scipy.io as io import scipy.io as io
import scipy.sparse as sparse
import h5py import h5py
CWD = os.getcwd() CWD = os.getcwd()
@ -75,16 +76,25 @@ def pre_dataset_sh():
gt_data[:, 1] = gt_data[:, 1] * rate_y gt_data[:, 1] = gt_data[:, 1] * rate_y
if is_portrait: if is_portrait:
print("Portrait img: \'{}\' -- rotating 90 deg clockwise...".format(img_path)) print("Portrait img: \'{}\' -- transposing...".format(img_path))
img_data = cv2.rotate(img_data, cv2.ROTATE_90_CLOCKWISE) img_data = cv2.transpose(img_data)
gt_data = gt_data[:, ::-1]
# Compute 0/1 counts from density map # Compute 0/1 counts from density map
kpoint = np.zeros((img_data.shape[0], img_data.shape[1])) assert img_data.shape[:2] == (768, 1152)
for i in range(len(gt_data)): coordinates = gt_data.round().astype(int) # To integer coords
if ( int(gt_data[i][1]) < img_data.shape[0] coordinates[:, 0] = np.clip(coordinates[:, 0], a_min=0, a_max=1151)
and int(gt_data[i][0]) < img_data.shape[1]): coordinates[:, 1] = np.clip(coordinates[:, 1], a_min=0, a_max=767)
kpoint[int(gt_data[i][1]), int(gt_data[i][0])] = 1 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()
fname = img_path.split("/")[-1] fname = img_path.split("/")[-1]
root_path = img_path.split("IMG_")[0].replace("images", "images_crop") root_path = img_path.split("IMG_")[0].replace("images", "images_crop")