github.com/janelia-flyem/dvid@v1.0.0/datatype/labelmap/read.go (about)

     1  package labelmap
     2  
     3  import (
     4  	"bytes"
     5  	"compress/gzip"
     6  	"encoding/binary"
     7  	"fmt"
     8  	"io"
     9  	"io/ioutil"
    10  	"log"
    11  	"sync"
    12  
    13  	"github.com/janelia-flyem/dvid/datastore"
    14  	"github.com/janelia-flyem/dvid/datatype/common/labels"
    15  	"github.com/janelia-flyem/dvid/datatype/imageblk"
    16  	"github.com/janelia-flyem/dvid/dvid"
    17  	"github.com/janelia-flyem/dvid/server"
    18  	"github.com/janelia-flyem/dvid/storage"
    19  )
    20  
    21  // low-level reading of blocks from io.Reader for use with POSTing of blocks and other ingestion.
    22  func readStreamedBlock(r io.Reader, scale uint8) (block *labels.Block, compressed []byte, bx, by, bz int32, err error) {
    23  	hdrBytes := make([]byte, 16)
    24  	var n int
    25  	n, err = io.ReadFull(r, hdrBytes)
    26  	if n == 0 || err != nil {
    27  		return
    28  	}
    29  	if n != 16 {
    30  		err = fmt.Errorf("error reading header bytes, only got %d bytes not 16", n)
    31  		return
    32  	}
    33  	bx = int32(binary.LittleEndian.Uint32(hdrBytes[0:4]))
    34  	by = int32(binary.LittleEndian.Uint32(hdrBytes[4:8]))
    35  	bz = int32(binary.LittleEndian.Uint32(hdrBytes[8:12]))
    36  	numBytes := int(binary.LittleEndian.Uint32(hdrBytes[12:16]))
    37  	if numBytes == 0 {
    38  		err = fmt.Errorf("illegal zero length block at block (%d, %d, %d) detected", bx, by, bz)
    39  		return
    40  	}
    41  	bcoord := dvid.ChunkPoint3d{bx, by, bz}.ToIZYXString()
    42  	compressed = make([]byte, numBytes)
    43  	n, err = io.ReadFull(r, compressed)
    44  	if n != numBytes || err != nil {
    45  		err = fmt.Errorf("error reading %d bytes for block %s: %d actually read (%v)", numBytes, bcoord, n, err)
    46  		return
    47  	}
    48  
    49  	gzipIn := bytes.NewBuffer(compressed)
    50  	var zr *gzip.Reader
    51  	zr, err = gzip.NewReader(gzipIn)
    52  	if err != nil {
    53  		err = fmt.Errorf("can't initiate gzip reader on compressed data of length %d: %v", len(compressed), err)
    54  		return
    55  	}
    56  	var uncompressed []byte
    57  	uncompressed, err = ioutil.ReadAll(zr)
    58  	if err != nil {
    59  		err = fmt.Errorf("can't read all %d bytes from gzipped block %s: %v", numBytes, bcoord, err)
    60  		return
    61  	}
    62  	if err = zr.Close(); err != nil {
    63  		err = fmt.Errorf("error on closing gzip on block read: %v", err)
    64  		return
    65  	}
    66  
    67  	block = new(labels.Block)
    68  	err = block.UnmarshalBinary(uncompressed)
    69  	return
    70  }
    71  
    72  // ComputeTransform determines the block coordinate and beginning + ending voxel points
    73  // for the data corresponding to the given Block.
    74  func (v *Labels) ComputeTransform(tkey storage.TKey, blockSize dvid.Point) (blockBeg, dataBeg, dataEnd dvid.Point, err error) {
    75  	var ptIndex *dvid.IndexZYX
    76  	_, ptIndex, err = DecodeBlockTKey(tkey)
    77  	if err != nil {
    78  		return
    79  	}
    80  
    81  	// Get the bounding voxel coordinates for this block.
    82  	minBlockVoxel := ptIndex.MinPoint(blockSize)
    83  	maxBlockVoxel := ptIndex.MaxPoint(blockSize)
    84  
    85  	// Compute the boundary voxel coordinates for the ExtData and adjust
    86  	// to our block bounds.
    87  	minDataVoxel := v.StartPoint()
    88  	maxDataVoxel := v.EndPoint()
    89  	begVolCoord, _ := minDataVoxel.Max(minBlockVoxel)
    90  	endVolCoord, _ := maxDataVoxel.Min(maxBlockVoxel)
    91  
    92  	// Adjust the DVID volume voxel coordinates for the data so that (0,0,0)
    93  	// is where we expect this slice/subvolume's data to begin.
    94  	dataBeg = begVolCoord.Sub(v.StartPoint())
    95  	dataEnd = endVolCoord.Sub(v.StartPoint())
    96  
    97  	// Compute block coord matching dataBeg
    98  	blockBeg = begVolCoord.Sub(minBlockVoxel)
    99  
   100  	return
   101  }
   102  
   103  // GetImage retrieves a 2d image from a version node given a geometry of labels.
   104  func (d *Data) GetImage(v dvid.VersionID, vox *Labels, supervoxels bool, scale uint8, roiname dvid.InstanceName) (*dvid.Image, error) {
   105  	r, err := imageblk.GetROI(v, roiname, vox)
   106  	if err != nil {
   107  		return nil, err
   108  	}
   109  	if err := d.GetLabels(v, supervoxels, scale, vox, r); err != nil {
   110  		return nil, err
   111  	}
   112  	return vox.GetImage2d()
   113  }
   114  
   115  // GetVolume retrieves a n-d volume from a version node given a geometry of labels.
   116  func (d *Data) GetVolume(v dvid.VersionID, vox *Labels, supervoxels bool, scale uint8, roiname dvid.InstanceName) ([]byte, error) {
   117  	r, err := imageblk.GetROI(v, roiname, vox)
   118  	if err != nil {
   119  		return nil, err
   120  	}
   121  	if err := d.GetLabels(v, supervoxels, scale, vox, r); err != nil {
   122  		return nil, err
   123  	}
   124  	return vox.Data(), nil
   125  }
   126  
   127  type getOperation struct {
   128  	version     dvid.VersionID
   129  	voxels      *Labels
   130  	blocksInROI map[string]bool
   131  	mapping     *VCache
   132  }
   133  
   134  // GetLabels copies labels from the storage engine to Labels, a requested subvolume or 2d image.
   135  // If supervoxels is true, the returned labels are not mapped but are the raw supervoxels.
   136  func (d *Data) GetLabels(v dvid.VersionID, supervoxels bool, scale uint8, vox *Labels, r *imageblk.ROI) error {
   137  	store, err := datastore.GetOrderedKeyValueDB(d)
   138  	if err != nil {
   139  		return fmt.Errorf("Data type imageblk had error initializing store: %v\n", err)
   140  	}
   141  	if r != nil && scale != 0 {
   142  		return fmt.Errorf("DVID does not currently support ROI masks on lower-scale GETs")
   143  	}
   144  
   145  	// Only do one large request at a time, although each request can start many goroutines.
   146  	if vox.NumVoxels() > 256*256*256 {
   147  		server.LargeMutationMutex.Lock()
   148  		defer server.LargeMutationMutex.Unlock()
   149  	}
   150  
   151  	var mapping *VCache
   152  	if !supervoxels {
   153  		if mapping, err = getMapping(d, v); err != nil {
   154  			return err
   155  		}
   156  	}
   157  	wg := new(sync.WaitGroup)
   158  	ctx := datastore.NewVersionedCtx(d, v)
   159  
   160  	okv := store.(storage.BufferableOps)
   161  	// extract buffer interface
   162  	req, hasbuffer := okv.(storage.KeyValueRequester)
   163  	if hasbuffer {
   164  		okv = req.NewBuffer(ctx)
   165  	}
   166  
   167  	for it, err := vox.NewIndexIterator(d.BlockSize()); err == nil && it.Valid(); it.NextSpan() {
   168  		indexBeg, indexEnd, err := it.IndexSpan()
   169  		if err != nil {
   170  			return err
   171  		}
   172  		begTKey := NewBlockTKey(scale, indexBeg)
   173  		endTKey := NewBlockTKey(scale, indexEnd)
   174  
   175  		// Get set of blocks in ROI if ROI provided
   176  		var chunkOp *storage.ChunkOp
   177  		if r != nil && r.Iter != nil {
   178  			ptBeg := indexBeg.Duplicate().(dvid.ChunkIndexer)
   179  			ptEnd := indexEnd.Duplicate().(dvid.ChunkIndexer)
   180  			begX := ptBeg.Value(0)
   181  			endX := ptEnd.Value(0)
   182  
   183  			blocksInROI := make(map[string]bool, (endX - begX + 1))
   184  			c := dvid.ChunkPoint3d{begX, ptBeg.Value(1), ptBeg.Value(2)}
   185  			for x := begX; x <= endX; x++ {
   186  				c[0] = x
   187  				curIndex := dvid.IndexZYX(c)
   188  				if r.Iter.InsideFast(curIndex) {
   189  					indexString := string(curIndex.Bytes())
   190  					blocksInROI[indexString] = true
   191  				}
   192  			}
   193  			chunkOp = &storage.ChunkOp{&getOperation{v, vox, blocksInROI, mapping}, wg}
   194  		} else {
   195  			chunkOp = &storage.ChunkOp{&getOperation{v, vox, nil, mapping}, wg}
   196  		}
   197  
   198  		if !hasbuffer {
   199  			// Send the entire range of key-value pairs to chunk processor
   200  			err = okv.ProcessRange(ctx, begTKey, endTKey, chunkOp, storage.ChunkFunc(d.ReadChunk))
   201  			if err != nil {
   202  				return fmt.Errorf("Unable to GET data %s: %v", ctx, err)
   203  			}
   204  		} else {
   205  			// Extract block list
   206  			var tkeys []storage.TKey
   207  			ptBeg := indexBeg.Duplicate().(dvid.ChunkIndexer)
   208  			ptEnd := indexEnd.Duplicate().(dvid.ChunkIndexer)
   209  			begX := ptBeg.Value(0)
   210  			endX := ptEnd.Value(0)
   211  
   212  			c := dvid.ChunkPoint3d{begX, ptBeg.Value(1), ptBeg.Value(2)}
   213  			for x := begX; x <= endX; x++ {
   214  				c[0] = x
   215  				tk := NewBlockTKeyByCoord(scale, c.ToIZYXString())
   216  				tkeys = append(tkeys, tk)
   217  			}
   218  
   219  			err = okv.(storage.RequestBuffer).ProcessList(ctx, tkeys, chunkOp, storage.ChunkFunc(d.ReadChunk))
   220  			if err != nil {
   221  				return fmt.Errorf("Unable to GET data %s: %v", ctx, err)
   222  			}
   223  		}
   224  	}
   225  
   226  	if hasbuffer {
   227  		// submit the entire buffer to the DB
   228  		err = okv.(storage.RequestBuffer).Flush()
   229  
   230  		if err != nil {
   231  			return fmt.Errorf("Unable to GET data %s: %v", ctx, err)
   232  		}
   233  	}
   234  
   235  	if err != nil {
   236  		return err
   237  	}
   238  	wg.Wait()
   239  	return nil
   240  }
   241  
   242  // ReadChunk reads a chunk of data as part of a mapped operation.
   243  // Only some multiple of the # of CPU cores can be used for chunk handling before
   244  // it waits for chunk processing to abate via the buffered server.HandlerToken channel.
   245  func (d *Data) ReadChunk(chunk *storage.Chunk) error {
   246  	server.CheckChunkThrottling()
   247  	go d.readChunk(chunk)
   248  	return nil
   249  }
   250  
   251  func (d *Data) readChunk(chunk *storage.Chunk) {
   252  	defer func() {
   253  		// After processing a chunk, return the token.
   254  		server.HandlerToken <- 1
   255  
   256  		// Notify the requestor that this chunk is done.
   257  		if chunk.Wg != nil {
   258  			chunk.Wg.Done()
   259  		}
   260  	}()
   261  
   262  	op, ok := chunk.Op.(*getOperation)
   263  	if !ok {
   264  		log.Fatalf("Illegal operation passed to readChunk() for data %s\n", d.DataName())
   265  	}
   266  
   267  	// Make sure our received chunk is valid.
   268  	if chunk == nil {
   269  		dvid.Errorf("Received nil chunk in readChunk.  Ignoring chunk.\n")
   270  		return
   271  	}
   272  	if chunk.K == nil {
   273  		dvid.Errorf("Received nil chunk key in readChunk.  Ignoring chunk.\n")
   274  		return
   275  	}
   276  
   277  	// If there's an ROI, if outside ROI, use blank buffer.
   278  	var zeroOut bool
   279  	_, indexZYX, err := DecodeBlockTKey(chunk.K)
   280  	if err != nil {
   281  		dvid.Errorf("Error processing voxel block: %s\n", err)
   282  		return
   283  	}
   284  	if op.blocksInROI != nil {
   285  		indexString := string(indexZYX.Bytes())
   286  		_, insideROI := op.blocksInROI[indexString]
   287  		if !insideROI {
   288  			zeroOut = true
   289  		}
   290  	}
   291  
   292  	// Initialize the block buffer using the chunk of data.  For voxels, this chunk of
   293  	// data needs to be uncompressed and deserialized.
   294  	var block labels.Block
   295  	if zeroOut || chunk.V == nil {
   296  		blockSize, ok := d.BlockSize().(dvid.Point3d)
   297  		if !ok {
   298  			dvid.Errorf("Block size for data %q is not 3d: %s\n", d.DataName(), d.BlockSize())
   299  			return
   300  		}
   301  		block = *labels.MakeSolidBlock(0, blockSize)
   302  	} else {
   303  		var data []byte
   304  		data, _, err = dvid.DeserializeData(chunk.V, true)
   305  		if err != nil {
   306  			dvid.Errorf("Unable to deserialize block in %q: %v\n", d.DataName(), err)
   307  			return
   308  		}
   309  		if err := block.UnmarshalBinary(data); err != nil {
   310  			dvid.Errorf("Unable to unmarshal labels Block compression in %q: %v\n", d.DataName(), err)
   311  			return
   312  		}
   313  	}
   314  
   315  	// Perform the operation.
   316  	if op.mapping != nil {
   317  		err = modifyBlockMapping(op.version, &block, op.mapping)
   318  		if err != nil {
   319  			dvid.Errorf("unable to modify block %s mapping: %v\n", indexZYX, err)
   320  		}
   321  	}
   322  	if err = op.voxels.readBlock(chunk.K, block, d.BlockSize()); err != nil {
   323  		dvid.Errorf("readBlock, key %v in %q: %v\n", chunk.K, d.DataName(), err)
   324  	}
   325  }
   326  
   327  // overwrites labels in header with their mapped values, so converts blocks from
   328  // supervoxels to body labels
   329  func modifyBlockMapping(v dvid.VersionID, block *labels.Block, vc *VCache) error {
   330  	mappedVersions := vc.getMappedVersionsDist(v)
   331  	for i, label := range block.Labels {
   332  		mapped, found := vc.mapLabel(label, mappedVersions)
   333  		if found {
   334  			block.Labels[i] = mapped
   335  		}
   336  	}
   337  	return nil
   338  }
   339  
   340  func (v *Labels) readBlock(tkey storage.TKey, block labels.Block, blockSize dvid.Point) error {
   341  	if blockSize.NumDims() > 3 {
   342  		return fmt.Errorf("DVID voxel blocks currently only supports up to 3d, not 4+ dimensions")
   343  	}
   344  	blockBeg, dataBeg, dataEnd, err := v.ComputeTransform(tkey, blockSize)
   345  	if err != nil {
   346  		return err
   347  	}
   348  
   349  	// TODO -- Refactor this function to make direct use of compressed label Block without
   350  	// conversion into full label array.
   351  	labels := v.Data()
   352  	blabels, _ := block.MakeLabelVolume()
   353  
   354  	// Compute the strides (in bytes)
   355  	bX := int64(blockSize.Value(0)) * 8
   356  	bY := int64(blockSize.Value(1)) * bX
   357  	dX := int64(v.Stride())
   358  
   359  	blockBegX := int64(blockBeg.Value(0))
   360  	blockBegY := int64(blockBeg.Value(1))
   361  	blockBegZ := int64(blockBeg.Value(2))
   362  
   363  	// Do the transfers depending on shape of the external voxels.
   364  	switch {
   365  	case v.DataShape().Equals(dvid.XY):
   366  		dataI := int64(dataBeg.Value(1))*dX + int64(dataBeg.Value(0))*8
   367  		blockI := blockBegZ*bY + blockBegY*bX + blockBegX*8
   368  		bytes := int64(dataEnd.Value(0)-dataBeg.Value(0)+1) * 8
   369  		for y := dataBeg.Value(1); y <= dataEnd.Value(1); y++ {
   370  			copy(labels[dataI:dataI+bytes], blabels[blockI:blockI+bytes])
   371  			blockI += bX
   372  			dataI += dX
   373  		}
   374  
   375  	case v.DataShape().Equals(dvid.XZ):
   376  		dataI := int64(dataBeg.Value(2))*dX + int64(dataBeg.Value(0))*8
   377  		blockI := blockBegZ*bY + blockBegY*bX + blockBegX*8
   378  		bytes := int64(dataEnd.Value(0)-dataBeg.Value(0)+1) * 8
   379  		for y := dataBeg.Value(2); y <= dataEnd.Value(2); y++ {
   380  			copy(labels[dataI:dataI+bytes], blabels[blockI:blockI+bytes])
   381  			blockI += bY
   382  			dataI += dX
   383  		}
   384  
   385  	case v.DataShape().Equals(dvid.YZ):
   386  		bz := blockBegZ
   387  		for y := int64(dataBeg.Value(2)); y <= int64(dataEnd.Value(2)); y++ {
   388  			dataI := y*dX + int64(dataBeg.Value(1))*8
   389  			blockI := bz*bY + blockBegY*bX + blockBegX*8
   390  			for x := dataBeg.Value(1); x <= dataEnd.Value(1); x++ {
   391  				copy(labels[dataI:dataI+8], blabels[blockI:blockI+8])
   392  				blockI += bX
   393  				dataI += 8
   394  			}
   395  			bz++
   396  		}
   397  
   398  	case v.DataShape().ShapeDimensions() == 2:
   399  		// TODO: General code for handling 2d ExtData in n-d space.
   400  		return fmt.Errorf("DVID currently does not support 2d in n-d space.")
   401  
   402  	case v.DataShape().Equals(dvid.Vol3d):
   403  		blockOffset := blockBegX * 8
   404  		dX := int64(v.Size().Value(0)) * 8
   405  		dY := int64(v.Size().Value(1)) * dX
   406  		dataOffset := int64(dataBeg.Value(0)) * 8
   407  		bytes := int64(dataEnd.Value(0)-dataBeg.Value(0)+1) * 8
   408  		blockZ := blockBegZ
   409  
   410  		for dataZ := int64(dataBeg.Value(2)); dataZ <= int64(dataEnd.Value(2)); dataZ++ {
   411  			blockY := blockBegY
   412  			for dataY := int64(dataBeg.Value(1)); dataY <= int64(dataEnd.Value(1)); dataY++ {
   413  				blockI := blockZ*bY + blockY*bX + blockOffset
   414  				dataI := dataZ*dY + dataY*dX + dataOffset
   415  				copy(labels[dataI:dataI+bytes], blabels[blockI:blockI+bytes])
   416  				blockY++
   417  			}
   418  			blockZ++
   419  		}
   420  
   421  	default:
   422  		return fmt.Errorf("Cannot readBlock() unsupported voxels data shape %s", v.DataShape())
   423  	}
   424  	return nil
   425  }