Coverage for collectra/utils.py: 46.04%
202 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 09:56 +0000
1"""Utility functions and helpers for the Collectra workflow system.
3This module provides common utility functions used throughout the Collectra
4application, including file operations, message formatting, image processing,
5and configuration management.
7The module includes functions for:
8 - Formatted console output messages (success, error, processing)
9 - File collection and filtering operations
10 - Configuration loading from directories and zip files
11 - Image processing and cropping operations
12 - Dynamic class path resolution
14Functions:
15 success_msg: Format success messages for console output
16 error_msg: Format error messages for console output
17 processing_msg: Format processing status messages
18 get_class_path: Get fully qualified class path from object
19 get_all_files: Collect files matching format from paths
20 unzip: Extract configuration from zip files
21 from_dir: Load configuration from directory
22 load_class_from_string: Dynamically load class from string
23 crop: Crop images using specified coordinates
24"""
26import importlib
27import os
28import shutil
29import zipfile
30from contextlib import contextmanager
31from functools import wraps
32from pathlib import Path
33from threading import Lock
34from typing import List
36import yaml
37from rich import print
38from rich.progress import track
39from yaml.emitter import Emitter, ScalarAnalysis
41from .logger import get_logger
43logger = get_logger(__name__)
46def threading_locked():
47 """Return a decorator that serializes calls to the decorated callable."""
48 lock = Lock()
50 def decorate(func):
51 @wraps(func)
52 def wrapped(*args, **kwargs):
53 with lock:
54 return func(*args, **kwargs)
56 return wrapped
58 return decorate
61def valid_raw_files() -> list[str]:
62 """Return a list of valid raw file extensions for processing."""
63 return [".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".gif", ".webp"]
66def resolve_files(inputs: list[Path], ext: list[str]) -> list[Path]:
67 """Resolve a list of file paths from the given inputs.
69 This function takes a list of Path objects, which can be either files or directories.
70 It recursively collects all files from the directories and returns a flat list of file paths.
72 Args:
73 inputs (list[Path]): List of Path objects representing files or directories.
75 Returns:
76 list[Path]: A flat list of Path objects representing all files found in the input paths.
77 """
78 files: list[Path] = []
79 for path in inputs:
80 path = Path(path)
81 if not path.exists():
82 logger.warning(f"Path does not exist: {path}")
83 continue
84 if path.is_dir():
85 if path.suffix.lower() in ext:
86 files.append(path)
87 else:
88 sub_files = [
89 file for file in path.rglob("*") if file.suffix.lower() in ext
90 ]
91 files.extend(sub_files)
92 elif path.is_file() and path.suffix.lower() in ext:
93 files.append(path)
94 return files
97def get_env(key: str, file: str = ".env") -> str:
98 from dotenv import dotenv_values
100 config = dotenv_values(Path.cwd() / file)
101 variable = config.get(key, "") or ""
102 return variable
105def traceback_error(e: Exception, message: str = "", verbose: bool = False):
106 if verbose:
107 logger.exception("An error occurred")
108 if message:
109 print(f"{message}\n")
110 print(str(e))
113def success_msg(message: str):
114 """Format a success message with green styling for console output.
116 Args:
117 message (str): The success message text to format.
119 Returns:
120 str: Formatted message with Rich markup for green styling.
121 """
122 return f"[green]Success[/green]: {message}"
125def error_msg(message: str):
126 """Format an error message with red styling for console output.
128 Args:
129 message (str): The error message text to format.
131 Returns:
132 str: Formatted message with Rich markup for red styling.
133 """
134 return f"[red]Error[/red]: {message}"
137def processing_msg(message: str):
138 """Format a processing status message with orange styling for console output.
140 Args:
141 message (str): The processing message text to format.
143 Returns:
144 str: Formatted message with Rich markup for orange styling.
145 """
146 return f"[dark_orange]Processing[/dark_orange]: {message}"
149def get_class_path(obj_or_class):
150 """Get the full dotted module path for a class or instance.
152 Extracts the fully qualified class path including module and class name,
153 which can be used for dynamic class loading and serialization.
155 Args:
156 obj_or_class: Either a class object or an instance of a class.
158 Returns:
159 str: Fully qualified class path (e.g., 'package.module.ClassName').
160 """
161 if hasattr(obj_or_class, "__class__"):
162 # It's an instance
163 cls = obj_or_class.__class__
164 else:
165 # It's already a class
166 cls = obj_or_class
168 return f"{cls.__module__}.{cls.__name__}"
171def get_all_files(data: list[str], format: str) -> List[Path]:
172 """Collect all files matching the specified format from given paths.
174 Recursively searches through directories and collects files that match
175 the specified format. Displays a progress bar during collection.
177 Args:
178 data (List[str]): List of file paths or directory paths to search.
179 format (str): File extension to filter by (without dot, e.g., 'jpg', 'png').
181 Returns:
182 List[Path]: List of Path objects for files matching the format.
184 Raises:
185 Exception: If no files are found matching the specified format.
186 """
187 files: list[Path] = []
188 for path in track(data, description="Collecting files"):
189 path = Path(path)
190 if path.is_dir():
191 sub_files = [Path(file) for file in path.glob(f"**/*{format}")]
192 files.extend(sub_files)
193 elif path.is_file() and path.suffix.replace(".", "") == format:
194 files.append(Path(path))
195 if len(files) == 0:
196 raise Exception(f"No files found with format '{format}'")
197 return files
200def unzip(path: Path, config: str = "pipeline.yaml") -> dict:
201 """Extract and load configuration data from a zip file.
203 Opens a zip file and reads the specified configuration file (usually YAML)
204 from within the archive, returning the parsed configuration data.
206 Args:
207 path (Path): Path to the zip file to extract from.
208 config (str, optional): Name of config file within zip. Defaults to "pipeline.yaml".
210 Returns:
211 dict: Parsed configuration data from the config file.
213 Raises:
214 ValueError: If the config file is empty, invalid, or missing.
215 zipfile.BadZipFile: If the zip file is corrupted or invalid.
216 """
217 with zipfile.ZipFile(path, "r") as zipf:
218 data = yaml.safe_load(zipf.read(config))
219 if not data:
220 raise ValueError(f"Config file is empty or invalid: {path}")
221 return data
224def from_dir(path: Path, config: str = "pipeline.yaml") -> dict:
225 """Load configuration data from a directory containing a config file.
227 Reads the specified configuration file from a directory and returns
228 the parsed YAML configuration data.
230 Args:
231 path (Path): Path to the directory containing the config file.
232 config (str, optional): Name of config file. Defaults to "pipeline.yaml".
234 Returns:
235 dict: Parsed configuration data from the config file.
237 Raises:
238 ValueError: If the config file is empty, invalid, or missing.
239 FileNotFoundError: If the config file doesn't exist in the directory.
240 yaml.YAMLError: If the YAML file is malformed.
241 """
242 pipeline = path / config
243 with open(pipeline, "r") as f:
244 data = yaml.safe_load(f)
245 if not data:
246 raise ValueError(f"Config file is empty or invalid: {pipeline}")
247 return data
250def load_class_from_string(path: str | None):
251 """Dynamically load a class from a string module path.
253 Takes a fully qualified class path string and imports the class
254 for instantiation. Used for loading task classes from configuration.
256 Args:
257 path (str): Fully qualified class path (e.g., 'module.submodule.ClassName').
259 Returns:
260 type: The loaded class object ready for instantiation.
262 Raises:
263 ImportError: If the module cannot be imported.
264 AttributeError: If the class does not exist in the module.
265 """
266 module_name, class_name = path.rsplit(".", 1)
267 module = importlib.import_module(module_name)
268 cls = getattr(module, class_name)
269 return cls
272@contextmanager
273def change_dir(path: Path):
274 original = Path.cwd()
275 os.chdir(path)
276 try:
277 yield
278 finally:
279 os.chdir(original)
282def remove_exif(image_path: Path, save_path: Path):
283 """Remove EXIF data from an image and save the cleaned image.
285 Opens an image file, removes any embedded EXIF metadata, and saves
286 the cleaned image to the specified path. TIFFs are preserved unchanged:
287 their tags describe scientific channel/page layouts that Pillow cannot
288 safely round-trip.
290 Args:
291 image_path (Path): Path to the original image file.
292 save_path (Path): Path to save the image without EXIF data.
293 """
294 image_path, save_path = Path(image_path), Path(save_path)
295 if image_path.suffix.lower() in {".tif", ".tiff"}:
296 if image_path.resolve() != save_path.resolve():
297 shutil.copy2(image_path, save_path)
298 return
300 from PIL import Image as PILImage
302 with PILImage.open(image_path) as img:
303 data = list(img.get_flattened_data())
304 image_no_exif = PILImage.new(img.mode, img.size)
305 image_no_exif.putdata(data)
306 image_no_exif.save(save_path, format=img.format if img.format else "JPEG")
309def get_distance_matrix(
310 path: str,
311 output: str,
312 model: str = "vgg19",
313):
314 """Compute the distance matrix for the embeddings"""
315 import h5py
316 import numpy as np
317 import pandas as pd
318 from rich.progress import Progress, SpinnerColumn, TextColumn
319 from scipy.spatial.distance import pdist, squareform
321 try:
322 input = Path(path)
323 export = Path(output)
324 export.parent.mkdir(parents=True, exist_ok=True)
325 df = pd.read_parquet(input)
326 embeddings = np.stack(df[model].values) # type: ignore
327 with Progress(
328 SpinnerColumn(),
329 TextColumn("[progress.description]{task.description}"),
330 transient=False,
331 ) as progress:
332 progress.add_task("Computing distance matrix...", total=None)
333 distvec = pdist(embeddings, metric="euclidean")
335 indices = df.filenames.to_list()
336 with h5py.File(export, "w") as f:
337 f.create_dataset("distmatrix", data=squareform(distvec))
338 f.create_dataset("index", data=np.array(indices, dtype="S"))
339 except Exception as e:
340 print(f"Error computing distance matrix: {e}")
343def farthest_first(
344 ref: Path, selected_ids: list = [], k: int = 400, within: bool = False
345) -> list:
346 """Perform farthest-first traversal to select k diverse samples."""
347 import random
349 import h5py
350 import numpy as np
352 with h5py.File(ref, "r") as f:
353 dist_matrix = f["distmatrix"][:]
354 indices: list[str] = [s.decode("utf-8") for s in f["index"][:]] # type: ignore
355 new_ids = []
356 if within and selected_ids:
357 iterator = [indices.index(sid) for sid in selected_ids]
358 selected_ids = []
359 else:
360 iterator = [i for i, _ in enumerate(indices)]
361 if len(selected_ids) == 0:
362 random.seed(42)
363 choice = random.choice(iterator)
364 selected_ids = [choice]
365 new_ids.append(choice)
366 k = k - 1
367 else:
368 selected_ids = [indices.index(sid) for sid in selected_ids]
370 for _ in track(range(k), description="Selecting diverse samples..."):
371 ids = [id for id in iterator if id not in selected_ids]
372 candidates = [(id, np.min(dist_matrix[id, selected_ids])) for id in ids]
373 next_id = max(candidates, key=lambda x: x[1])[0]
374 selected_ids.append(next_id)
375 new_ids.append(next_id)
376 k -= 1
377 new_ids = [indices[sid] for sid in new_ids]
378 return new_ids
381def search(
382 reference: str,
383 existing: str | list = "",
384 k: int = 100,
385 output: str = "",
386 new: bool = True,
387 within: bool = False,
388) -> list[str]:
389 ref = Path(reference)
390 if not ref.exists() or not ref.is_file():
391 raise ValueError(f"Reference distance matrix file not found: {ref}")
392 selected_ids = existing if isinstance(existing, list) else []
393 if (
394 isinstance(existing, str)
395 and Path(existing).is_file()
396 and Path(existing).suffix == ".txt"
397 ):
398 f = open(existing, "r")
399 selected_ids = list(set([line.strip() for line in f if line.strip()]))
400 f.close()
401 new_ids = farthest_first(ref, selected_ids=selected_ids, k=k, within=within)
402 new_ids = list(set(new_ids))
403 duplicates = [dup for dup in new_ids if dup in selected_ids]
404 if duplicates and not within:
405 raise ValueError(f"Duplicate IDs found in selection: {duplicates}")
406 if not new:
407 new_ids = list(set(selected_ids + new_ids))
408 if output:
409 Path(output).write_text("\n".join(new_ids))
410 return new_ids
413def is_image(file_path: Path) -> bool:
414 """Check if a file is an image based on its extension."""
415 image_extensions = [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".gif", ".webp"]
416 return file_path.suffix.lower() in image_extensions
419# Keep original function
420_orig_analyze_scalar = Emitter.analyze_scalar
423def analyze_scalar_allow_block(self, scalar):
424 analysis = _orig_analyze_scalar(self, scalar)
425 # For multi-line scalars, force block style to be allowed
426 if isinstance(analysis, ScalarAnalysis) and "\n" in scalar:
427 analysis.allow_block = True
428 return analysis
431Emitter.analyze_scalar = analyze_scalar_allow_block
434class LiteralDumper(yaml.SafeDumper):
435 pass
438def str_presenter(dumper, data):
439 if "\n" in data:
440 # Force literal block scalar (|)
441 return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")
442 return dumper.represent_scalar("tag:yaml.org,2002:str", data)
445LiteralDumper.add_representer(str, str_presenter)
448def write_yaml(data, filepath: Path | str):
449 with open(filepath, "w") as f:
450 for file_key, data_item in data.items():
451 yaml.dump(
452 {file_key: data_item},
453 f,
454 Dumper=LiteralDumper,
455 sort_keys=False,
456 allow_unicode=True,
457 )
458 f.write("\n")