1"""
2Utility functions to read/write image data from/to file
3"""
4
5import csv
6import json
7import os
8import logging
9import warnings
10from xml.etree import ElementTree
11
12import numpy as np
13import tifffile as tiff
14from PIL import Image
15
16from .decoding import decode_tar, get_data_format, fix_jp2_image, get_jp2_bit_depth
17from .constants import MimeType
18from .os_utils import create_parent_folder
19
20
21warnings.simplefilter('ignore', Image.DecompressionBombWarning)
22LOGGER = logging.getLogger(__name__)
23
24CSV_DELIMITER = ';'
25
26
27def read_data(filename, data_format=None):
28    """ Read image data from file
29
30    This function reads input data from file. The format of the file
31    can be specified in ``data_format``. If not specified, the format is
32    guessed from the extension of the filename.
33
34    :param filename: filename to read data from
35    :type filename: str
36    :param data_format: format of filename. Default is `None`
37    :type data_format: MimeType
38    :return: data read from filename
39    :raises: exception if filename does not exist
40    """
41    if not os.path.exists(filename):
42        raise FileNotFoundError(f'Filename {filename} does not exist')
43
44    if not isinstance(data_format, MimeType):
45        data_format = get_data_format(filename)
46
47    reader = _get_reader(data_format)
48
49    try:
50        return reader(filename)
51    except BaseException as exception:
52        # In case a procedure would read a lot of files and one would be corrupt this helps us figure out which one
53        LOGGER.debug('Failed to read from file: %s', filename)
54        raise exception
55
56
57def _get_reader(data_format):
58    """ Provides a function for reading data in a given data format
59    """
60    if data_format is MimeType.TIFF:
61        return read_tiff_image
62    if data_format is MimeType.JP2:
63        return read_jp2_image
64    if data_format.is_image_format():
65        return read_image
66    try:
67        return {
68            MimeType.TAR: read_tar,
69            MimeType.TXT: read_text,
70            MimeType.RAW: _read_binary,
71            MimeType.CSV: read_csv,
72            MimeType.JSON: read_json,
73            MimeType.XML: read_xml,
74            MimeType.GML: read_xml,
75            MimeType.SAFE: read_xml
76        }[data_format]
77    except KeyError as exception:
78        raise ValueError(f'Reading data format {data_format} is not supported') from exception
79
80
81def read_tar(filename):
82    """ Read a tar from file
83    """
84    with open(filename, 'rb') as file:
85        return decode_tar(file)
86
87
88def read_tiff_image(filename):
89    """ Read data from TIFF file
90
91    :param filename: name of TIFF file to be read
92    :type filename: str
93    :return: data stored in TIFF file
94    """
95    return tiff.imread(filename)
96
97
98def read_jp2_image(filename):
99    """ Read data from JPEG2000 file
100
101    :param filename: name of JPEG2000 file to be read
102    :type filename: str
103    :return: data stored in JPEG2000 file
104    """
105    # Other option:
106    # return glymur.Jp2k(filename)[:]
107    image = read_image(filename)
108
109    with open(filename, 'rb') as file:
110        bit_depth = get_jp2_bit_depth(file)
111
112    return fix_jp2_image(image, bit_depth)
113
114
115def read_image(filename):
116    """ Read data from PNG or JPG file
117
118    :param filename: name of PNG or JPG file to be read
119    :type filename: str
120    :return: data stored in JPG file
121    """
122    return np.array(Image.open(filename))
123
124
125def read_text(filename):
126    """ Read data from text file
127
128    :param filename: name of text file to be read
129    :type filename: str
130    :return: data stored in text file
131    """
132    with open(filename, 'r') as file:
133        return file.read()
134
135
136def _read_binary(filename):
137    """ Reads data in bytes
138    """
139    with open(filename, 'rb') as file:
140        return file.read()
141
142
143def read_csv(filename, delimiter=CSV_DELIMITER):
144    """ Read data from CSV file
145
146    :param filename: name of CSV file to be read
147    :type filename: str
148    :param delimiter: type of CSV delimiter. Default is ``;``
149    :type delimiter: str
150    :return: data stored in CSV file as list
151    """
152    with open(filename, 'r') as file:
153        return list(csv.reader(file, delimiter=delimiter))
154
155
156def read_json(filename):
157    """ Read data from JSON file
158
159    :param filename: name of JSON file to be read
160    :type filename: str
161    :return: data stored in JSON file
162    """
163    with open(filename, 'rb') as file:
164        return json.load(file)
165
166
167def read_xml(filename):
168    """ Read data from XML or GML file
169
170    :param filename: name of XML or GML file to be read
171    :type filename: str
172    :return: data stored in XML file
173    """
174    return ElementTree.parse(filename)
175
176
177def read_numpy(filename):
178    """ Read data from numpy file
179
180    :param filename: name of numpy file to be read
181    :type filename: str
182    :return: data stored in file as numpy array
183    """
184    return np.load(filename)
185
186
187def write_data(filename, data, data_format=None, compress=False, add=False):
188    """ Write image data to file
189
190    Function to write image data to specified file. If file format is not provided
191    explicitly, it is guessed from the filename extension. If format is TIFF, geo
192    information and compression can be optionally added.
193
194    :param filename: name of file to write data to
195    :type filename: str
196    :param data: image data to write to file
197    :type data: numpy array
198    :param data_format: format of output file. Default is `None`
199    :type data_format: MimeType
200    :param compress: whether to compress data or not. Default is `False`
201    :type compress: bool
202    :param add: whether to append to existing text file or not. Default is `False`
203    :type add: bool
204    :raises: exception if numpy format is not supported or file cannot be written
205    """
206    create_parent_folder(filename)
207
208    if not isinstance(data_format, MimeType):
209        data_format = get_data_format(filename)
210
211    if data_format is MimeType.TIFF:
212        return write_tiff_image(filename, data, compress)
213    if data_format.is_image_format():
214        return write_image(filename, data)
215    if data_format is MimeType.TXT:
216        return write_text(filename, data, add=add)
217
218    try:
219        return {
220            MimeType.RAW: write_bytes,
221            MimeType.CSV: write_csv,
222            MimeType.JSON: write_json,
223            MimeType.XML: write_xml,
224            MimeType.GML: write_xml
225        }[data_format](filename, data)
226    except KeyError as exception:
227        raise ValueError(f'Writing data format {data_format} is not supported') from exception
228
229
230def write_tiff_image(filename, image, compress=False):
231    """ Write image data to TIFF file
232
233    :param filename: name of file to write data to
234    :type filename: str
235    :param image: image data to write to file
236    :type image: numpy array
237    :param compress: whether to compress data. If `True`, lzma compression is used. Default is `False`
238    :type compress: bool
239    """
240    if compress:
241        return tiff.imsave(filename, image, compress='lzma')  # lossless compression, works very well on masks
242    return tiff.imsave(filename, image)
243
244
245def write_jp2_image(filename, image):
246    """ Write image data to JPEG2000 file
247
248    :param filename: name of JPEG2000 file to write data to
249    :type filename: str
250    :param image: image data to write to file
251    :type image: numpy array
252    :return: jp2k object
253    """
254    # Other options:
255    # return glymur.Jp2k(filename, data=image)
256    # cv2.imwrite(filename, image)
257    return write_image(filename, image)
258
259
260def write_image(filename, image):
261    """ Write image data to PNG, JPG file
262
263    :param filename: name of PNG or JPG file to write data to
264    :type filename: str
265    :param image: image data to write to file
266    :type image: numpy array
267    """
268    data_format = get_data_format(filename)
269    if data_format is MimeType.JPG:
270        LOGGER.warning('Warning: jpeg is a lossy format therefore saved data will be modified.')
271    return Image.fromarray(image).save(filename)
272
273
274def write_text(filename, data, add=False):
275    """ Write image data to text file
276
277    :param filename: name of text file to write data to
278    :type filename: str
279    :param data: image data to write to text file
280    :type data: numpy array
281    :param add: whether to append to existing file or not. Default is `False`
282    :type add: bool
283    """
284    write_type = 'a' if add else 'w'
285    with open(filename, write_type) as file:
286        print(data, end='', file=file)
287
288
289def write_csv(filename, data, delimiter=CSV_DELIMITER):
290    """ Write image data to CSV file
291
292    :param filename: name of CSV file to write data to
293    :type filename: str
294    :param data: image data to write to CSV file
295    :type data: numpy array
296    :param delimiter: delimiter used in CSV file. Default is ``;``
297    :type delimiter: str
298    """
299    with open(filename, 'w') as file:
300        csv_writer = csv.writer(file, delimiter=delimiter)
301        for line in data:
302            csv_writer.writerow(line)
303
304
305def write_json(filename, data):
306    """ Write data to JSON file
307
308    :param filename: name of JSON file to write data to
309    :type filename: str
310    :param data: data to write to JSON file
311    :type data: list, tuple
312    """
313    with open(filename, 'w') as file:
314        json.dump(data, file, indent=4, sort_keys=True)
315
316
317def write_xml(filename, element_tree):
318    """ Write data to XML or GML file
319
320    :param filename: name of XML or GML file to write data to
321    :type filename: str
322    :param element_tree: data as ElementTree object
323    :type element_tree: xmlElementTree
324    """
325    return element_tree.write(filename)
326    # this will write declaration tag in first line:
327    # return element_tree.write(filename, encoding='utf-8', xml_declaration=True)
328
329
330def write_numpy(filename, data):
331    """ Write data as numpy file
332
333    :param filename: name of numpy file to write data to
334    :type filename: str
335    :param data: data to write to numpy file
336    :type data: numpy array
337    """
338    return np.save(filename, data)
339
340
341def write_bytes(filename, data):
342    """ Write binary data into a file
343
344    :param filename:
345    :param data:
346    :return:
347    """
348    with open(filename, 'wb') as file:
349        file.write(data)
350