pyosv.io.writer

 1from ..utils.paths import get_path_gui
 2
 3import matplotlib.pyplot as plt
 4import numpy as np
 5import rasterio
 6import netCDF4
 7
 8
 9def write(image : np.ndarray, path : str, meta : dict) -> None:
10    '''
11        Save an image and its metadata given a path.
12
13        Supported data format
14
15        RASTERIO_EXTENSIONS   = ['.tif', '.tiff', '.geotiff']  
16        MATPLOTLIB_EXTENSIONS = ['.png', '.jpg', 'jpeg', 'jp2']
17        NETCDF4_EXTENSIONS    = ['.nc']
18
19        Data must always be in channel last format.
20
21        If image extension is in MATPLOTLIB_EXTENSIONS, metadata can be None.
22        If image extension is in NETCDF4_EXTENSIONS, metadata can be None.
23
24        Parameters:
25        -----------
26            - image : np.ndarray
27                the WxHxC image to be saved, with W width, H height and B bands (channel last)
28            - path : str 
29                position of the image, if None the function will ask for the image path using a menu
30            - meta : dict
31                metadata for the image to be saved
32        
33        Returns:
34        --------
35        Nothing, the image will be saved
36
37        Usage:
38        ------
39
40        ```python
41        import numpy as np
42
43        img         = np.array(  
44            [[
45                [0.1, 0.2, 0.3],  
46                [0.4, 0.5, 0.6],  
47                [0.7, 0.8, 0.9]
48             ],
49             [
50                [1.1, 1.2, 1.3],  
51                [1.4, 1.5, 1.6],  
52                [1.7, 1.8, 1.9]
53             ]
54            ]  
55        ) 
56
57        # Making channels last
58        img = np.moveaxis(img, 0, -1)
59
60        write(img, 'path/to/save/img.png')
61        
62        ```
63        Output:
64        -------
65        Nothing, the image will be saved
66
67
68    '''
69
70    RASTERIO_EXTENSIONS   = ['.tif', '.tiff']
71    MATPLOTLIB_EXTENSIONS = ['.png', '.jpg', 'jpeg', 'jp2']
72    NETCDF4_EXTENSIONS    = ['.nc']
73
74
75    if path is None:
76        path = get_path_gui()
77
78    if any(frmt in path for frmt in RASTERIO_EXTENSIONS):
79        if image!=None:
80            meta.update({'driver':'GTiff',
81                            'width':image.shape[0],
82                            'height':image.shape[1],
83                            'count':image.shape[2],
84                            'dtype':'float64'})
85
86        with rasterio.open(fp=path, mode='w',**meta) as dst:
87            for count in range(image.shape[2]):
88                dst.write(image[:,:,count], count)
89
90    elif any(frmt in path for frmt in MATPLOTLIB_EXTENSIONS):
91        plt.imsave(path, image)
92    elif any(frmt in path for frmt in NETCDF4_EXTENSIONS):
93        raise Exception('Error: [under dev] currently netCDF4 files can not be saved!')
94    else:
95        raise Exception('Error: file can not be saved, format not supported!')