dataset.go

package dldataset

import (
  "fmt"
  "github.com/c3sr/dldataset/generateDataset"
  "github.com/c3sr/go-python3"
  "github.com/pkg/errors"
  "gopkg.in/yaml.v2"
  "io"
  "io/ioutil"
  "net/http"
  "os"
  "path/filepath"
  "runtime"
)

type Dataset struct {
  dataInMemory map[int]interface{}
  method       map[string]string
  files        []string
  count        int
}

func NewDataset(datasetPath string, count int) (*Dataset, error) {
  datasetPath, err := filepath.Abs(datasetPath)
  if err != nil {
    return nil, err
  }

  datasetManifest, err := ioutil.ReadFile(datasetPath)
  if err != nil {
    return nil, err
  }

  dataset := new(Dataset)
  if err := yaml.Unmarshal(datasetManifest, &(dataset.method)); err != nil {
    return nil, err
  }

  runtime.LockOSThread()
  defer runtime.UnlockOSThread()

  pyState := python3.PyGILState_Ensure()
  defer python3.PyGILState_Release(pyState)

  python3.PyRun_SimpleString(dataset.method["init"])
  pyMain := python3.PyImport_AddModule("__main__")
  pyDict := python3.PyModule_GetDict(pyMain)
  pyInit := python3.PyDict_GetItemString(pyDict, "init")

  pyCount := python3.PyLong_FromLong(count)
  defer pyCount.DecRef()

  pyRes := pyInit.CallFunctionObjArgs(pyCount)
  defer pyRes.DecRef()

  dataset.count = python3.PyLong_AsLong(pyRes)

  if dataset.count < 0 {
    return nil, fmt.Errorf("Error when initializing dataset.")
  }

  return dataset, nil
}

func NewImageDatasetFromUrls(urls []string, overrideDataDir string) (*Dataset, error) {
  dataset := new(Dataset)

  images, err := downloadImagesFromUrls(urls, overrideDataDir)
  dataset.files = images
  if err != nil {
    return nil, err
  }

  datasetManifest := generateDataset.Create(images)

  if err := yaml.Unmarshal([]byte(datasetManifest), &(dataset.method)); err != nil {
    return nil, err
  }

  runtime.LockOSThread()
  defer runtime.UnlockOSThread()

  pyState := python3.PyGILState_Ensure()
  defer python3.PyGILState_Release(pyState)

  python3.PyRun_SimpleString(dataset.method["init"])
  pyMain := python3.PyImport_AddModule("__main__")
  pyDict := python3.PyModule_GetDict(pyMain)
  pyInit := python3.PyDict_GetItemString(pyDict, "init")

  pyCount := python3.PyLong_FromLong(len(urls))
  defer pyCount.DecRef()

  pyRes := pyInit.CallFunctionObjArgs(pyCount)
  defer pyRes.DecRef()

  dataset.count = python3.PyLong_AsLong(pyRes)

  if dataset.count < 0 {
    return nil, fmt.Errorf("Error when initializing dataset.")
  }

  return dataset, nil
}

func downloadImagesFromUrls(urls []string, overrideDataDir string) (images []string, err error) {
  if overrideDataDir == "" {
    overrideDataDir = os.ExpandEnv("$HOME/data/mlharness/image_url/")
  }
  images = make([]string, len(urls))
  prefix := generateDataset.GetUniquePrefix()

  for i, url := range urls {
    filename := fmt.Sprintf("%s%d%s", prefix, i, filepath.Ext(url))
    filepath := filepath.Join(overrideDataDir, filename)
    err = downloadFile(filepath, url)
    if err != nil {
      break
    }
    images[i] = filepath
  }
  return
}

func downloadFile(pathToFile string, url string) (err error) {
  resp, err := http.Get(url)
  if err != nil {
    return
  }

  defer resp.Body.Close()

  err = ensureBaseDir(pathToFile)
  if err != nil {
    return
  }

  out, err := os.Create(pathToFile)
  if err != nil {
    return
  }

  defer out.Close()
  _, err = io.Copy(out, resp.Body)
  return
}

func ensureBaseDir(path string) error {
  baseDir := filepath.Dir(path)
  info, err := os.Stat(baseDir)
  if err == nil && info.IsDir() {
    return nil
  }
  return os.MkdirAll(baseDir, 0755)
}

func (d *Dataset) CleanupFiles() (err error) {
  for _, file := range d.files {
    if err = os.Remove(file); err != nil {
      return errors.Wrapf(err, "failed to remove dataset file: %v", file)
    }
  }

  return
}

func (d *Dataset) Count() int {
  return d.count
}

func (d *Dataset) Load(input []int) error {
  runtime.LockOSThread()
  defer runtime.UnlockOSThread()

  pyState := python3.PyGILState_Ensure()
  defer python3.PyGILState_Release(pyState)

  python3.PyRun_SimpleString(d.method["load"])
  pyMain := python3.PyImport_AddModule("__main__")
  pyDict := python3.PyModule_GetDict(pyMain)
  pyLoad := python3.PyDict_GetItemString(pyDict, "load")

  pyList := python3.PyList_New(0)
  for _, in := range input {
    python3.PyList_Append(pyList, python3.PyLong_FromLong(in))
  }
  defer pyList.DecRef()

  pyRes := pyLoad.CallFunctionObjArgs(pyList)
  pyRes.DecRef()

  return nil
}

func (d *Dataset) Unload(input []int) error {
  runtime.LockOSThread()
  defer runtime.UnlockOSThread()

  pyState := python3.PyGILState_Ensure()
  defer python3.PyGILState_Release(pyState)

  python3.PyRun_SimpleString(d.method["unload"])
  pyMain := python3.PyImport_AddModule("__main__")
  pyDict := python3.PyModule_GetDict(pyMain)
  pyUnload := python3.PyDict_GetItemString(pyDict, "unload")

  pyList := python3.PyList_New(0)
  for _, in := range input {
    python3.PyList_Append(pyList, python3.PyLong_FromLong(in))
  }
  defer pyList.DecRef()

  pyRes := pyUnload.CallFunctionObjArgs(pyList)
  pyRes.DecRef()

  if d.dataInMemory == nil {
    return nil
  }
  if len(input) == 0 {
    d.dataInMemory = nil
  } else {
    for _, sample := range input {
      delete(d.dataInMemory, sample)
    }
  }

  return nil
}

func (d *Dataset) Set(data map[int]interface{}) error {
  d.dataInMemory = data
  return nil
}

func (d *Dataset) Get(input []int) (map[int]interface{}, error) {
  res := make(map[int]interface{})
  for _, sample := range input {
    if _, exist := d.dataInMemory[sample]; !exist {
      return nil, fmt.Errorf("sample id %d not loaded.", sample)
    }
    res[sample] = d.dataInMemory[sample]
  }
  return res, nil
}

func (d *Dataset) GetAll() (map[int]interface{}, error) {
  if d.dataInMemory == nil {
    return nil, fmt.Errorf("No data in memory.")
  }
  return d.dataInMemory, nil
}