code_public/lidar_rendu
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add --gpu flag to select specific GPU(s) for processing

Browse Source
This commit is contained in:
Antoine Jacquin
2026-05-31 16:06:58 +02:00
parent 266214fe3e
commit b5b6787956
3 changed files with 141 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
lidar_pipeline/pipeline.py
View File

@ -63,7 +63,7 @@ from .visualizations import (
generate_roughness, generate_wavelet,
generate_svf, generate_aniso_open, generate_paths,
)
from .gpu import gpu_cleanup, num_gpus, safe_gpu_call
from .gpu import gpu_cleanup, num_gpus, restrict_gpus, safe_gpu_call
from .ign import generate_ign_overlay
from .rendering import tif_to_png
@ -107,7 +107,7 @@ VIZ_STEPS = [
class LidarArchaeoPipeline:
"""Orchestrates the LiDAR archaeological analysis pipeline."""
def __init__(self, input_dir, output_dir, resolution=0.5, workers=1, force=False, ground_method='auto', force_classify=False, keep_tif=False, quality=98, only_viz=None, skip_viz=None, output_format='avif'):
def __init__(self, input_dir, output_dir, resolution=0.5, workers=1, force=False, ground_method='auto', force_classify=False, keep_tif=False, quality=98, only_viz=None, skip_viz=None, output_format='avif', gpu_ids=None):
self.input_dir = Path(input_dir)
self.output_dir = Path(output_dir)
# Accept single float or comma-separated string for multi-resolution
@ -127,6 +127,7 @@ class LidarArchaeoPipeline:
self.only_viz = only_viz
self.skip_viz = skip_viz
self.output_format = output_format
self.gpu_ids = gpu_ids
self.temp_dir = self.output_dir / "temp"
if not self.input_dir.exists():
@ -448,6 +449,10 @@ class LidarArchaeoPipeline:
results = {}
t_pipeline_start = time.time()
# Restrict visible GPUs in the main process before spawning workers
if self.gpu_ids is not None:
restrict_gpus(self.gpu_ids)
if self.workers > 1 and len(files) > 1:
n_gpus = num_gpus() or 1
if n_gpus > 1:
@ -462,7 +467,7 @@ class LidarArchaeoPipeline:
# Pass resolutions as comma-separated string for multiprocessing serialization
resolutions_str = ','.join(str(r) for r in self.resolutions)
future_to_file = {
executor.submit(_process_file_standalone, str(laz_file), str(self.input_dir), str(self.output_dir), resolutions_str, self.force, self.ground_method, self.force_classify, self.keep_tif, self.quality, self.only_viz, self.skip_viz, self.output_format, gpu_id % n_gpus): laz_file
executor.submit(_process_file_standalone, str(laz_file), str(self.input_dir), str(self.output_dir), resolutions_str, self.force, self.ground_method, self.force_classify, self.keep_tif, self.quality, self.only_viz, self.skip_viz, self.output_format, gpu_id % n_gpus, gpu_ids=self.gpu_ids): laz_file
for gpu_id, laz_file in enumerate(files)
}
done = 0
@ -531,14 +536,18 @@ class LidarArchaeoPipeline:
logger.warning(f" Note: Impossible de supprimer les fichiers temporaires: {e}")
def _process_file_standalone(laz_file_str, input_dir, output_dir, resolution, force=False, ground_method='auto', force_classify=False, keep_tif=False, quality=98, only_viz=None, skip_viz=None, output_format='avif', gpu_id=None):
def _process_file_standalone(laz_file_str, input_dir, output_dir, resolution, force=False, ground_method='auto', force_classify=False, keep_tif=False, quality=98, only_viz=None, skip_viz=None, output_format='avif', gpu_id=None, gpu_ids=None):
"""Standalone function for multiprocessing — creates its own pipeline instance.
Each worker gets its own temp directory to avoid file conflicts.
When multiple GPUs are available, each worker is assigned a GPU via
CuPy's Device API to balance load across GPUs.
CUDA_VISIBLE_DEVICES to balance load across GPUs.
"""
# Assign GPU to this worker using CuPy's Device API
# Restrict visible GPUs first, then pick one for this worker
if gpu_ids is not None:
from .gpu import restrict_gpus
restrict_gpus(gpu_ids)
if gpu_id is not None and gpu_id >= 0:
from .gpu import set_active_gpu
set_active_gpu(gpu_id)