hycore

  1"""
  2A functional data structure for creating, managing and analysing hyperspectral drillcore datasets using python and hylite.
  3"""
  4
  5from .coreshed import *
  6
  7import os
  8from pathlib import Path
  9import shutil
 10
 11sandbox = None
 12
 13def loadShed( path ):
 14    """
 15    Load a shed Object
 16    :param path: Open the specified Shed.
 17    :return: A hycore.Shed instance.
 18    """
 19
 20    # check shed exists and is a directory
 21    pth = os.path.splitext(path)[0] + ".shed"
 22    if not os.path.exists(pth) and os.path.isdir(pth):
 23        raise FileNotFoundError("%s is not a valid hycore.Shed" % path)
 24    name = os.path.basename( os.path.splitext( path )[0] )
 25
 26    # load header (if one exists)
 27    from hylite import io
 28    header = None
 29    hdr = os.path.splitext(path)[0] + ".hdr"
 30    if os.path.exists(hdr):
 31        header = io.loadHeader( hdr )
 32
 33    # create and return Shed instance
 34    from hycore.coreshed import Shed
 35    S = Shed(name, os.path.dirname( pth ) , header )
 36
 37    # load holes and boxes to ensure correct dtype (Hole or Box, not HyCollection)
 38    _ = S.getBoxes()
 39
 40    return Shed(name, os.path.dirname( pth ) , header )
 41
 42def get_sandbox( fill=False, vis=False, mosaic=False ):
 43    """
 44    Return a sandbox directory on the local file system for reading / writing temporary files.
 45    :param fill: True if a dummy hyperspectral core database should be added to the sandbox directory. If this is
 46                  set, a shed will be returned rather than the directory path.
 47    :param vis: True if .png previews should be generated in the shed.
 48    :param mosaic: True if .png mosiacs should be generated for each drillcore.
 49    :return: A Path defining the sandbox directory or a Shed instance (if fill = True).
 50    """
 51    global sandbox
 52    if sandbox is None:
 53        sandbox = Path( os.path.join(os.getcwd(), 'sandbox') )
 54    os.makedirs( sandbox, exist_ok = True ) # if we are getting it, the sandbox dir should exist!
 55
 56    from hycore.generator import genTray
 57    from hycore.coreshed import Shed
 58
 59    if fill:
 60        # create a dummy coreshed!
 61        S = Shed('eldorado', sandbox)
 62
 63        # add fake holes and data
 64        h1 = [False, False, True]  # True / False are lithology codes
 65        h2 = [True, False, True]
 66        h3 = [False]  # include a hole with only one tray for testing purposes
 67        for h, n, c, b in zip([h1, h2, h3], [4, 5, 4], ['H01', 'H03', 'H02'], [1.0, 0.4, 1.0]):
 68            trays = [genTray(n, l) for l in h]
 69            ids = [i + 1 for i in range(len(h))]
 70            for (timg, A, mask), i in zip(trays, ids):
 71                # add Fenix image
 72                timg.data[..., timg.get_band_index(2180.):timg.get_band_index(
 73                    2210.)] *= b  # adjust brightness of 2200 nm band to provide visual difference between different holes
 74                t = S.addHSI('FENIX', c, '%03d' % i, timg)
 75
 76                # add mask
 77                t.mask = hylite.HyImage(mask.astype(np.uint8))
 78
 79                # add a result
 80                from hylite.analyse import band_ratio
 81                R = band_ratio(t.FENIX, 2230., 2200.)
 82                G = band_ratio(t.FENIX, 2190., 2200.)
 83                B = band_ratio(t.FENIX, 2140., 2160.)
 84                stack = hylite.HyImage(np.clip(np.dstack([(R.data - 1.0) / 0.2,
 85                                                          (G.data - 1.) / 0.2,
 86                                                          (B.data - 1.0) / 0.2]), 0, 1))
 87                stack.data = (stack.data * 255).astype(np.uint8)  # convert to RGB
 88                t.results.BR_Clays = stack
 89
 90                def saveLegend(red, green, blue, path):
 91                    import matplotlib.pyplot as plt
 92                    tr = [np.nan, np.nan, np.nan]
 93                    plt.figure(figsize=(4, 0.5))
 94                    plt.fill(tr, tr, label=red, color='r')
 95                    plt.fill(tr, tr, label=green, color='g')
 96                    plt.fill(tr, tr, label=blue, color='b')
 97                    plt.legend(ncol=3, loc='center')
 98                    plt.axis('off')
 99                    os.makedirs(os.path.dirname(path), exist_ok=True)
100                    plt.tight_layout()
101                    plt.savefig(path, dpi=300)
102                    plt.close()
103
104                saveLegend(r"R= ($\frac{2230}{2200}$)",
105                           r"G= ($\frac{2190}{2200}$)",
106                           r"B= ($\frac{2140}{2160}$)",
107                           os.path.join(t.results.getDirectory(), 'LEG_Clays.png'))
108
109                # add fake image as though from another sensor
110                lwir = timg.copy()
111                lwir.data = 1. - lwir.data
112                lwir.set_wavelengths(timg.get_wavelengths() * 10 - 13000)
113                t = S.addHSI('LWIR', c, '%03d' % i, lwir)
114                t.start = (i - 1) * n
115                t.end = i * n
116                t.save()
117                t.free()
118
119        # add a box that is non-contiguous (has a gap)
120        t = S.addHSI('FENIX', 'H01', '004', timg)
121        t = S.addHSI('LWIR', 'H01', '004', lwir)
122        t.start = 15
123        t.end = 20
124        t.mask = hylite.HyImage(mask.astype(np.uint8))
125        t.results.BR_Clays = stack
126        t.save()
127        t.free()
128        saveLegend(r"R= ($\frac{2230}{2200}$)",
129                   r"G= ($\frac{2190}{2200}$)",
130                   r"B= ($\frac{2140}{2160}$)",
131                   os.path.join(t.results.getDirectory(), 'LEG_Clays.png'))
132
133        if vis:
134            for b in S.getBoxes():
135                b.updateVis(qaqc=False, thumb='FENIX')
136
137        if mosaic:
138            from hycore.templates import unwrap_tray, Template, Canvas, compositeStack
139            for i, h in enumerate(S.getHoles()):
140                # get pole mosaic of core
141                T = h.getTemplate(method='pole', res=2e-3)
142
143                # apply template to generate strips for various products
144                for f in ['FENIX.png', 'LWIR.png', 'BR_Clays.png']:
145                    m = T.apply(f, (0,1,2))
146                    h.results.set(f, m)
147                h.results.save()
148                h.save()
149                h.free()
150
151            # add some example annotations to the final hole
152            h.annotate("Vein 1", "I think there's gold here!", depth_from=0.5, group='Gold')
153            h.annotate("Vein 2", "No gold here.", depth_from=1, depth_to=2, group='Gold')
154            h.annotate("Fault 1", "What a bloody mess!", depth_from=0.5, depth_to=2, group='Faults')
155            h.annotate("Fault 2", "https://en.wikipedia.org/wiki/San_Andreas_Fault",
156                       depth_from=2.1, depth_to=2.1, group='Faults', type='link')
157            h.save()
158
159        S.save()
160        S.free()
161        return S
162    else:
163        return sandbox
164
165
166def set_sandbox( path ):
167    """
168    Set the sandbox directory for hycore to use.
169    :param path: A path to the new sandbox directory
170    :return: None
171    """
172    global sandbox
173    sandbox = sandbox
174
175def empty_sandbox():
176    """
177    Remove the sandbox directory (if it exists).
178    :return: None
179    """
180    shutil.rmtree( get_sandbox() )