updates

Author: Abdul Dakkak

vision/cifar10.go

@@ -2,8 +2,8 @@ package vision
 
 import (
 	"bufio"
-	"bytes"
 	"encoding/binary"
+	"image"
 	"io"
 	"os"
 	"path"
@@ -16,7 +16,7 @@ import (
 	"github.com/rai-project/config"
 	"github.com/rai-project/dldataset"
 	"github.com/rai-project/downloadmanager"
-	"github.com/rai-project/image"
+	"github.com/rai-project/image/types"
 	context "golang.org/x/net/context"
 )
 
@@ -41,15 +41,15 @@ type CIFAR10 struct {
 
 type CIFAR10LabeledImage struct {
 	label string
-	data  image.RGBImage
+	data  *types.RGBImage
 }
 
 func (l CIFAR10LabeledImage) Label() string {
 	return l.label
 }
 
-func (l CIFAR10LabeledImage) Data() (io.Reader, error) {
-	return bytes.NewBuffer(l.data), nil
+func (l CIFAR10LabeledImage) Data() (interface{}, error) {
+	return l.data, nil
 }
 
 func (*CIFAR10) Name() string {
@@ -235,26 +235,17 @@ func (d *CIFAR10) readEntry(ctx context.Context, reader io.Reader) (*CIFAR10Labe
 
 	pixelByteSize := int64(d.pixelByteSize)
 	pixelBytesReader := io.LimitReader(reader, pixelByteSize)
-	pixelBytes := make([]byte, pixelByteSize)
-	err = binary.Read(pixelBytesReader, binary.LittleEndian, &pixelBytes)
+
+	img := types.NewRGBImage(image.Rect(0, 0, d.imageDimensions[0], d.imageDimensions[1]))
+
+	err = binary.Read(pixelBytesReader, binary.LittleEndian, img.Pix)
 	if err == io.EOF {
 		return nil, err
 	}
 	if err != nil {
 		return nil, errors.New("unable to read label")
 	}
 
-	img := image.NewRGBImage(image.Rect(0, 0, d.imageDimensions[0], d.imageDimensions[1]))
-	var idx int
-	for y := 0; y < d.imageDimensions[0]; y++ {
-		for x := 0; x < d.imageDimensions[1]; x++ {
-			for c := 0; c < d.imageDimensions[2]; c++ {
-				img.Pix[ii] = float32(pixelBytes[idx])
-				idx++
-			}
-		}
-	}
-
 	return &CIFAR10LabeledImage{
 		label: d.labels[labelIdx],
 		data:  img,

vision/cifar100.go

@@ -2,8 +2,8 @@ package vision
 
 import (
 	"bufio"
-	"bytes"
 	"encoding/binary"
+	"image"
 	"io"
 	"os"
 	"path"
@@ -16,6 +16,7 @@ import (
 	"github.com/rai-project/config"
 	"github.com/rai-project/dldataset"
 	"github.com/rai-project/downloadmanager"
+	"github.com/rai-project/image/types"
 	context "golang.org/x/net/context"
 )
 
@@ -44,7 +45,7 @@ type CIFAR100 struct {
 type CIFAR100LabeledImage struct {
 	coarseLabel string
 	fineLabel   string
-	data        image.RGBImage
+	data        *types.RGBImage
 }
 
 func (l CIFAR100LabeledImage) CoarseLabel() string {
@@ -59,8 +60,8 @@ func (l CIFAR100LabeledImage) Label() string {
 	return l.FineLabel()
 }
 
-func (l CIFAR100LabeledImage) Data() (io.Reader, error) {
-	return bytes.NewBuffer(l.data), nil
+func (l CIFAR100LabeledImage) Data() (interface{}, error) {
+	return l.data, nil
 }
 
 func (*CIFAR100) Name() string {
@@ -270,26 +271,17 @@ func (d *CIFAR100) readEntry(ctx context.Context, reader io.Reader) (*CIFAR100La
 
 	pixelByteSize := int64(d.pixelByteSize)
 	pixelBytesReader := io.LimitReader(reader, pixelByteSize)
-	pixelBytes := make([]byte, pixelByteSize)
-	err = binary.Read(pixelBytesReader, binary.LittleEndian, &pixelBytes)
+
+	img := types.NewRGBImage(image.Rect(0, 0, d.imageDimensions[0], d.imageDimensions[1]))
+
+	err = binary.Read(pixelBytesReader, binary.LittleEndian, img.Pix)
 	if err == io.EOF {
 		return nil, err
 	}
 	if err != nil {
 		return nil, errors.Wrap(err, "unable to read label")
 	}
 
-	img := image.NewRGBImage(image.Rect(0, 0, d.imageDimensions[0], d.imageDimensions[1]))
-	var idx int
-	for y := 0; y < d.imageDimensions[0]; y++ {
-		for x := 0; x < d.imageDimensions[1]; x++ {
-			for c := 0; c < d.imageDimensions[2]; c++ {
-				img.Pix[ii] = float32(pixelBytes[idx])
-				idx++
-			}
-		}
-	}
-
 	return &CIFAR100LabeledImage{
 		coarseLabel: d.coarseLabels[coarseLabelIdx],
 		fineLabel:   d.fineLabels[fineLabelIdx],

vision/mnist.go

@@ -1,8 +1,7 @@
 package vision
 
 import (
-	"bytes"
-	"io"
+	"image"
 	"path"
 	"strconv"
 	"strings"
@@ -12,6 +11,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/rai-project/config"
 	"github.com/rai-project/dldataset"
+	"github.com/rai-project/image/types"
 	mnistLoader "github.com/unixpickle/mnist"
 )
 
@@ -25,15 +25,15 @@ var mnist *MNIST
 
 type MNISTLabeledImage struct {
 	label string
-	data  []byte
+	data  *types.RGBImage
 }
 
 func (l MNISTLabeledImage) Label() string {
 	return l.label
 }
 
-func (l MNISTLabeledImage) Data() (io.Reader, error) {
-	return bytes.NewBuffer(l.data), nil
+func (l MNISTLabeledImage) Data() (interface{}, error) {
+	return l.data, nil
 }
 
 func (*MNIST) Name() string {
@@ -86,17 +86,24 @@ func (d *MNIST) Get(ctx context.Context, name string) (dldataset.LabeledData, er
 	}
 
 	elem := dataset.Samples[idx]
-	data := make([]byte, len(elem.Intensities))
+
+	img := types.NewRGBImage(image.Rect(0, 0, dataset.Width, dataset.Height))
+	data := img.Pix
+
 	for ii, intensity := range elem.Intensities {
 		if intensity == 1 {
-			data[ii] = byte(1)
+			data[3*ii+0] = byte(1)
+			data[3*ii+1] = byte(1)
+			data[3*ii+2] = byte(1)
 		} else {
-			data[ii] = byte(0)
+			data[3*ii+0] = byte(0)
+			data[3*ii+1] = byte(0)
+			data[3*ii+2] = byte(0)
 		}
 	}
 
 	return &MNISTLabeledImage{
-		data:  data,
+		data:  img,
 		label: strconv.Itoa(elem.Label),
 	}, nil
 }