add utils

utils.py

@@ -0,0 +1,155 @@
+import numpy as np
+import imageio
+import pickle
+import tensorflow as tf
+import matplotlib.pyplot as plt
+import cartopy.crs as ccrs
+import matplotlib.ticker as mticker
+from cartopy.mpl.gridliner import LATITUDE_FORMATTER, LONGITUDE_FORMATTER
+import shapefile as shp
+from matplotlib import animation
+from IPython.display import HTML
+
+def get_array(source, scaler_dict):
+
+    reader = imageio.get_reader(source)
+
+    source_video = []
+
+    try:
+        for im in reader:
+            source_video.append(im)
+    except RuntimeError:
+        pass
+
+    reader.close()
+
+    scaler_path = scaler_dict[source[-14:]]
+
+    with open(scaler_path, 'rb') as f:
+        sc = pickle.load(f)
+
+    data = np.array(source_video)[:,:,:,0]
+    data = sc.inverse_transform(data)
+    data = np.swapaxes(data, 0, 2)
+    data = np.swapaxes(data, 0, 1)
+
+    X = data[:,:,0:12]
+    y = data[:,:,12:]
+
+    return X, y
+
+def get_slices(values, slices):
+
+    dim_size = len(values)
+    idx_step = int(dim_size/slices)
+
+    slices_list = []
+
+    for i in range(idx_step, dim_size, idx_step):
+
+        slices_list.append(np.round(values[i], 2))
+    
+    return slices_list
+
+def save_video(X, threshold=0, file_path = 'data/video.mp4'):
+
+    # Get vmax
+    var = X.copy()
+
+    var[np.isnan(var)] = 0
+    var[var<=0] = 0
+
+    counts, bins = np.histogram(var[:])
+
+    value = counts[counts>np.median(counts)][-1]
+    idx = np.where(counts==value)[0][0]
+
+    vmax = np.round(bins[idx])
+    
+    # Latitude and longitude
+    lon   = np.loadtxt('data/longitude.txt')
+    lat   = np.loadtxt('data/latitude.txt')
+
+    area  = [lon.min(),lon.max(),lat.min(),lat.max()]
+
+    lat_list = get_slices(lat, 4)
+    lon_list = get_slices(lon, 6)
+
+    # Visualization
+    ims = []
+    fig = plt.figure(figsize=(7,5))
+    ax = plt.axes(projection=ccrs.PlateCarree())
+
+    gl = ax.gridlines(crs=ccrs.PlateCarree(),
+                    draw_labels=True,
+                    linewidth=0.3,
+                    color='black',
+                    linestyle='--')
+
+    gl.top_labels = False
+    gl.right_labels = False
+    gl.xlines = True
+    gl.xlocator = mticker.FixedLocator(lon_list)
+    gl.ylocator = mticker.FixedLocator(lat_list)
+    gl.xformatter = LONGITUDE_FORMATTER
+    gl.yformatter = LATITUDE_FORMATTER
+    gl.xlabel_style = {'size':10, 'color':'black'}
+    gl.ylabel_style = {'size':10, 'color':'black'}
+
+    frames = X
+    frames[frames<=threshold] = np.nan
+
+    barra = np.arange(0, vmax+1, 5)
+
+    for i in range(frames.shape[2]):
+        im = plt.imshow(frames[..., i],
+                        cmap=plt.cm.rainbow,
+                        vmin=0,
+                        vmax=vmax,
+                        extent=area,
+                        origin='lower',
+                        animated=True)
+        
+        ims.append([im])
+
+    cbar = plt.colorbar(ax=ax, pad=0.02, aspect=16, shrink=0.77)
+    cbar.set_ticks(barra)
+    cbar.set_label('mm/h')
+
+    shapeID = shp.Reader("data/shapefile/regiao_sul.shp")
+
+    for shape in shapeID.shapeRecords():
+        point = np.array( shape.shape.points )
+        dummy = plt.plot( point[:,0] , point[:,1], color="black", linewidth=0.5 ) # 1
+
+    ani = animation.ArtistAnimation(fig, ims, interval=500, blit=True, repeat_delay=1000)
+
+    FFwriter = animation.FFMpegWriter(fps=2)
+    ani.save(file_path, writer = FFwriter)
+
+    # ani.save(f'data/vis.gif', writer='pillow', fps=6)
+
+    plt.close(ani._fig)
+    HTML(ani.to_html5_video())
+
+def make_predictions(X):
+
+    filepath = "models/model.h5"
+    model = tf.keras.models.load_model(filepath)
+
+    X[np.isnan(X)] = 0
+
+    X = np.expand_dims(X, axis=0)
+
+    scaler_path = 'models/scaler.pkl'
+    with open(scaler_path, 'rb') as f:
+        sc = pickle.load(f)
+
+    X = sc.transform(X)
+    ypred = model.predict(X)
+    print(ypred.shape)
+    ypred = sc.inverse_transform(ypred)[0]
+    print(ypred.shape)
+
+    return ypred