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 cv2
import scipy.io as io
import scipy.sparse as sparse
import h5py
CWD = os.getcwd()
@ -75,16 +76,25 @@ def pre_dataset_sh():
gt_data[:, 1] = gt_data[:, 1] * rate_y
if is_portrait:
print("Portrait img: \'{}\' -- rotating 90 deg clockwise...".format(img_path))
img_data = cv2.rotate(img_data, cv2.ROTATE_90_CLOCKWISE)
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
kpoint = np.zeros((img_data.shape[0], img_data.shape[1]))
for i in range(len(gt_data)):
if ( int(gt_data[i][1]) < img_data.shape[0]
and int(gt_data[i][0]) < img_data.shape[1]):
kpoint[int(gt_data[i][1]), int(gt_data[i][0])] = 1
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()
fname = img_path.split("/")[-1]
root_path = img_path.split("IMG_")[0].replace("images", "images_crop")