volume-visualization/volumes/models.py

import io
import os

from mongoengine import Document, EmbeddedDocument, signals, GridFSProxy
from mongoengine.fields import BooleanField, IntField, FileField, ListField, EmbeddedDocumentField, StringField

from django.conf import settings

from bson.objectid import ObjectId

import numpy as np

import logging

logger = logging.getLogger(__name__)


class OriginalSize(EmbeddedDocument):
    sizex = IntField()
    sizey = IntField()
    sizez = IntField()


class TextureInfo(EmbeddedDocument):
    slicesOverX = IntField()
    slicesOverY = IntField()
    numberOfSlices = IntField()
    imgMin = IntField()
    imgMax = IntField()
    multipleFrames = BooleanField()
    frameFrom = IntField()
    frameTo = IntField()
    slicesPerSprite = IntField()
    numberOfSprites = IntField()

    thresholdIndex_otsu = IntField()
    thresholdIndex_yen  =  IntField()
    thresholdIndex_isodata  =  IntField()

    @staticmethod
    def from_info_object(infoObject):
        textureInfo = TextureInfo()
        textureInfo.slicesOverX = infoObject['slicesOverX']
        textureInfo.slicesOverY = infoObject['slicesOverY']
        textureInfo.numberOfSlices = infoObject['numberOfSlices']
        textureInfo.imgMin = infoObject['imgMin']
        textureInfo.imgMax = infoObject['imgMax']
        textureInfo.numberOfSprites = infoObject['numberOfSprites']
        textureInfo.slicesPerSprite = infoObject['slicesPerSprite']

        textureInfo.multipleFrames = infoObject['multipleFrames']
        textureInfo.frameFrom = infoObject['frameFrom']
        textureInfo.frameTo = infoObject['frameTo']

        textureInfo.thresholdIndex_otsu = infoObject['thresholdIndex_otsu']
        textureInfo.thresholdIndex_yen  =  infoObject['thresholdIndex_yen']
        textureInfo.thresholdIndex_isodata = infoObject['thresholdIndex_isodata']

        return textureInfo


class Texture(EmbeddedDocument):
    size = IntField()
    info = EmbeddedDocumentField(TextureInfo)
    contenttype = StringField()
    sprites = ListField(ListField(FileField()))
    originalSize = EmbeddedDocumentField(OriginalSize)

    def deleteContent(self):
        for sprite in self.sprites:
            for fileItem in sprite:
                fileItem.delete()

    def addSpriteBuffers(self, spriteBuffers):
        proxies = []
        for spriteBuffer in spriteBuffers:
            gridFSProxy = GridFSProxy()
            gridFSProxy.put(spriteBuffer)
            proxies.append(gridFSProxy)

        self.sprites.append(proxies)


class RawFrame(EmbeddedDocument):
    dataName = StringField()
    frameNumber = IntField()
    dtype = StringField()
    originalSize = EmbeddedDocumentField(OriginalSize)

    def setNumpyArray(self, array, frameNumber, pathToStore):
        if len(array.shape) != 3:
            raise ValueError('array has not dimension of 3')

        originalSize = OriginalSize()
        originalSize.sizez, originalSize.sizey, originalSize.sizex = array.shape
        self.originalSize = originalSize
        self.dtype = str(array.dtype)
        self.frameNumber = frameNumber

        logger.debug('RawFrame.setNumpyArray: write frame to buffer')
        mmapId = ObjectId()
        self.dataName = str(mmapId) + '.npy'

        filepath = self.__getpathtofile(pathToStore)
        mmap = np.memmap(filepath, dtype=array.dtype, mode='w+', shape=array.shape)

        logger.debug('RawFrame.setNumpyArray: put buffer to file')
        mmap[:] = array[:]

    def getNumpyArray(self, path):
        filepath = self.__getpathtofile(path)
        shape = (self.originalSize.sizez, self.originalSize.sizey, self.originalSize.sizex)
        mmap = np.memmap(filepath, dtype=self.dtype, mode='r', shape=shape)
        return mmap

    def deleteContent(self, path):
        try:
            os.remove(self.__getpathtofile(path))
        except OSError:
            logger.error('file %s could not be deleted' % self.__getpathtofile(path))

    def __getpathtofile(self, path):
        return path + self.dataName


class RawFrameGridFsProxy(EmbeddedDocument):
    data = FileField()
    frameNumber = IntField()
    dtype = StringField()
    originalSize = EmbeddedDocumentField(OriginalSize)

    def setNumpyArray(self, array, frameNumber, pathToStore):
        if len(array.shape) != 3:
            raise ValueError('array has not dimension of 3')

        self.dtype = str(array.dtype)
        self.frameNumber = frameNumber

        buff = io.BytesIO()
        buff.write(array)
        buff.seek(0, 0)
        gridFSProxy = GridFSProxy()
        gridFSProxy.put(buff)
        self.data = gridFSProxy

    def getNumpyArray(self, path):
        array = np.frombuffer(self.data.read(), dtype=np.dtype(self.dtype))
        array = np.reshape(array, (self.originalSize.sizez, self.originalSize.sizey, self.originalSize.sizex))

        self.data.seek(0, 0)

        return array

    def deleteContent(self, path):
        self.data.delete()

    def __getpathtofile(self, path):
        return path + self.dataName


class Volume(Document):
    name = StringField(max_length=100)
    path = StringField(max_length=1000)
    rawData = BooleanField()
    width = IntField()
    height = IntField()
    littleEndian = BooleanField()
    slices = IntField()
    sliceFrom = IntField()
    sliceTo = IntField()
    textures = ListField(EmbeddedDocumentField(Texture))
    rawFrames = ListField(EmbeddedDocumentField(RawFrame))
    generateSlicesStatus = StringField()
    imageSequence = BooleanField()
    automaticProcessing = BooleanField()

    @classmethod
    def pre_save(cls, sender, document, **kwargs):
        if document.sliceFrom is not None and document.sliceTo is not None:
            document.slices = document.sliceTo - document.sliceFrom

    def clear_textures(self):
        for texture in self.textures:
            self.remove_texture(texture.size)

    def add_slice_maps(self, size, sprites, originalSize, infoObject):
        textureSize = int(size)

        texture = Texture()
        texture.sprites = []

        texture.info = TextureInfo.from_info_object(infoObject)
        texture.contenttype = None
        texture.originalSize = originalSize

        texture.size = textureSize

        for sprite in sprites:
            if texture.contenttype is None:
                texture.contenttype = sprite.contenttype

            texture.addSpriteBuffers(sprite.buffers)

        self.insert_texture(texture)

    def insert_texture(self, textureToAdd):
        for texture in self.textures:
            if texture.size == textureToAdd.size:
                self.textures.remove(texture)

        return Volume.objects(id=self.id).update_one(push__textures=textureToAdd)

    def remove_texture(self, textureSize):
        for texture in self.textures:
            if texture.size == textureSize:
                texture.deleteContent()
                break

        Volume.objects(id=self.id).update_one(pull__textures={'size': textureSize})

    def add_frame(self, frame, frameNumber):
        rawFrame = RawFrame()
        rawFrame.setNumpyArray(frame, frameNumber, self.get_mmap_path())

        if settings.REDUCE_HARDDISK_ACCESS_FOR_DEV is True:
            logger.debug('not saving rawframes in volume --> reducing harddisk access')
        else:
            Volume.objects(id=self.id).update_one(push__rawFrames=rawFrame)

        return rawFrame

    def get_frame_numpy_array(self, frameNumber):
        rawFrame = self.rawFrames[frameNumber]
        return rawFrame.getNumpyArray(self.get_mmap_path())

    def remove_frames(self):
        for frame in self.rawFrames:
            frame.deleteContent(self.get_mmap_path())

        Volume.objects(id=self.id).update_one(set__rawFrames=[])

    def get_mmap_path(self):
        path = settings.PATH_MMAP + str(self.id) + '/'
        if os.path.exists(path) is False:
            os.makedirs(path)

        return path

    def delete(self):
        self.remove_frames()

        for texture in self.textures:
            self.remove_texture(texture.size)

        return super(Volume, self).delete()

    @property
    def rendering_possible(self):
        return 'running' not in self.generateSlicesStatus and len(self.textures) > 0

    @staticmethod
    def read_sprite(volumeId, spriteId):
        spriteId = ObjectId(spriteId)
        volume = Volume.objects(id=volumeId)[0]

        for texture in volume.textures:
            for sprite in texture.sprites:
                for fileItem in sprite:
                    if fileItem._id == spriteId:
                        return fileItem, texture.contenttype

        raise IndexError('sprite with size %s could not be found' % spriteId)

    @staticmethod
    def read_texture(id, textureSize):
        volume = Volume.objects(id=id)[0]

        if textureSize is None:
            logger.debug('returning texture0 with its size: %d' % volume.textures[0].size)
            return volume.textures[0]

        textureSize = int(textureSize)
        for texture in volume.textures:
            if texture.size == textureSize:
                logger.debug('texture found, returing')
                return texture

        raise IndexError('texture with size %s could not be found' % textureSize)


signals.pre_save.connect(Volume.pre_save, sender=Volume)