|
44 | 44 | "metadata": {}, |
45 | 45 | "outputs": [], |
46 | 46 | "source": [ |
47 | | - "from mdio import segy_to_mdio, MDIOReader" |
| 47 | + "from mdio import MDIOReader\n", |
| 48 | + "from mdio import segy_to_mdio" |
48 | 49 | ] |
49 | 50 | }, |
50 | 51 | { |
|
263 | 264 | "ascii": false, |
264 | 265 | "bar_format": null, |
265 | 266 | "colour": null, |
266 | | - "elapsed": 0.014499902725219727, |
| 267 | + "elapsed": 0.014499902725219728, |
267 | 268 | "initial": 0, |
268 | 269 | "n": 0, |
269 | 270 | "ncols": null, |
|
366 | 367 | } |
367 | 368 | ], |
368 | 369 | "source": [ |
369 | | - "import os\n", |
| 370 | + "from pathlib import Path\n", |
370 | 371 | "\n", |
371 | 372 | "\n", |
372 | | - "def get_dir_size(path: str) -> int:\n", |
| 373 | + "def get_dir_size(path: Path) -> int:\n", |
373 | 374 | " \"\"\"Get size of a directory recursively.\"\"\"\n", |
374 | 375 | " total = 0\n", |
375 | | - " with os.scandir(path) as it:\n", |
376 | | - " for entry in it:\n", |
377 | | - " if entry.is_file():\n", |
378 | | - " total += entry.stat().st_size\n", |
379 | | - " elif entry.is_dir():\n", |
380 | | - " total += get_dir_size(entry.path)\n", |
| 376 | + " for entry in path.iterdir():\n", |
| 377 | + " if entry.is_file():\n", |
| 378 | + " total += entry.stat().st_size\n", |
| 379 | + " elif entry.is_dir():\n", |
| 380 | + " total += get_dir_size(entry)\n", |
381 | 381 | " return total\n", |
382 | 382 | "\n", |
383 | 383 | "\n", |
384 | | - "def get_size(path: str) -> int:\n", |
| 384 | + "def get_size(path: Path) -> int:\n", |
385 | 385 | " \"\"\"Get size of a folder or a file.\"\"\"\n", |
386 | | - " if os.path.isfile(path):\n", |
387 | | - " return os.path.getsize(path)\n", |
388 | | - "\n", |
389 | | - " elif os.path.isdir(path):\n", |
| 386 | + " if path.is_file():\n", |
| 387 | + " return path.stat().st_size\n", |
| 388 | + " if path.is_dir():\n", |
390 | 389 | " return get_dir_size(path)\n", |
| 390 | + " return 0 # Handle case where path is neither file nor directory\n", |
391 | 391 | "\n", |
392 | 392 | "\n", |
393 | 393 | "print(f\"SEG-Y:\\t{get_size('volve.segy') / 1024 / 1024:.2f} MB\")\n", |
|
466 | 466 | ], |
467 | 467 | "source": [ |
468 | 468 | "import matplotlib.pyplot as plt\n", |
| 469 | + "from matplotlib.axes import Axes\n", |
| 470 | + "from matplotlib.image import AxesImage\n", |
469 | 471 | "from mpl_toolkits.axes_grid1 import make_axes_locatable\n", |
| 472 | + "from numpy.typing import NDArray\n", |
470 | 473 | "\n", |
471 | 474 | "vmin = -3 * std\n", |
472 | 475 | "vmax = 3 * std\n", |
473 | 476 | "\n", |
474 | | - "imshow_kw = dict(vmin=vmin, vmax=vmax, cmap=\"gray_r\", interpolation=\"bilinear\")\n", |
| 477 | + "imshow_kw = {\"vmin\": -3 * std, \"vmax\": 3 * std, \"cmap\": \"gray_r\", \"interpolation\": \"bilinear\"}\n", |
475 | 478 | "\n", |
476 | 479 | "\n", |
477 | | - "def attach_colorbar(image, axis):\n", |
| 480 | + "def attach_colorbar(image: AxesImage, axis: Axes) -> None:\n", |
| 481 | + " \"\"\"Attach a colorbar to an axis.\"\"\"\n", |
478 | 482 | " divider = make_axes_locatable(axis)\n", |
479 | 483 | " cax = divider.append_axes(\"top\", size=\"2%\", pad=0.05)\n", |
480 | 484 | " plt.colorbar(image, cax=cax, orientation=\"horizontal\")\n", |
481 | 485 | " cax.xaxis.set_ticks_position(\"top\")\n", |
482 | 486 | " cax.tick_params(labelsize=8)\n", |
483 | 487 | "\n", |
484 | 488 | "\n", |
485 | | - "def plot_image_and_cbar(data, axis, title):\n", |
| 489 | + "def plot_image_and_cbar(data: NDArray, axis: Axes, title: str) -> None:\n", |
| 490 | + " \"\"\"Plot an image with a colorbar.\"\"\"\n", |
486 | 491 | " image = axis.imshow(data.T, **imshow_kw)\n", |
487 | 492 | " attach_colorbar(image, axis)\n", |
488 | 493 | " axis.set_title(title, y=-0.15)\n", |
489 | 494 | "\n", |
490 | 495 | "\n", |
491 | | - "def plot_inlines_with_diff(orig, compressed, title):\n", |
| 496 | + "def plot_inlines_with_diff(orig: NDArray, compressed: NDArray, title: str) -> None:\n", |
| 497 | + " \"\"\"Plot lossless and lossy inline with their differences.\"\"\"\n", |
492 | 498 | " fig, ax = plt.subplots(1, 4, sharey=\"all\", sharex=\"all\", figsize=(12, 5))\n", |
493 | 499 | "\n", |
494 | 500 | " diff = orig[200] - compressed[200]\n", |
|
545 | 551 | } |
546 | 552 | ], |
547 | 553 | "source": [ |
548 | | - "import skimage\n", |
549 | | - "\n", |
550 | | - "\n", |
551 | | - "def get_metrics(image_true, image_test):\n", |
552 | | - " \"\"\"Get four metrics\"\"\"\n", |
553 | | - " psnr = skimage.metrics.peak_signal_noise_ratio(\n", |
554 | | - " image_true[200], image_test[200], data_range=max_ - min_\n", |
555 | | - " )\n", |
556 | | - " ssim = skimage.metrics.structural_similarity(\n", |
557 | | - " image_true[200], image_test[200], data_range=max_ - min_\n", |
558 | | - " )\n", |
559 | | - " mse = skimage.metrics.mean_squared_error(image_true[200], image_test[200])\n", |
560 | | - " nrmse = skimage.metrics.normalized_root_mse(image_true[200], image_test[200])\n", |
| 554 | + "from numpy.typing import NDArray\n", |
| 555 | + "from skimage.metrics import mean_squared_error\n", |
| 556 | + "from skimage.metrics import normalized_root_mse\n", |
| 557 | + "from skimage.metrics import peak_signal_noise_ratio\n", |
| 558 | + "from skimage.metrics import structural_similarity\n", |
| 559 | + "\n", |
| 560 | + "\n", |
| 561 | + "def get_metrics(image_true: NDArray, image_test: NDArray) -> tuple[float, float, float, float]:\n", |
| 562 | + " \"\"\"Get four image quality metrics.\"\"\"\n", |
| 563 | + " psnr = peak_signal_noise_ratio(image_true[200], image_test[200], data_range=max_ - min_)\n", |
| 564 | + " ssim = structural_similarity(image_true[200], image_test[200], data_range=max_ - min_)\n", |
| 565 | + " mse = mean_squared_error(image_true[200], image_test[200])\n", |
| 566 | + " nrmse = normalized_root_mse(image_true[200], image_test[200])\n", |
561 | 567 | "\n", |
562 | 568 | " return psnr, ssim, mse, nrmse\n", |
563 | 569 | "\n", |
|
0 commit comments