3 files changed, 16 insertions, 3 deletions
diff --git a/app/client/modules/pix2pixhd/pix2pixhd.actions.js b/app/client/modules/pix2pixhd/pix2pixhd.actions.js
index 811f913..a17eeab 100644
--- a/app/client/modules/pix2pixhd/pix2pixhd.actions.js
+++ b/app/client/modules/pix2pixhd/pix2pixhd.actions.js
@@ -202,7 +202,7 @@ export const list_epochs = (checkpoint_name) => (dispatch) => {
   })
 }
 
-export const count_sequence_dataset = (checkpoint_name) => (dispatch) => {
+export const count_dataset = (checkpoint_name) => (dispatch) => {
   const module = pix2pixhdModule.name
   util.allProgress([
     actions.socket.count_directory({ module, dir: 'sequences/' + checkpoint_name + '/' }),
@@ -218,8 +218,9 @@ export const count_sequence_dataset = (checkpoint_name) => (dispatch) => {
     const [sequenceCount, datasetCount] = res //, datasets, results, output, datasetUsage, lossReport] = res
     console.log(sequenceCount, datasetCount)
     dispatch({
-      type: types.pix2pixhd.load_results,
-      results: {
+      type: types.pix2pixhd.load_dataset_count,
+      data: {
+        name: checkpoint_name,
         sequenceCount,
         datasetCount,
       }
diff --git a/app/client/modules/pix2pixhd/pix2pixhd.reducer.js b/app/client/modules/pix2pixhd/pix2pixhd.reducer.js
index c3d52a3..5a2afc0 100644
--- a/app/client/modules/pix2pixhd/pix2pixhd.reducer.js
+++ b/app/client/modules/pix2pixhd/pix2pixhd.reducer.js
@@ -8,6 +8,11 @@ const pix2pixhdInitialState = {
   folder_id: 0,
   data: null,
   results: null,
+  checkpoint: {
+    name: '',
+    sequenceCount: 0,
+    datasetCount: 0,
+  }
 }
 
 const pix2pixhdReducer = (state = pix2pixhdInitialState, action) => {
@@ -21,6 +26,11 @@ const pix2pixhdReducer = (state = pix2pixhdInitialState, action) => {
         ...state,
         results: action.results,
       }
+    case types.pix2pixhd.load_dataset_count:
+      return {
+        ...state,
+        checkpoint: action.data,
+      }
     case types.file.destroy:
       console.log('file destroy', state.results)
       return {
diff --git a/app/client/modules/pix2pixhd/views/pix2pixhd.train.js b/app/client/modules/pix2pixhd/views/pix2pixhd.train.js
index 05ad638..06caa5a 100644
--- a/app/client/modules/pix2pixhd/views/pix2pixhd.train.js
+++ b/app/client/modules/pix2pixhd/views/pix2pixhd.train.js
@@ -52,6 +52,7 @@ class Pix2PixHDTrain extends Component {
     if (prevState.checkpoint_name !== this.state.checkpoint_name) {
       this.setState({ epoch: 'latest' })
       this.props.actions.list_epochs(this.state.checkpoint_name)
+      this.props.actions.count_dataset(this.state.checkpoint_name)
     }
   }
   handleChange(name, value){
@@ -152,6 +153,7 @@ class Pix2PixHDTrain extends Component {
             </Group>
             <Group title='Augmentation Grid'>
               <AugmentationGrid
+                checkpoint={this.props.pix2pixhd.checkpoint}
                 take={[1,2,3,4,5,10,15,20,25,50,75,100,200,300,400,500,1000]}
                 make={[1,2,3,4,5,10,15,20,25,50,75,100,200,]}
                 onAugment={(augment_take, augment_make) => {