Can evaluate the inpainting model

Author: GreenWizard2015

Core/CDataSamplerInpainting.py

@@ -50,6 +50,7 @@ def sample(self, **kwargs):
         timesteps = kwargs.get('timesteps', None)
         N = kwargs.get('N', self._batchSize) // len(self._keys)
         indexes = []
+        added = False
         for _ in range(N):
             added = False
             while not added:
@@ -62,7 +63,7 @@ def sample(self, **kwargs):
                     indexes.extend(sampledSteps)
                     added = True
                 continue
-
+        if not added: return None, 0
         return self._indexes2XY(indexes, kwargs)
 
     def sampleById(self, idx, **kwargs):
@@ -97,7 +98,7 @@ def sampleByIds(self, ids, **kwargs):
 
         res = None
         if 0 < len(sampledSteps):
-            res = self._indexes2XY(sampledSteps, kwargs)
+            res, _ = self._indexes2XY(sampledSteps, kwargs)
         return res, rejected, accepted
 
     def _indexes2XY(self, indexesAndTime, kwargs):
@@ -216,7 +217,7 @@ def _indexes2XY(self, indexesAndTime, kwargs):
             assert B == v.shape[0], f'Invalid batch size for X[{k}]: {v.shape[0]} != {B} ({v.shape})'
         for k, v in Y.items():
             assert B == v.shape[0], f'Invalid batch size for Y[{k}]: {v.shape[0]} != {B} ({v.shape})'
-        return (X, Y)
+        return (X, Y), B
 
     def merge(self, samples, expected_batch_size):
         X = {}

Core/CDatasetLoader.py

@@ -43,6 +43,8 @@ def __init__(self, folder, samplerArgs, sampling, stats, sampler_class):
       i: len(ds.validSamples())
       for i, ds in enumerate(self._datasets)
     }
+    # ignore datasets with no valid samples
+    validSamples = {k: v for k, v in validSamples.items() if 0 < v}
     dtype = np.uint8 if len(self._datasets) < 256 else np.uint32
     # create an array of dataset indices to sample from
     sampling = ESampling(sampling)
@@ -101,13 +103,15 @@ def sample(self, **kwargs):
     samples = []
     totalSamples = 0
     # find the datasets ids and the number of samples to take from each dataset
-    datasetIds, counts = self._getBatchStats(batchSize)
-    for datasetId, N in zip(datasetIds, counts):
-      dataset = self._datasets[datasetId]
-      sampled = dataset.sample(N=N, **kwargs)
-      samples.append(sampled)
-      totalSamples += N
-      continue
+    while totalSamples < batchSize:
+      datasetIds, counts = self._getBatchStats(batchSize - totalSamples)
+      for datasetId, N in zip(datasetIds, counts):
+        dataset = self._datasets[datasetId]
+        sampled, N = dataset.sample(N=N, **kwargs)
+        if 0 < N:
+          samples.append(sampled)
+        totalSamples += N
+        continue
 
     first_dataset = self._datasets[0]
     return first_dataset.merge(samples, batchSize)

Core/CInpaintingTrainer.py

@@ -45,32 +45,36 @@ def __init__(self, timesteps, model='simple', KP=5, **kwargs):
   def compile(self):
     self._optimizer = NNU.createOptimizer()
 
+  def _calcLoss(self, x, y, training):
+    losses = {}
+    x = self._model.replaceByEmbeddings(x)
+    latents = self._encoder(x, training=training)['latent']
+    decoderArgs = {
+      'keyPoints': latents,
+      'time': y['time'],
+      'userId': x['userId'],
+      'placeId': x['placeId'],
+      'screenId': x['screenId'],
+    }
+    predictions = self._decoder(decoderArgs, training=training)
+    losses = {}
+    for k in predictions.keys():
+      pred = predictions[k]
+      gt = y[k]
+      tf.assert_equal(tf.shape(pred), tf.shape(gt))
+      loss = tf.losses.mse(gt, pred)
+      losses[f"loss-{k}"] = tf.reduce_mean(loss)
+      
+    # calculate total loss and final loss
+    losses['loss'] = sum(losses.values())
+    return losses, losses['loss']
+  
   def _trainStep(self, Data):
     print('Instantiate _trainStep')
     ###############
     x, y = Data
-    losses = {}
     with tf.GradientTape() as tape:
-      x = self._model.replaceByEmbeddings(x)
-      latents = self._encoder(x, training=True)['latent']
-      decoderArgs = {
-        'keyPoints': latents,
-        'time': y['time'],
-        'userId': x['userId'],
-        'placeId': x['placeId'],
-        'screenId': x['screenId'],
-      }
-      predictions = self._decoder(decoderArgs, training=True)
-      losses = {}
-      for k in predictions.keys():
-        pred = predictions[k]
-        gt = y[k]
-        tf.assert_equal(tf.shape(pred), tf.shape(gt))
-        loss = tf.losses.mse(gt, pred)
-        losses[f"loss-{k}"] = tf.reduce_mean(loss)
-        
-      # calculate total loss and final loss
-      losses['loss'] = loss = sum(losses.values())
+      losses, loss = self._calcLoss(x, y, training=True)
 
     self._optimizer.minimize(loss, tape.watched_variables(), tape=tape)
     ###############
@@ -84,18 +88,10 @@ def fit(self, data):
 
   def _eval(self, xy):
     print('Instantiate _eval')
-    x, (y,) = xy
-    x = self._replaceByEmbeddings(x)
-    y = y[:, :, 0]
-    predictions = self._model(x, training=False)
-    points = predictions['result'][:, :, :]
-    tf.assert_equal(tf.shape(points), tf.shape(y))
-
-    loss = self._pointLoss(y, points)
-    tf.assert_equal(tf.shape(loss), tf.shape(y)[:2])
-    _, dist = NNU.normVec(points - y)
-    return loss, points, dist
+    x, y = xy
+    losses, loss = self._calcLoss(x, y, training=False)
+    return loss
 
   def eval(self, data):
-    loss, sampled, dist = self._eval(data)
-    return loss.numpy(), sampled.numpy(), dist.numpy()
+    loss = self._eval(data)
+    return loss.numpy()

Core/CTestInpaintingLoader.py

@@ -0,0 +1,26 @@
+import tensorflow as tf
+import numpy as np
+import os, glob
+from functools import lru_cache
+
+class CTestInpaintingLoader(tf.keras.utils.Sequence):
+  def __init__(self, testFolder):
+    self._batchesNpz = [
+      f for f in glob.glob(os.path.join(testFolder, 'test-*.npz'))
+    ]
+    self.on_epoch_end()
+    return
+  
+  def on_epoch_end(self):
+    return
+
+  def __len__(self):
+    return len(self._batchesNpz)
+  
+  def __getitem__(self, idx):
+    with np.load(self._batchesNpz[idx]) as res:
+      res = {k: v for k, v in res.items()}
+      
+    X = {k.replace('X_', ''): v for k, v in res.items() if 'X_' in k}
+    Y = {k.replace('Y_', ''): v for k, v in res.items() if 'Y_' in k}
+    return(X, Y)

Core/CTestLoader.py

@@ -5,25 +5,12 @@
 
 class CTestLoader(tf.keras.utils.Sequence):
   def __init__(self, testFolder):
-    self._folder = testFolder
     self._batchesNpz = [
       f for f in glob.glob(os.path.join(testFolder, 'test-*.npz'))
     ]
     self.on_epoch_end()
     return
 
-  @property
-  def folder(self):
-    return self._folder
-  
-  @lru_cache(maxsize=1)
-  def parametersIDs(self):
-    batch, _ = self[0]
-    userId = batch['userId'][0, 0, 0]
-    placeId = batch['placeId'][0, 0, 0]
-    screenId = batch['screenId'][0, 0, 0]
-    return userId, placeId, screenId
-    
   def on_epoch_end(self):
     return
 
@@ -35,15 +22,4 @@ def __getitem__(self, idx):
       res = {k: v for k, v in res.items()}
 
     Y = res.pop('y')
-    return(res, (Y, ))
-
-if __name__ == '__main__':
-  folder = os.path.dirname(__file__)
-  ds = CTestLoader(os.path.join(folder, 'test'))
-  print(len(ds))
-  batch, (y,) = ds[0]
-  for k, v in batch.items():
-    print(k, v.shape)
-  print()
-  print(y.shape)
-  pass
+    return(res, (Y, ))

scripts/create-test-dataset-inpainting.py

@@ -0,0 +1,148 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-.
+import argparse, os, sys
+# add the root folder of the project to the path
+ROOT_FOLDER = os.path.abspath(os.path.dirname(__file__) + '/../')
+sys.path.append(ROOT_FOLDER)
+
+import numpy as np
+import Core.Utils as Utils
+from Core.CSamplesStorage import CSamplesStorage
+from Core.CDataSamplerInpainting import CDataSamplerInpainting
+from collections import defaultdict
+import glob
+import json
+import shutil
+import tensorflow as tf
+
+BATCH_SIZE = 128 * 4
+
+def samplesStream(params, take, filename, stats):
+  if not isinstance(take, list): take = [take]
+  # filename is "{placeId}/{userId}/{screenId}/train.npz"
+  # extract the placeId, userId, and screenId
+  parts = os.path.split(filename)[0].split(os.path.sep)
+  placeId, userId, screenId = parts[-3], parts[-2], parts[-1]
+  # use the stats to get the numeric values of the placeId, userId, and screenId  
+  ds = CDataSamplerInpainting(
+    CSamplesStorage(
+      placeId=stats['placeId'].index(placeId),
+      userId=stats['userId'].index(userId),
+      screenId=stats['screenId'].index('%s/%s' % (placeId, screenId))
+    ),
+    defaults=params, 
+    batch_size=BATCH_SIZE, minFrames=params['timesteps'],
+    keys=take
+  )
+  ds.addBlock(Utils.datasetFrom(filename))
+  
+  N = ds.totalSamples
+  for i in range(0, N, BATCH_SIZE):
+    indices = list(range(i, min(i + BATCH_SIZE, N)))
+    batch, rejected, accepted = ds.sampleByIds(indices)
+    if batch is None: continue
+
+    # main batch
+    x, y = batch
+    for idx in range(len(x['points'])):
+      resX = {}
+      for k, v in x.items():
+        item = v[idx, None]
+        if tf.is_tensor(item): item = item.numpy()
+        resX[f'X_{k}'] = item
+        continue
+
+      resY = {}
+      for k, v in y.items():
+        item = v[idx, None]
+        if tf.is_tensor(item): item = item.numpy()
+        resY[f'Y_{k}'] = item
+        continue
+        
+      yield dict(**resX, **resY)
+      continue
+    continue
+  return
+
+def batches(*params):
+  data = defaultdict(list)
+  for sample in samplesStream(*params):
+    for k, v in sample.items():
+      data[k].append(v)
+      continue
+
+    if BATCH_SIZE <= len(data['X_points']):
+      yield data
+      data = defaultdict(list)
+    continue
+
+  if 0 < len(data['X_points']):
+    # copy data to match batch size
+    for k, v in data.items():
+      while len(v) < BATCH_SIZE: v.extend(v)
+      data[k] = v[:BATCH_SIZE]
+      continue
+    yield data
+  return
+############################################
+def generateTestDataset(params, filename, stats, outputFolder):
+  # generate test dataset
+  ONE_MB = 1024 * 1024
+  totalSize = 0
+  if not os.path.exists(outputFolder):
+    os.makedirs(outputFolder, exist_ok=True)
+  for bIndex, batch in enumerate(batches(params, ['clean'], filename, stats)):
+    fname = os.path.join(outputFolder, 'test-%d.npz' % bIndex)
+    # concatenate all arrays
+    batch = {k: np.concatenate(v, axis=0) for k, v in batch.items()}
+    np.savez_compressed(fname, **batch)
+    # get fname size
+    size = os.path.getsize(fname)
+    totalSize += size
+    print('%d | Size: %.1f MB | Total: %.1f MB' % (bIndex + 1, size / ONE_MB, totalSize / ONE_MB))
+    continue
+  print('Done')
+  return
+
+def main(args):
+  PARAMS = [
+    dict(      
+      timesteps=args.steps,
+      stepsSampling='uniform',
+      # no augmentations by default
+      pointsNoise=0.01, pointsDropout=0.0,
+      eyesDropout=0.1, eyesAdditiveNoise=0.01, brightnessFactor=1.5, lightBlobFactor=1.5,
+      targets=dict(keypoints=3, total=10),
+    ),
+  ]
+  folder = os.path.join(ROOT_FOLDER, 'Data', 'remote')
+
+  stats = None
+  with open(os.path.join(folder, 'stats.json'), 'r') as f:
+    stats = json.load(f)
+
+  # remove all content from the output folder
+  shutil.rmtree(args.output, ignore_errors=True)
+  # recursively find the train file
+  trainFilename = glob.glob(os.path.join(folder, '**', 'test.npz'), recursive=True)
+  print('Found test files:', len(trainFilename))
+  for idx, filename in enumerate(trainFilename):
+    print('Processing', filename)
+    for params in PARAMS:
+      targetFolder = os.path.join(args.output, 'test-%d' % idx)
+      generateTestDataset(params, filename, stats, outputFolder=targetFolder)
+      continue
+  return
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--steps', type=int, default=5, help='Number of timesteps')
+  parser.add_argument('--batch-size', type=int, default=512, help='Batch size of the test dataset')
+  parser.add_argument(
+    '--output', type=str, help='Output folder',
+    default=os.path.join(ROOT_FOLDER, 'Data', 'test-inpainting')
+  )
+  args = parser.parse_args()
+  BATCH_SIZE = args.batch_size # TODO: fix this hack
+  main(args)
+  pass