mirror of
https://github.com/sbrl/research-rainfallradar
synced 2024-11-25 10:32:59 +00:00
batch data; use generator
This commit is contained in:
parent
d306853c42
commit
cc6679c609
1 changed files with 21 additions and 18 deletions
|
@ -7,6 +7,7 @@ import re
|
|||
|
||||
from loguru import logger
|
||||
import tensorflow as tf
|
||||
from lib.dataset.batched_iterator import batched_iterator
|
||||
|
||||
from lib.vis.segmentation_plot import segmentation_plot
|
||||
from lib.io.handle_open import handle_open
|
||||
|
@ -87,10 +88,12 @@ def run(args):
|
|||
sys.stderr.write(">>> Complete\n")
|
||||
|
||||
def do_png(args, ai, dataset, model_code):
|
||||
i = 0
|
||||
for rainfall, water in dataset:
|
||||
water_predict = ai.embed(rainfall)
|
||||
model_params = json.loads(readfile(args.params))
|
||||
|
||||
i = 0
|
||||
for rainfall, water in batched_iterator(dataset, tensors_in_item=2, batch_size=model_params["batch_size"]):
|
||||
water_predict_batch = ai.embed(rainfall)
|
||||
for water_predict in water_predict_batch:
|
||||
# [ width, height, softmax_probabilities ] → [ batch, width, height ]
|
||||
water_predict = tf.math.argmax(water_predict, axis=-1)
|
||||
# [ width, height ]
|
||||
|
|
Loading…
Reference in a new issue