Heatmap_3D_Multiple.py 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  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. # import seaborn as sns
  7. # import matplotlib.patches as patches
  8. # import matplotlib.colors as colors
  9. # from sklearn import preprocessing
  10. # from pylab import *
  11. # from matplotlib import style
  12. from matplotlib.colors import LinearSegmentedColormap
  13. from mpl_toolkits.mplot3d import Axes3D
  14. WIDTH = int(70)
  15. HEIGHT = int(35)
  16. OBSTACLE_PATH = "Assets/Data_image/obstacle.pkl"
  17. POSITION_PATH1 = ue.Application.dataPath + '/Data_position/80/Walk1.csv'
  18. POSITION_PATH2 = ue.Application.dataPath + '/Data_position/80/Walk2.csv'
  19. POSITION_PATH3 = ue.Application.dataPath + '/Data_position/80/Walk4.csv'
  20. HEATMAP_PATH = "Assets/Data_image/80/heatmap3DMultiple1.png"
  21. # 1. Get position data from csv file
  22. data1 = pd.read_csv(POSITION_PATH1, sep=';', usecols=["Position x", "Position z"], decimal=',', dtype={'Position x': float, 'Position z': float})
  23. data1 = data1.round(0)
  24. data2 = pd.read_csv(POSITION_PATH2, sep=';', usecols=["Position x", "Position z"], decimal=',', dtype={'Position x': float, 'Position z': float})
  25. data2 = data2.round(0)
  26. data3 = pd.read_csv(POSITION_PATH3, sep=';', usecols=["Position x", "Position z"], decimal=',', dtype={'Position x': float, 'Position z': float})
  27. data3 = data3.round(0)
  28. # 2. Group by positions and count appearance
  29. data_count1 = data1.groupby(['Position x', 'Position z']).size().reset_index(name='counts')
  30. data_count2 = data2.groupby(['Position x', 'Position z']).size().reset_index(name='counts')
  31. data_count3 = data3.groupby(['Position x', 'Position z']).size().reset_index(name='counts')
  32. # 3.1 Assign x, y, z, width, depth, height
  33. x1 = data_count1["Position x"].tolist()
  34. y1 = data_count1["Position z"].tolist()
  35. z1 = np.zeros_like(len(x1))
  36. dz1 = data_count1["counts"].tolist() # Change
  37. x2 = data_count2["Position x"].tolist()
  38. y2 = data_count2["Position z"].tolist()
  39. z2 = np.zeros_like(len(x2))
  40. dz2 = data_count2["counts"].tolist() # Change
  41. x3 = data_count3["Position x"].tolist()
  42. y3 = data_count3["Position z"].tolist()
  43. z3 = np.zeros_like(len(x3))
  44. dz3 = data_count3["counts"].tolist() # Change
  45. # 3.2 Add offset to day2 and day3
  46. x2[:] = [a+0.5 for a in x2[:]]
  47. x3[:] = [a+0.5 for a in x3[:]]
  48. y3[:] = [a+0.5 for a in y3[:]]
  49. # 4. Create figure and axes
  50. fig = plt.figure(figsize=(10,10))
  51. ax = fig.add_subplot(111, projection='3d')
  52. # 5.1 Create custom colormap Day1, Day2, Day3
  53. cmap1 = LinearSegmentedColormap.from_list(name='day1', colors=[(0.40,0.76,0.65), (0.11,0.62,0.47)])
  54. cmap2 = LinearSegmentedColormap.from_list(name='day2', colors=[(0.99,0.55,0.38), (0.85,0.37,0.01)])
  55. cmap3 = LinearSegmentedColormap.from_list(name='day3', colors=[(0.55,0.63,0.80), (0.46,0.44,0.70)])
  56. # 5.2 Initialize array for coloring the bars
  57. dz_array1 = np.array(data_count1['counts'])
  58. fracs1 = dz_array1.astype(float) / dz_array1.max()
  59. color_values1 = cmap1(fracs1.tolist())
  60. dz_array2 = np.array(data_count2['counts'])
  61. fracs2 = dz_array2.astype(float) / dz_array2.max()
  62. color_values2 = cmap2(fracs2.tolist())
  63. dz_array3 = np.array(data_count3['counts'])
  64. fracs3 = dz_array3.astype(float) / dz_array3.max()
  65. color_values3 = cmap3(fracs3.tolist())
  66. # 6. Create the bars
  67. img = ax.bar3d(x1, y1, z1, 0.5, 0.5, dz1, color=color_values1, shade=False)
  68. img = ax.bar3d(x2, y2, z2, 0.5, 0.5, dz2, color=color_values2, shade=False)
  69. img = ax.bar3d(x3, y3, z3, 0.5, 0.5, dz3, color=color_values3, shade=False)
  70. # 7. Create Colorbar
  71. # color_map1 = cm.ScalarMappable(cmap=cmap1)
  72. # color_map1.set_array(dz1)
  73. # fig.colorbar(color_map1)
  74. # color_map2 = cm.ScalarMappable(cmap=cmap2)
  75. # color_map2.set_array(dz2)
  76. # fig.colorbar(color_map2)
  77. # color_map3 = cm.ScalarMappable(cmap=cmap3)
  78. # color_map3.set_array(dz3)
  79. # fig.colorbar(color_map3)
  80. # 8.
  81. plt.show()
  82. # 10. Save 3D Heatmap
  83. # heatmap.get_figure().savefig(HEATMAP_PATH, transparent=True)
  84. fig.savefig(HEATMAP_PATH, transparent=True)