check if dataset is generated

3 files changed, 54 insertions, 49 deletions

diff --git a/app/client/modules/pix2pixhd/pix2pixhd.actions.js b/app/client/modules/pix2pixhd/pix2pixhd.actions.js
index 44a42c4..d067017 100644
--- a/app/client/modules/pix2pixhd/pix2pixhd.actions.js
+++ b/app/client/modules/pix2pixhd/pix2pixhd.actions.js
@@ -202,11 +202,11 @@ export const list_epochs = (checkpoint_name) => (dispatch) => {
   })
 }
 
-export const count_dataset = (checkpoint_name) => (dispatch) => {
+export const count_dataset = (dataset) => (dispatch) => {
   const module = pix2pixhdModule.name
   util.allProgress([
-    actions.socket.list_directory({ module, dir: 'sequences/' + checkpoint_name + '/' }),
-    actions.socket.count_directory({ module, dir: 'datasets/' + checkpoint_name + '/train_A/' }),
+    actions.socket.list_directory({ module, dir: 'sequences/' + dataset + '/' }),
+    actions.socket.count_directory({ module, dir: 'datasets/' + dataset + '/train_A/' }),
   ], (percent, i, n) => {
     console.log('pix2pixhd load progress', i, n)
     // dispatch({
@@ -221,7 +221,7 @@ export const count_dataset = (checkpoint_name) => (dispatch) => {
     dispatch({
       type: types.pix2pixhd.load_dataset_count,
       data: {
-        name: checkpoint_name,
+        name: dataset,
         sequence,
         sequenceCount,
         datasetCount,
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,
diff --git a/app/client/modules/pix2pixhd/views/sequence.editor.js b/app/client/modules/pix2pixhd/views/sequence.editor.js
index 33ed629..135dfb3 100644
--- a/app/client/modules/pix2pixhd/views/sequence.editor.js
+++ b/app/client/modules/pix2pixhd/views/sequence.editor.js
@@ -18,26 +18,6 @@ const initialState = {
   title: null,
 }
 
-/*
-  so there are two things you could do with this
-  1) create an entirely new dataset
-  2) add frames to an existing dataset
-  ...
-  method 1)
-  - requires minimal setup, just a server script which..
-     - creates the new sequence folder
-     - symlinks the frames
-     - runs build_dataset
-     - create appropriate file object, using parent dataset's folder_id
-     - enqueue initial training
-       ---> we can then train on basic subsets, with more control, like we do already...
-     - tell parent dataset thing the new sequence name
-  ...
-  method 2)
-  - requires sequence editor to be aware of its own dataset
-  - requires sequence editor to know whether a sequence is original or not
-*/
-
 class SequenceEditor extends Component {
   state = { ...initialState }
 
@@ -80,9 +60,10 @@ class SequenceEditor extends Component {
   }
 
   render() {
-    const { app, remote, checkpoint, folder_id } = this.props
+    const { app, pix2pixhd, remote, checkpoint, folder_id, generated } = this.props
     const { cursor, selection, title } = this.state
     const path = "sequences/" + checkpoint.name
+    console.log(checkpoint, pix2pixhd)
     return (
       <div className='sequenceEditor row'>
         {selection
@@ -106,6 +87,11 @@ class SequenceEditor extends Component {
                 Create
               </Button>
             </Group>
+            <Group title='Salt dataset'>
+              {generated
+                ? 'Salting is only available on generated datasets.'
+                : 'Salting coming soon!'}
+            </Group>
           </div>
         : <div className='form'><Group title='New dataset'>Please select some frames</Group></div>
         }
@@ -145,6 +131,7 @@ function Frame ({ label, path, frame }) {
 
 const mapStateToProps = state => ({
   app: state.system.app,
+  pix2pixhd: state.module.pix2pixhd,
 })
 
 const mapDispatchToProps = (dispatch, ownProps) => ({