Picks refinement without bench + Python 3.12

.github/workflows/pre-commit-workflow.yaml

@@ -12,18 +12,17 @@ on:
 
 jobs:
   pre_commit:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-24.04
     steps:
     - name: Checkout
-      uses: actions/checkout@v3
+      uses: actions/checkout@v4
     - name: Set up Python
-      uses: actions/setup-python@v3
+      uses: actions/setup-python@v5
       with:
-        python-version: "3.9.0"
+        python-version: "3.12.0"
     - name: Install dependencies
       run: |
-        apt-get get update && apt-get install cmake
-        make install_precommit
+        pip install pre-commit
     - name: Pre-commit tests
       run: |
-        make run_precommit
+        pre-commit run -a

.github/workflows/tests-workflow.yaml

@@ -1,24 +1,39 @@
 name: Tests
 
-
 on:
   pull_request:
-    branches:
-      - main
+    branches: [main]
   push:
-    branches:
-      - main
-
+    branches: [main]
 
 jobs:
   tests:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-24.04
+    env:
+      QT_QPA_PLATFORM: offscreen
+      LIBGL_ALWAYS_INDIRECT: "1"
     steps:
-    - name: Checkout
-      uses: actions/checkout@v3
-    - name: Install package
-      run: |
-        make docker_build
-    - name: Tests
-      run: |
-        make docker_tests
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: System deps for PyQt5 (xcb/offscreen)
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y \
+            cmake \
+            libgl1 libglib2.0-0 libsm6 libxext6 libxrender1 \
+            libx11-6 libxcb1 libxkbcommon0 libxkbcommon-x11-0 \
+            libfontconfig1 libfreetype6
+
+      - name: Install project
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e .
+
+      - name: Tests
+        run: make run_tests

.github/workflows/tomls.yaml

@@ -12,14 +12,14 @@ on:
 
 jobs:
   tests:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-24.04
     steps:
     - name: Checkout
-      uses: actions/checkout@v3
+      uses: actions/checkout@v4
     - name: Set up Python
-      uses: actions/setup-python@v3
+      uses: actions/setup-python@v5
       with:
-        python-version: "3.8.0"
+        python-version: "3.12.0"
     - name: Install lib
       run: |
         pip install tomli

.gitignore

@@ -143,3 +143,4 @@ dmypy.json
 fb.onnx
 
 local_dev
+/debug/

.pre-commit-config.yaml

@@ -1,12 +1,10 @@
 default_language_version:
-  python: python3.9
+  python: python3.12
 files: \.py$|tests/.*\.py$
-default_stages:
-- commit
 repos:
 # general hooks to verify or beautify code
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.6.0
+  rev: v6.0.0
   hooks:
   - id: check-added-large-files
     args: [--maxkb=5000]
@@ -23,7 +21,7 @@ repos:
 
 # autoformat code with black formatter
 - repo: https://github.com/psf/black
-  rev: 22.3.0
+  rev: 25.9.0
   hooks:
   - id: black
     files: first_breaks|tests
@@ -32,7 +30,7 @@ repos:
 
 # beautify and sort imports
 - repo: https://github.com/pycqa/isort
-  rev: 5.12.0
+  rev: 6.1.0
   hooks:
   - id: isort
     files: first_breaks|tests
@@ -41,7 +39,7 @@ repos:
 
 # check code style
 - repo: https://github.com/pycqa/flake8
-  rev: 3.8.4
+  rev: 7.3.0
   hooks:
   - id: flake8
     files: first_breaks|tests
@@ -50,7 +48,7 @@ repos:
 
 # static type checking
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v1.5.1
+  rev: v1.18.2
   hooks:
   - id: mypy
     files: first_breaks|tests

Makefile

@@ -1,28 +1,11 @@
-IMAGE_NAME ?= first-breaks-picking:latest
-
-
-.PHONY: install_precommit
-install_precommit:
-	python -m pip install --upgrade pre-commit==3.5.0
-
 .PHONY: run_precommit
 run_precommit:
-	pre-commit install && pre-commit run -a
+	pre-commit run -a
 
 .PHONY: run_tests
 run_tests:
 	pytest -sv --disable-warnings tests
 
-
-.PHONY: docker_build
-docker_build:
-	DOCKER_BUILDKIT=1 docker build -t $(IMAGE_NAME) .
-
-.PHONY: docker_tests
-docker_tests: docker_build
-	docker run -t $(IMAGE_NAME) make run_tests
-
-
 .PHONY: build_wheel
 build_wheel:
 	python -m pip install --upgrade pip

README.md

@@ -476,7 +476,7 @@ picks_ms = np.random.uniform(low=0,
                              size=sgy.ntr)
 picks = Picks(values=picks_ms, unit="ms", dt_mcs=sgy.dt_mcs, color=(0, 100, 100))
 export_image(sgy, image_filename,
-             picks=picks)
+             picks_list=[picks])
 ```
 [code-block-end]:plot-sgy-custom-picks
 

first_breaks/_pytorch/picker_torch.py

@@ -80,7 +80,7 @@ def change_settings(  # type: ignore
         device: Optional[str] = None,
         segmentation_hw: Optional[Tuple[int, int]] = None,
         num_workers: Optional[int] = None,
-        batch_size: Optional[int] = None
+        batch_size: Optional[int] = None,
     ) -> None:
         if args:
             raise ValueError("Use named arguments instead of positional")

first_breaks/const.py

@@ -28,6 +28,10 @@ def get_cache_folder() -> Path:
 MODEL_ONNX_PATH = CACHE_FOLDER / "fb.onnx"
 MODEL_ONNX_URL = "https://oml.daloroserver.com/download/seis/fb.onnx"
 MODEL_ONNX_HASH = "7e39e017b01325180e36885eccaeb17a"
+MODEL_ONNX_HASHES = [
+    MODEL_ONNX_HASH,
+    "afc03594f49b88ea61b5cf6ba8245be4",  # model with heatmap
+]
 
 TIMEOUT = 60
 

first_breaks/data_models/independent.py

@@ -5,7 +5,7 @@
 import numpy as np
 from pydantic import UUID4, BaseModel, Field, field_validator
 
-TColor = Union[Tuple[int, int, int, int], Tuple[int, int, int]]
+TColor = Union[Tuple[int, int, int], Tuple[int, int, int, int], Tuple[int, ...]]
 TNormalize = Union[Literal["trace", "gather"], float, int, np.ndarray, None]
 
 

first_breaks/desktop/combobox_with_mapping.py

@@ -30,7 +30,7 @@ def __init__(
         current_label: Optional[str] = None,
         current_value: Optional[str] = None,
         *args: Any,
-        **kwargs: Any
+        **kwargs: Any,
     ):
         super().__init__(*args, **kwargs)
 

first_breaks/desktop/graph.py

@@ -2,7 +2,7 @@
 import os
 import warnings
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
 import pyqtgraph as pg
@@ -423,7 +423,7 @@ def export(
         fill_black: Optional[str] = DEFAULTS.fill_black,
         time_window: Optional[Tuple[float, float]] = None,
         traces_window: Optional[Tuple[float, float]] = None,
-        picks: Optional[Picks] = None,
+        picks_list: Optional[Sequence[Picks]] = None,
         task: Optional[Task] = None,
         show_processing_region: bool = True,
         contour_color: TColor = DEFAULTS.region_contour_color,
@@ -443,9 +443,6 @@ def export(
         if args:
             raise need_kwargs_exception
 
-        if picks is not None and task is not None:
-            raise ValueError("'picks' and 'task' are mutually exclusive. Use only one of them or none")
-
         if width is None:
             if traces_window is None:
                 num_traces = sgy.num_traces
@@ -467,10 +464,12 @@ def export(
             refresh_view=True,
         )
 
-        if task:
+        if task and task.picks is not None:
             self.plot_picks(task.picks)
-        elif picks is not None:
-            self.plot_picks(picks)
+        if picks_list:
+            assert all(len(picks) == sgy.num_traces for picks in picks_list)
+            for picks in picks_list:
+                self.plot_picks(picks)
 
         if task is not None and show_processing_region:
             self.plot_processing_region(
@@ -537,7 +536,7 @@ def export_image(
     fill_black: Optional[str] = DEFAULTS.fill_black,
     time_window: Optional[Tuple[float, float]] = None,
     traces_window: Optional[Tuple[float, float]] = None,
-    picks: Optional[Picks] = None,
+    picks_list: Optional[Sequence[Picks]] = None,
     show_processing_region: bool = True,
     contour_color: TColor = DEFAULTS.region_contour_color,
     poly_color: TColor = DEFAULTS.region_poly_color,
@@ -586,7 +585,7 @@ def export_image(
         fill_black=fill_black,
         time_window=time_window,
         traces_window=traces_window,
-        picks=picks,
+        picks_list=picks_list,
         task=task,
         show_processing_region=show_processing_region,
         contour_color=contour_color,

first_breaks/desktop/main_gui.py

@@ -1,7 +1,7 @@
 import sys
 import warnings
 from pathlib import Path
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from PyQt5.QtCore import QSize, Qt, QThreadPool, pyqtSignal
 from PyQt5.QtGui import QCloseEvent
@@ -19,7 +19,12 @@
     QWidget,
 )
 
-from first_breaks.const import DEMO_SGY_PATH, HIGH_DPI, MODEL_ONNX_HASH, MODEL_ONNX_PATH
+from first_breaks.const import (
+    DEMO_SGY_PATH,
+    HIGH_DPI,
+    MODEL_ONNX_HASHES,
+    MODEL_ONNX_PATH,
+)
 from first_breaks.data_models.independent import ExceptionOptional
 from first_breaks.desktop.graph import GraphWidget
 from first_breaks.desktop.last_folder_manager import last_folder_manager
@@ -48,11 +53,11 @@ class FileState:
     file_changed = 2
 
     @classmethod
-    def get_file_state(cls, fname: Union[str, Path], fhash: str) -> int:
+    def get_file_state(cls, fname: Union[str, Path], fhashes: List[str]) -> int:
         if not Path(fname).is_file():
             return cls.file_not_exists
         else:
-            return cls.valid_file if calc_hash(fname) == fhash else cls.file_changed
+            return cls.valid_file if calc_hash(fname) in fhashes else cls.file_changed
 
 
 class ReadyToProcess:
@@ -195,7 +200,7 @@ def __init__(self, use_open_gl: bool = True, show: bool = True):  # type: ignore
         self.settings: Optional[Dict[str, Any]] = None
         self.last_folder: Optional[Union[str, Path]] = None
         self.picks_from_file_in_ms: Optional[Tuple[Union[int, float], ...]] = None
-        self.picker_hash = MODEL_ONNX_HASH
+        self.picker_hashes = MODEL_ONNX_HASHES
 
         if show:
             self.show()
@@ -255,7 +260,10 @@ def run_processing_region(self) -> None:
     def show_processing_region(self) -> None:
         for picks in self.picks_manager.picks_mapping.values():
             if picks.created_by_nn and picks.active:
-                tps, max_time = picks.picking_parameters.traces_per_gather, picks.picking_parameters.maximum_time
+                tps, max_time = (
+                    picks.picking_parameters.traces_per_gather,
+                    picks.picking_parameters.maximum_time,
+                )
                 self.graph.plot_processing_region(tps, max_time)
                 break
 
@@ -290,11 +298,14 @@ def load_nn(self, filename: Optional[Union[str, Path]] = None) -> None:
         if not filename:
             options = QFileDialog.Options()
             filename, _ = QFileDialog.getOpenFileName(
-                self, "Select file with NN weights", directory=last_folder_manager.get_last_folder(), options=options
+                self,
+                "Select file with NN weights",
+                directory=last_folder_manager.get_last_folder(),
+                options=options,
             )
 
         if filename:
-            if FileState.get_file_state(filename, self.picker_hash) == FileState.valid_file:
+            if FileState.get_file_state(filename, self.picker_hashes) == FileState.valid_file:
                 self.nn_manager.init_net(weights=filename)
                 self.button_load_nn.setEnabled(False)
                 self.ready_to_process.model_loaded = True