path: root/app/client/modules/pix2pixhd/views/pix2pixhd.train.js

import { h, Component } from 'preact'
import { bindActionCreators } from 'redux'
import { connect } from 'react-redux'
import util from '../../../util'

import * as pix2pixhdActions from '../pix2pixhd.actions'
import * as pix2pixhdTasks from '../pix2pixhd.tasks'

import {
  Loading,
  FileList, FileRow,
  Select, SelectGroup, Group, Param, Button,
  TextInput, NumberInput,
  CurrentTask, TaskList,
  AugmentationGrid,
} from '../../../common'
import DatasetForm from '../../../dataset/dataset.form'
import NewDatasetForm from '../../../dataset/dataset.new'
import UploadStatus from '../../../dataset/upload.status'

import DatasetComponent from '../../../dataset/dataset.component'
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,
    processed: false,
  }
  constructor(props){
    super(props)
    this.handleChange = this.handleChange.bind(this)
  }
  componentDidMount(){
    const id = this.props.match.params.id || localStorage.getItem('pix2pixhd.last_id')
    const dataset = localStorage.getItem('pix2pixhd.last_dataset')
    console.log('load dataset:', id, dataset)
    const { match, pix2pixhd, actions } = this.props
    if (id === 'new') return
    if (id) {
      if (parseInt(id)) localStorage.setItem('pix2pixhd.last_id', id)
      if (! pix2pixhd.folder || pix2pixhd.folder.id !== id) {
        actions.load_directories(id)
      }
    } else {
      this.props.history.push('/pix2pixhd/new/')
    }
    if (dataset) {
      this.setState({ dataset })
    }
  }
  componentDidUpdate(prevProps, prevState){
    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 processed = dataset.input
          .map(f => this.props.pix2pixhd.data.fileLookup[f])
          .reduce((a, b) => {
            console.log(b)
            return b.processed || a
          }, false)
        console.log(dataset.input, processed)
        dataset.processed = processed
        this.setState({ processed })
      }
    }
    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){
    console.log('name', name, 'value', value)
    this.setState({ [name]: value })
  }
  interrupt(){
    this.props.actions.queue.stop_task('gpu')
  }
  render(){
    if (this.props.pix2pixhd.loading) {
      return <Loading progress={this.props.pix2pixhd.progress} />
    }
    const { pix2pixhd, match, history } = this.props
    const { folderLookup, datasetLookup } = (pix2pixhd.data || {})
    const folder = (folderLookup || {})[pix2pixhd.folder_id] || {}
    const { checkpoint } = pix2pixhd
    // console.log(pix2pixhd)

    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 => {
        const dataset = datasetLookup[name]
        // console.log(folder)
        if (dataset.isBuilt) { // && dataset.checkpoints.length
          return name
        }
        return null
      }).filter(n => !!n)
      return {
        name: folder.name,
        options: datasets.sort(),
      }
    }).filter(n => !!n && !!n.options.length).sort((a,b) => a.name.localeCompare(b.name))

    let epochs = this.props.pix2pixhd.epochs
    if (!epochs || !epochs.length) epochs = ["latest"]
    // console.log('state', this.props.pix2pixhd.data.epochs)
    // console.log(this.state.dataset, this.state.epoch)
    // console.log(queue)
    return (
      <div className='app pix2pixhd'>
        <div className='heading'>
          <h1>pix2pixhd training</h1>
        </div>
        <div className='columns'>
          <div className='column'>
            <Group title='Sequence'>
              <SelectGroup
                name='dataset'
                title='Sequence name'
                options={sequenceGroups}
                onChange={this.handleChange}
                placeholder='Pick a dataset'
                value={this.state.dataset}
              />
              <Select
                title="Epoch"
                name="epoch"
                options={epochs}
                onChange={this.handleChange}
                value={this.state.epoch}
              />
            </Group>
          </div>
        </div>
        {checkpoint && checkpoint.sequence && checkpoint.sequence.length
          ? this.renderEditor()
          : checkpoint && <div>Sequence empty, augmentation impossible</div>}
      </div>
    )
  }

  renderEditor(){
    const { pix2pixhd, queue, remote } = this.props
    const { dataset, processed } = this.state
    let { checkpoint, folder_id } = pix2pixhd

    const currentDataset = this.props.pix2pixhd.data.datasetLookup[this.state.dataset]
    if (currentDataset && currentDataset.folder_id) {
      folder_id = currentDataset.folder_id || folder_id
      console.log("dataset folder id", folder_id)
      // const processed = currentDataset.input
      //   .map(f => this.props.pix2pixhd.data.fileLookup[f])
      //   .reduce((a, b) => {
      //     console.log(b)
      //     return b.processed || a
      //   }, false)
      // console.log("is processed?", processed)
    }
    console.log(currentDataset)
    console.log("so folder id is..", folder_id)

    return (
      <div>
        <Group title='Sequence Editor'>
          <SequenceEditor
            folder_id={folder_id}
            module={pix2pixhdModule}
            checkpoint={checkpoint}
            processed={processed}
          />
        </Group>
        <div className='columns'>
          <div className='column'>
            <Group title='Augmentation Grid'>
              <AugmentationGrid
                checkpoint={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) => {
                  remote.augment_task(dataset, {
                    ...this.state,
                    augment_take,
                    augment_make,
                  })
                }}
                onTrain={() => {
                  remote.train_task(dataset, folder_id, 1)
                  setTimeout(() => { // auto-generate epoch demo
                    remote.augment_task(dataset, {
                      ...this.state,
                      augment_take: 10,
                      augment_make: 149,
                      no_symlinks: true,
                      mov: true,
                      folder_id: pix2pixhd.data.resultsFolder.id
                    })
                  }, 250)
                }}
              />
            </Group>
            <Group title='Augment'>
              <NumberInput
                name="augment_take"
                title="Pick N random frames"
                value={this.state.augment_take}
                onChange={this.handleChange}
                type="int"
                min="1"
                max="1000"
              />
              <NumberInput
                name="augment_make"
                title="Generate N recursively"
                value={this.state.augment_make}
                onChange={this.handleChange}
                type="int"
                min="1"
                max="1000"
              />
              <Button
                title="Augment dataset"
                value="Augment"
                onClick={() => this.props.remote.augment_task(dataset, this.state)}
              />
              <Button
                title="Make a movie without augmenting"
                value="Generate"
                onClick={() => {
                  this.props.remote.augment_task(dataset, {
                    ...this.state,
                    no_symlinks: true,
                    mov: true,
                    folder_id: pix2pixhd.data.resultsFolder.id
                  })
                }}
              />
            </Group>

            <Group title='Status'>
              <Button
                title="GPU"
                value={this.props.runner.gpu.status === 'IDLE' ? "Idle" : "Interrupt"}
                onClick={() => this.interrupt()}
              />
              <CurrentTask />
            </Group>
          </div>
          <div className='column'>
            <Group title='Upcoming Tasks'>
              <TaskList tasks={queue.queue.map(id => queue.tasks[id])} sort="date asc" />
            </Group>
          </div>
        </div>
      </div>
    )
  }
}

const mapStateToProps = state => ({
  pix2pixhd: state.module.pix2pixhd,
  runner: state.system.runner,
  queue: state.queue,
})

const mapDispatchToProps = (dispatch, ownProps) => ({
  actions: bindActionCreators(pix2pixhdActions, dispatch),
  remote: bindActionCreators(pix2pixhdTasks, dispatch),
})

export default connect(mapStateToProps, mapDispatchToProps)(Pix2PixHDTrain)