123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293 |
- from matplotlib.pyplot import fill
- import UnityEngine as ue
- import numpy as np
- import seaborn as sns
- import matplotlib.pylab as plt
- import matplotlib.patches as patches
- from matplotlib.colors import LinearSegmentedColormap
- import pandas as pd
- WIDTH = int(70)
- HEIGHT = int(35)
- OBSTACLE_PATH = "Assets/Data_image/obstacle.pkl"
- POSITION_PATH = ue.Application.dataPath + '/Data_position/Walk2.csv'
- HEATMAP_PATH = "Assets/Data_image/heatmap2.png"
- # Generate only if obstacles change
- def set_obstacles():
- global list_obstacle_length
- positions = pd.DataFrame(np.zeros((HEIGHT, WIDTH)))
- obstacles = ue.Object.FindObjectsOfType(ue.GameObject)
- for obstacle in obstacles:
- if(obstacle.layer == 15 or obstacle.layer == 12):
- startWidth = int(obstacle.transform.position.x - obstacle.transform.localScale.x / 2)
- endWidth = int(obstacle.transform.position.x + obstacle.transform.localScale.x / 2)
- startHeight = int(obstacle.transform.position.z - obstacle.transform.localScale.z / 2)
- endHeight = int(obstacle.transform.position.z + obstacle.transform.localScale.z / 2)
- for currentW in range(startWidth, endWidth, 1):
- for currentH in range(startHeight, endHeight, 1):
- positions[currentW][currentH] = -500
- positions.to_pickle(OBSTACLE_PATH)
- def set_patches(plt):
- obstacles = ue.Object.FindObjectsOfType(ue.GameObject)
- for obstacle in obstacles:
- if(obstacle.layer == 15):
- startWidth = int(obstacle.transform.position.x - obstacle.transform.localScale.x / 2)
- endWidth = int(obstacle.transform.position.x + obstacle.transform.localScale.x / 2)
- startHeight = int(obstacle.transform.position.z - obstacle.transform.localScale.z / 2)
- endHeight = int(obstacle.transform.position.z + obstacle.transform.localScale.z / 2)
- plt.gca().add_patch(
- patches.Rectangle(
- (startWidth, startHeight),
- endWidth - startWidth,
- endHeight - startHeight,
- fill=False,
- color='black'
- )
- )
- # 1. Get position data from csv file
- data = pd.read_csv(POSITION_PATH, sep=';', usecols=["Position x", "Position z"], decimal=',', dtype={'Position x': float, 'Position z': float})
- data = data.round(0)
- # 2. Group by positions and count appearance
- data_count = data.groupby(['Position x', 'Position z']).size().reset_index(name='counts')
- # 3. Create wide-form DataFrame for generating heatmap
- positions = data_count.loc[:,:].reset_index().pivot(index='Position z', columns='Position x', values='counts')
- # 4. Fill missing values
- positions.fillna(0, inplace=True)
- # 5. reindex DataFrame (70,35) size of Surface; (70, 35) first x width then z height
- positions = positions.reindex_axis(range(0, HEIGHT), axis=0, fill_value=0)
- positions = positions.reindex_axis(range(0, WIDTH), axis=1, fill_value=0)
- # 6. Get obstacles (obstacles, market stalls) and paste specific value in positions
- # Therefore save pkl file in folder and read from it afterwards
- # set_obstacles()
- positions_heatmap = pd.read_pickle(OBSTACLE_PATH)
- # 6.1 Merge positions data with obstacles data
- positions_heatmap.where(positions_heatmap != 0, positions, inplace=True)
- # 6.2 Debug Output
- # positions_heatmap = pd.DataFrame(positions_heatmap)
- # positions_heatmap.to_html('Assets/Data_image/positions_heatmap.html')
- # 7. Plot the heatmap
- cmap = LinearSegmentedColormap.from_list(name='greenToRed', colors=['grey', 'limegreen', 'chartreuse', 'yellow', 'darkorange', 'red'])
- heatmap = sns.heatmap(positions_heatmap, cmap=cmap, cbar=False, square=True, yticklabels=False, xticklabels=False)
- heatmap.invert_yaxis()
- plt.xlabel('')
- plt.ylabel('')
- # 9. Mark the Market stalls
- set_patches(plt)
- # 9.1
- plt.show()
- # 10. Save Heatmap
- heatmap.get_figure().savefig(HEATMAP_PATH)
|