1 files changed, 42 insertions, 24 deletions

diff --git a/app/client/modules/pix2pixhd/views/pix2pixhd.train.js b/app/client/modules/pix2pixhd/views/pix2pixhd.train.js
index c5e8e67..de32fcd 100644
--- a/app/client/modules/pix2pixhd/views/pix2pixhd.train.js
+++ b/app/client/modules/pix2pixhd/views/pix2pixhd.train.js
@@ -24,16 +24,17 @@ import SequenceEditor from './sequence.editor'
 import pix2pixhdModule from '../pix2pixhd.module'
 
 class Pix2PixHDTrain extends Component {
+  state = {
+    dataset: 'PLACEHOLDER',
+    epoch: 'latest',
+    augment_name: '',
+    augment_take: 100,
+    augment_make: 20,
+    generated: false,
+  }
   constructor(props){
     super(props)
     this.handleChange = this.handleChange.bind(this)
-    this.state = {
-      checkpoint_name: 'PLACEHOLDER',
-      epoch: 'latest',
-      augment_name: '',
-      augment_take: 100,
-      augment_make: 20,
-    }
   }
   componentDidMount(){
     const id = this.props.match.params.id || localStorage.getItem('pix2pixhd.last_id')
@@ -50,15 +51,30 @@ class Pix2PixHDTrain extends Component {
       this.props.history.push('/pix2pixhd/new/')
     }
     if (dataset) {
-      this.setState({ checkpoint_name: dataset })
+      this.setState({ dataset })
     }
   }
   componentDidUpdate(prevProps, prevState){
-    if (prevState.checkpoint_name !== this.state.checkpoint_name) {
-      localStorage.setItem('pix2pixhd.last_dataset', 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)
+    if ((!prevProps.pix2pixhd.data && this.props.pix2pixhd.data)
+      || (prevProps.pix2pixhd.data && prevState.dataset !== this.state.dataset)) {
+      const dataset = this.props.pix2pixhd.data.datasetLookup[this.state.dataset]
+      if (dataset) {
+        const generated = dataset.input
+          .map(f => this.props.pix2pixhd.data.fileLookup[f])
+          .reduce((a,b) => {
+            return b.generated || a
+          }, false)
+        dataset.generated = generated
+        this.setState({ generated })
+      }
+    }
+    if (prevState.dataset !== this.state.dataset) {
+      localStorage.setItem('pix2pixhd.last_dataset', this.state.dataset)
+      this.setState({
+        epoch: 'latest',
+      })
+      this.props.actions.list_epochs(this.state.dataset)
+      this.props.actions.count_dataset(this.state.dataset)
     }
   }
   handleChange(name, value){
@@ -78,7 +94,7 @@ class Pix2PixHDTrain extends Component {
     const { checkpoint } = pix2pixhd
     // console.log(pix2pixhd)
 
-    const checkpointGroups = Object.keys(folderLookup).map(id => {
+    const sequenceGroups = Object.keys(folderLookup).map(id => {
       const folder = this.props.pix2pixhd.data.folderLookup[id]
       if (folder.name === 'results') return
       const datasets = folder.datasets.map(name => {
@@ -95,7 +111,7 @@ class Pix2PixHDTrain extends Component {
     }).filter(n => !!n && !!n.options.length).sort((a,b) => a.name.localeCompare(b.name))
 
     // console.log('state', this.props.pix2pixhd.data.epochs)
-    // console.log(this.state.checkpoint_name, this.state.epoch)
+    // console.log(this.state.dataset, this.state.epoch)
     // console.log(queue)
     return (
       <div className='app pix2pixhd'>
@@ -106,12 +122,12 @@ class Pix2PixHDTrain extends Component {
           <div className='column'>
             <Group title='Sequence'>
               <SelectGroup
-                name='checkpoint_name'
+                name='dataset'
                 title='Sequence name'
-                options={checkpointGroups}
+                options={sequenceGroups}
                 onChange={this.handleChange}
                 placeholder='Pick a dataset'
-                value={this.state.checkpoint_name}
+                value={this.state.dataset}
               />
               <Select
                 title="Epoch"
@@ -131,8 +147,9 @@ class Pix2PixHDTrain extends Component {
   }
 
   renderEditor(){
-    const { pix2pixhd, queue } = this.props
+    const { pix2pixhd, queue, remote } = this.props
     const { checkpoint, folder_id } = pix2pixhd
+    const { dataset, generated } = this.state
     return (
       <div>
         <Group title='Sequence Editor'>
@@ -140,6 +157,7 @@ class Pix2PixHDTrain extends Component {
             folder_id={folder_id}
             module={pix2pixhdModule}
             checkpoint={checkpoint}
+            generated={generated}
           />
         </Group>
         <div className='columns'>
@@ -150,16 +168,16 @@ class Pix2PixHDTrain extends Component {
                 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) => {
-                  this.props.remote.augment_task(this.state.checkpoint_name, {
+                  remote.augment_task(dataset, {
                     ...this.state,
                     augment_take,
                     augment_make,
                   })
                 }}
                 onTrain={() => {
-                  this.props.remote.train_task(this.state.checkpoint_name, folder_id, 1)
+                  remote.train_task(dataset, folder_id, 1)
                   setTimeout(() => { // auto-generate epoch demo
-                    this.props.remote.augment_task(this.state.checkpoint_name, {
+                    remote.augment_task(dataset, {
                       ...this.state,
                       augment_take: 10,
                       augment_make: 149,
@@ -193,13 +211,13 @@ class Pix2PixHDTrain extends Component {
               <Button
                 title="Augment dataset"
                 value="Augment"
-                onClick={() => this.props.remote.augment_task(this.state.checkpoint_name, this.state)}
+                onClick={() => this.props.remote.augment_task(dataset, this.state)}
               />
               <Button
                 title="Make a movie without augmenting"
                 value="Generate"
                 onClick={() => {
-                  this.props.remote.augment_task(this.state.checkpoint_name, {
+                  this.props.remote.augment_task(dataset, {
                     ...this.state,
                     no_symlinks: true,
                     mov: true,