forked from agilescientific/seisplot
-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
439 lines (373 loc) · 14.3 KB
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
# -*- coding: utf-8 -*-
"""
Utility functions for seisplot.
:copyright: 2016-22 Agile Scientific
:license: Apache 2.0
"""
import os
import errno
import sys
import re
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
import utils
LABELS = {'i': 'inline', 'x': 'xline', 't': 'time [ms]'}
DEFAULTS = {'ndim': 2,
'line': 'inline',
'direction': 'inline',
'number': 0.5,
'sidelabel': 'right',
'tpi': 10,
'plot_width': 17,
'ips': 1,
'plot_height': 11,
'trange': [0, 0],
'skip': 2,
'display': 'vd',
'title': '_filename',
'subtitle': '_date',
'credit': True,
'gain': 1.0,
'percentile': 99.0,
'colour': [0, 0, 0],
'opacity': 1.0,
'lineweight': 0.2,
'grid_time': False,
'grid_traces': False,
'grid_colour': [0, 0, 0],
'grid_alpha': 0.15,
'cmap': 'Greys',
'interpolation': 'bicubic',
'highlight_colour': [0, 0, 0],
'fontsize': 10,
'watermark_text': '', # None by default
'watermark_size': 14,
'watermark_family': 'sans-serif',
'watermark_style': 'normal',
'watermark_weight': 'normal',
'watermark_colour': 'white',
'watermark_alpha': 0.5,
'watermark_angle': 33,
'watermark_cols': 6,
'watermark_rows': 0, # automatic
'stain_paper': None,
'coffee_rings': 0,
'distort': False,
'scribble': False,
'segy_library': 'segyio',
}
HEADERS = [
'for_3d_poststack_data_this_field_is_for_cross_line_number', # sawtooth for 3d
'for_3d_poststack_data_this_field_is_for_in_line_number', # steps for 3d
'trace_number_within_the_ensemble', # nothing for 2d, same as crossline number for 3d
'trace_sequence_number_within_line', # nothing for 2d, zero-based sawtooth for 3d
'trace_sequence_number_within_segy_file', # zero-based monotonic for 2d and 3d
'ensemble_number', # trace-number-based monotonic for 2d
'original_field_record_number',
'energy_source_point_number',
'trace_number_within_the_original_field_record',
# 'trace_identification_code',
]
def listdir(directory, match=None):
"""
Wrapper for `os.listdir()` that returns full paths. A bit like
`utils.walk()` but not recursive. Case insensitive.
Args:
directory (str): The directory to list.
Yields:
str: Full path to each file in turn.
"""
for f in os.listdir(directory):
if match:
if not re.search(match, f, flags=re.IGNORECASE):
continue
yield os.path.join(directory, f)
def rgb_to_hex(rgb):
"""
Utility function to convert (r,g,b) triples to hex.
http:https://ageo.co/1CFxXpO
Args:
rgb (tuple): A sequence of RGB values in the
range 0-255 or 0-1.
Returns:
str: The hex code for the colour.
"""
r, g, b = rgb[:3]
if (r < 0) or (g < 0) or (b < 0):
raise Exception("RGB values must all be 0-255 or 0-1")
if (r > 255) or (g > 255) or (b > 255):
raise Exception("RGB values must all be 0-255 or 0-1")
if (0 < r < 1) or (0 < g < 1) or (0 < b < 1):
if (r > 1) or (g > 1) or (b > 1):
raise Exception("RGB values must all be 0-255 or 0-1")
if (0 <= r <= 1) and (0 <= g <= 1) and (0 <= b <= 1):
rgb = tuple([int(round(val * 255)) for val in [r, g, b]])
else:
rgb = (int(r), int(g), int(b))
result = '#%02x%02x%02x' % rgb
return result.lower()
def path_bits(path):
"""
Returns the directory, the stem of the filename, and the extension.
"""
d, f = os.path.split(path)
s, e = os.path.splitext(f)
return d, s, e
def get_pattern_from_stream(stream, pattern_function):
"""
Return the first non-zero-based monotonic header.
If there isn't one, return the first monotonic header.
"""
candidates = []
for header in HEADERS:
data = [t.header.__dict__[header] for t in stream.traces]
if pattern_function(data):
candidates.append(data)
for candidate in candidates:
if candidate[0] > 0:
return candidate
if candidates:
return candidates[0]
else:
return None
def get_trace_indices(y, ntraces, random=False):
"""
Get a subset of trace positions.
Args:
y (int): The total number of traces to choose from. 1D or 2D.
ntraces (int): The number of traces to choose.
random (bool): Whether to choose random traces, or
to choose regularly-spaced.
Yields:
ndarray. 1D array of ints or tuples, with the coords of the
random traces.
"""
total = np.product(y)
if random:
x = 0.05 + 0.9*np.random.random(ntraces) # avoids edges
ti = np.sort(x * total)
else:
n = ntraces + 1
ti = np.arange(1./n, 1., 1./n) * total
if len(y) > 1:
# Re-form the 2D indices.
ti = [(t%y[0], t//y[0]) for t in ti]
return np.floor(ti).astype(int)
def max_opacity(image, maxi):
"""
Adjust the maximum opacity of an image.
"""
data = np.array(image)
adj = maxi*255/np.amax(data)
data[..., 3] = adj * data[..., 3]
result = Image.fromarray(data)
return result
def stain_paper(image):
"""
Add a staining image to the paper.
"""
fname = "resources/stained_and_folded_paper_.png"
paper = Image.open(fname)
# Adjust opacity.
paper2 = max_opacity(paper, 0.25)
# Crop paper to image size
# If paper is bigger than image.
left = int((paper2.size[0] - image.size[0]) / 2)
right = left + image.size[0]
upper = int((paper2.size[1] - image.size[1]) / 2)
lower = upper + image.size[1]
box = (left, upper, right, lower)
paper2 = paper2.crop(box=box)
# Do it.
image.paste(paper2, (0, 0), paper2)
return
def add_scribble(image):
"""
Add a scribble.
"""
files = list(utils.listdir('resources', 'scribble_[0-9]+.png'))
fname = np.random.choice(files)
scribble = Image.open(fname)
# Adjust opacity.
scribble2 = max_opacity(scribble, 0.85)
# Shrink a bit
scribble2.thumbnail((256, 256), Image.ANTIALIAS)
# Rotate the stain by a small random amount.
angle = (np.random.random() - 0.5) * 20
scribble2 = scribble2.rotate(angle, resample=Image.BICUBIC, expand=True)
# Find a random place near the upper corner.
x = np.random.randint(image.size[0]-(1.5*scribble2.size[0]), image.size[0]-(0.8*scribble2.size[0]))
y = np.random.randint(0, scribble2.size[1])
# Do it.
image.paste(scribble2, (x, y), scribble2)
return
def add_a_ring(image):
"""
Add a coffee cup ring.
"""
fname = "resources/coffee_stain.png"
stain = Image.open(fname)
# Adjust opacity.
stain2 = max_opacity(stain, 0.35)
# Rotate the stain by a random amount.
angle = np.random.random() * 360
stain2 = stain2.rotate(angle, resample=Image.BICUBIC, expand=True)
# Find a random place for it.
x = np.random.randint(-stain2.size[0]/2, image.size[0]-stain2.size[0]/2)
y = np.random.randint(-stain2.size[1]/2, image.size[1]-stain2.size[1]/2)
# Do it.
image.paste(stain2, (x, y), stain2)
return
def add_rings(image, n=0):
"""
Add the required number of coffee rings.
"""
if n:
for i in range(n):
add_a_ring(image)
return
def add_subplot_axes(ax, rect, axisbg='w'):
"""
Facilitates the addition of a small subplot within another plot.
From: http:https://stackoverflow.com/questions/17458580/
embedding-small-plots-inside-subplots-in-matplotlib
License: CC-BY-SA
Args:
ax (axis): A matplotlib axis.
rect (list): A rect specifying [left pos, bot pos, width, height]
Returns:
axis: The sub-axis in the specified position.
"""
def axis_to_fig(axis):
fig = axis.figure
def transform(coord):
a = axis.transAxes.transform(coord)
return fig.transFigure.inverted().transform(a)
return transform
fig = plt.gcf()
left, bottom, width, height = rect
trans = axis_to_fig(ax)
x1, y1 = trans((left, bottom))
x2, y2 = trans((left + width, bottom + height))
subax = fig.add_axes([x1, y1, x2 - x1, y2 - y1])
x_labelsize = subax.get_xticklabels()[0].get_size()
y_labelsize = subax.get_yticklabels()[0].get_size()
x_labelsize *= rect[2] ** 0.5
y_labelsize *= rect[3] ** 0.5
subax.xaxis.set_tick_params(labelsize=x_labelsize)
subax.yaxis.set_tick_params(labelsize=y_labelsize)
return subax
def make_patch_spines_invisible(ax):
"""
Removes spines from patches.
"""
ax.set_frame_on(True)
ax.patch.set_visible(False)
for sp in ax.spines.values():
sp.set_visible(False)
return ax
# Sadly, Python fails to provide the following magic number for us.
ERROR_INVALID_NAME = 123
'''
Windows-specific error code indicating an invalid pathname.
See Also
----------
https://msdn.microsoft.com/en-us/library/windows/desktop/ms681382%28v=vs.85%29.aspx
Official listing of all such codes.
'''
def is_pathname_valid(pathname: str) -> bool:
'''
http:https://stackoverflow.com/questions/9532499/check-whether-a-path-is-valid-in-python-without-creating-a-file-at-the-paths-ta
by Cecil Curry, mindshines.com
`True` if the passed pathname is a valid pathname for the current OS;
`False` otherwise.
'''
# If this pathname is either not a string or is but is empty, this pathname
# is invalid.
try:
if not isinstance(pathname, str) or not pathname:
return False
# Strip this pathname's Windows-specific drive specifier (e.g., `C:\`)
# if any. Since Windows prohibits path components from containing `:`
# characters, failing to strip this `:`-suffixed prefix would
# erroneously invalidate all valid absolute Windows pathnames.
_, pathname = os.path.splitdrive(pathname)
# Directory guaranteed to exist. If the current OS is Windows, this is
# the drive to which Windows was installed (e.g., the "%HOMEDRIVE%"
# environment variable); else, the typical root directory.
root_dirname = os.environ.get('HOMEDRIVE', 'C:') \
if sys.platform == 'win32' else os.path.sep
assert os.path.isdir(root_dirname) # ...Murphy and her ironclad Law
# Append a path separator to this directory if needed.
root_dirname = root_dirname.rstrip(os.path.sep) + os.path.sep
# Test whether each path component split from this pathname is valid or
# not, ignoring non-existent and non-readable path components.
for pathname_part in pathname.split(os.path.sep):
try:
os.lstat(root_dirname + pathname_part)
# If an OS-specific exception is raised, its error code
# indicates whether this pathname is valid or not. Unless this
# is the case, this exception implies an ignorable kernel or
# filesystem complaint (e.g., path not found or inaccessible).
#
# Only the following exceptions indicate invalid pathnames:
#
# * Instances of the Windows-specific "WindowsError" class
# defining the "winerror" attribute whose value is
# "ERROR_INVALID_NAME". Under Windows, "winerror" is more
# fine-grained and hence useful than the generic "errno"
# attribute. When a too-long pathname is passed, for example,
# "errno" is "ENOENT" (i.e., no such file or directory) rather
# than "ENAMETOOLONG" (i.e., file name too long).
# * Instances of the cross-platform "OSError" class defining the
# generic "errno" attribute whose value is either:
# * Under most POSIX-compatible OSes, "ENAMETOOLONG".
# * Under some edge-case OSes (e.g., SunOS, *BSD), "ERANGE".
except OSError as exc:
if hasattr(exc, 'winerror'):
if exc.winerror == ERROR_INVALID_NAME:
return False
elif exc.errno in {errno.ENAMETOOLONG, errno.ERANGE}:
return False
# If a "TypeError" exception was raised, it almost certainly has the
# error message "embedded NUL character" indicating an invalid pathname.
except TypeError as exc:
return False
# If no exception was raised, all path components and hence this
# pathname itself are valid. (Praise be to the curmudgeonly python.)
else:
return True
# If any other exception was raised, this is an unrelated fatal issue
# (e.g., a bug). Permit this exception to unwind the call stack.
#
# Did we mention this should be shipped with Python already?
def is_path_creatable(pathname: str) -> bool:
'''
`True` if the current user has sufficient permissions to create the passed
pathname; `False` otherwise.
'''
# Parent directory of the passed path. If empty, we substitute the current
# working directory (CWD) instead.
dirname = os.path.dirname(pathname) or os.getcwd()
return os.access(dirname, os.W_OK)
def is_path_exists_or_creatable(pathname: str) -> bool:
'''
`True` if the passed pathname is a valid pathname for the current OS _and_
either currently exists or is hypothetically creatable; `False` otherwise.
This function is guaranteed to _never_ raise exceptions.
'''
try:
# To prevent "os" module calls from raising undesirable exceptions on
# invalid pathnames, is_pathname_valid() is explicitly called first.
return is_pathname_valid(pathname) and (
os.path.exists(pathname) or is_path_creatable(pathname))
# Report failure on non-fatal filesystem complaints (e.g., connection
# timeouts, permissions issues) implying this path to be inaccessible. All
# other exceptions are unrelated fatal issues and should not be caught.
except OSError:
return False
def chunks(s: str, n: int) -> str:
"""Produce `n`-character chunks from string `s`."""
for start in range(0, len(s), n):
yield s[start:start + n]