Heatmap_3D.py 3.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798
  1. import UnityEngine as ue
  2. import pandas as pd
  3. import numpy as np
  4. import matplotlib.pyplot as plt
  5. import matplotlib.cm as cm
  6. from matplotlib.colors import LinearSegmentedColormap
  7. from mpl_toolkits.mplot3d import Axes3D
  8. WIDTH = int(70)
  9. HEIGHT = int(35)
  10. OBSTACLE_PATH = "Assets/Data_image/obstacle.pkl"
  11. POSITION_PATH = ue.Application.dataPath + "/Data_position/80/Walk1.csv"
  12. HEATMAP_PATH = "Assets/Data_image/80/heatmap3D1.png"
  13. def get_data_hight():
  14. dataset = pd.read_csv(POSITION_PATH, sep=';', usecols=["Delta Time", "Position x", "Position z"], decimal=',', dtype={'Delta Time': float, 'Position x': float, 'Position z': float})
  15. dataset["Position x"] = dataset["Position x"].round(0)
  16. dataset["Position z"] = dataset["Position z"].round(0)
  17. dataset["Delta Time"] = dataset["Delta Time"].round(2)
  18. dataset_human_on_same_spot = dataset.groupby(["Delta Time", "Position x", "Position z"]).size().reset_index(name='AmountOneSpot')
  19. dataset_human_amount_per_deltaTime = dataset_human_on_same_spot.groupby(['Delta Time']).size().reset_index(name='counts')
  20. solution = np.zeros((HEIGHT, WIDTH))
  21. old = np.zeros((HEIGHT, WIDTH))
  22. current = np.zeros((HEIGHT, WIDTH))
  23. minStart = 0
  24. for i in range(dataset_human_amount_per_deltaTime.shape[0]):
  25. for count in range(minStart, minStart + int(dataset_human_amount_per_deltaTime['counts'][i])):
  26. x = int(dataset_human_on_same_spot['Position x'][count])
  27. y = int(dataset_human_on_same_spot['Position z'][count])
  28. if i == 0:
  29. old[y][x] = dataset_human_on_same_spot['AmountOneSpot'][count]
  30. else :
  31. current[y][x] = dataset_human_on_same_spot['AmountOneSpot'][count]
  32. if i > 0:
  33. solution = solution + np.absolute(old - current)
  34. old = current
  35. current = np.zeros((HEIGHT, WIDTH))
  36. minStart += int(dataset_human_amount_per_deltaTime['counts'][i])
  37. return solution
  38. # 1. Get position data from csv file
  39. data = pd.read_csv(POSITION_PATH, sep=';', usecols=["Position x", "Position z"], decimal=',', dtype={'Position x': float, 'Position z': float})
  40. data = data.round(0)
  41. # 2. Group by positions and count appearance
  42. data_count = data.groupby(['Position x', 'Position z']).size().reset_index(name='counts')
  43. # 3.1 Assign x, y, z
  44. x = data_count["Position x"].tolist()
  45. y = data_count["Position z"].tolist()
  46. z = np.zeros_like(len(x))
  47. # 3.2 Assign height of the bars
  48. dz_dynamic = get_data_hight()
  49. dz = []
  50. for i in range(len(x)):
  51. dz.append(dz_dynamic[int(y[i])][int(x[i])])
  52. # 4. Create figure and axes
  53. fig = plt.figure()
  54. ax = fig.add_subplot(111, projection='3d')
  55. # 5.1 Create custom colormap Day1, Day2, Day3
  56. cmap = LinearSegmentedColormap.from_list(name='day1', colors=[(0.40,0.76,0.65), (0.11,0.62,0.47)])
  57. # cmap = LinearSegmentedColormap.from_list(name='day2', colors=['grey', (0.99,0.55,0.38), (0.85,0.37,0.01)])
  58. # cmap = LinearSegmentedColormap.from_list(name='day3', colors=['grey', (0.55,0.63,0.80), (0.46,0.44,0.70)])
  59. # 5.2 Initialize array for coloring the bars
  60. dz_array = np.array(data_count['counts'])
  61. fracs = dz_array.astype(float) / dz_array.max()
  62. color_values = cmap(fracs.tolist())
  63. # 6. Create the bars
  64. # for i in range(len(x)):
  65. # img = ax.bar3d(x[i], y[i], z, 1, 1, dz[i], color=color_values[i], shade=False)
  66. img = ax.bar3d(x, y, z, 1, 1, dz, color=color_values, shade=False)
  67. # 7. Create Colorbar
  68. color_intens = data_count['counts'].tolist()
  69. color_map = cm.ScalarMappable(cmap=cmap)
  70. color_map.set_array(color_intens)
  71. fig.colorbar(color_map)
  72. # 8.
  73. ax.set_xlabel('Width')
  74. ax.set_ylabel('Height')
  75. ax.set_zlabel('Human dynamic')
  76. plt.show()
  77. # 10. Save 3D Heatmap
  78. # heatmap.get_figure().savefig(HEATMAP_PATH, transparent=True)
  79. fig.savefig(HEATMAP_PATH, transparent=True)