diff --git a/.github/workflows/release-notes.yml b/.github/workflows/release-notes.yml new file mode 100644 index 00000000..3b275048 --- /dev/null +++ b/.github/workflows/release-notes.yml @@ -0,0 +1,88 @@ +name: Auto-generate Release Notes + +on: + push: + tags: + - 'v*' + +permissions: + contents: write + pull-requests: read + +jobs: + release-notes: + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Generate release notes + id: release_notes + uses: actions/github-script@v7 + with: + script: | + const tag = context.ref.replace('refs/tags/', ''); + const previousTag = await github.rest.repos.getLatestRelease({ + owner: context.repo.owner, + repo: context.repo.repo, + }).then(res => res.data.tag_name).catch(() => null); + + const compareUrl = previousTag + ? `https://github.com/${context.repo.owner}/${context.repo.repo}/compare/${previousTag}...${tag}` + : `https://github.com/${context.repo.owner}/${context.repo.repo}/releases/tag/${tag}`; + + const commits = await github.rest.repos.compareCommits({ + owner: context.repo.owner, + repo: context.repo.repo, + base: previousTag || tag + '^', + head: tag, + }).then(res => res.data.commits).catch(() => []); + + let releaseNotes = `## Changes in ${tag}\n\n`; + + if (commits.length > 0) { + const features = commits.filter(c => c.commit.message.toLowerCase().includes('feat')); + const fixes = commits.filter(c => c.commit.message.toLowerCase().includes('fix')); + const other = commits.filter(c => !c.commit.message.toLowerCase().includes('feat') && !c.commit.message.toLowerCase().includes('fix')); + + if (features.length > 0) { + releaseNotes += '### Features\n'; + features.forEach(c => { + const msg = c.commit.message.split('\n')[0]; + releaseNotes += `- ${msg} ([${c.sha.substring(0, 7)}](${c.html_url}))\n`; + }); + releaseNotes += '\n'; + } + + if (fixes.length > 0) { + releaseNotes += '### Bug Fixes\n'; + fixes.forEach(c => { + const msg = c.commit.message.split('\n')[0]; + releaseNotes += `- ${msg} ([${c.sha.substring(0, 7)}](${c.html_url}))\n`; + }); + releaseNotes += '\n'; + } + + if (other.length > 0) { + releaseNotes += '### Other Changes\n'; + other.forEach(c => { + const msg = c.commit.message.split('\n')[0]; + releaseNotes += `- ${msg} ([${c.sha.substring(0, 7)}](${c.html_url}))\n`; + }); + } + } + + core.setOutput('release_notes', releaseNotes); + + - name: Create Release + uses: actions/create-release@v1 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + with: + tag_name: ${{ github.ref }} + release_name: Release ${{ github.ref }} + body: ${{ steps.release_notes.outputs.release_notes }} + draft: false + prerelease: false diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 00000000..bf830826 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,39 @@ +name: Tests + +on: + pull_request: + branches: [ main, dev ] + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ['3.12', '3.13'] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -e . + pip install pytest pytest-cov + + - name: Run pytest + run: | + pytest tests/ -v --cov=animl --cov-report=xml --cov-report=term + + - name: Upload coverage to Codecov + uses: codecov/codecov-action@v3 + with: + files: ./coverage.xml + flags: unittests + name: codecov-umbrella + fail_ci_if_error: false + verbose: true \ No newline at end of file diff --git a/.gitignore b/.gitignore index 957edba1..045bddd5 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,4 @@ examples/ -tests/ # Byte-compiled / optimized / DLL files __pycache__/ @@ -21,7 +20,6 @@ lib/ lib64/ parts/ sdist/ -tests/ var/ wheels/ pip-wheel-metadata/ diff --git a/README.md b/README.md index f4da11bc..2be2bf19 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -# animl-py 3.3.0 +# animl-py 3.3.1 AniML comprises a variety of machine learning tools for analyzing ecological data. This Python package includes a set of functions to classify subjects within camera trap field data and can handle both images and videos. This package is also available in R: [animl](https://github.com/conservationtechlab/animl) diff --git a/pyproject.toml b/pyproject.toml index 79450e25..9aad0c16 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "animl-lite" -version = "3.3.0" +version = "3.3.1" description = "Tools for classifying camera trap images" readme = "README.md" @@ -13,8 +13,9 @@ requires-python = ">=3.12,<3.14" dependencies = [ "numpy>=2.0.2", "pandas>=2.2.2,<3.0.0", - "pillow>=11.0.0", + "pillow>=12.3.0", "pyexiftool>=0.5.6", + "pyyaml>=6.0.3", "opencv-python>=4.12.0.88", "tqdm>=4.66.5", ] diff --git a/src/animl/__init__.py b/src/animl/__init__.py index 2b82d322..55917338 100644 --- a/src/animl/__init__.py +++ b/src/animl/__init__.py @@ -1,4 +1,4 @@ -__version__ = '3.3.0' +__version__ = '3.3.1' from animl import classification from animl import detection @@ -36,22 +36,24 @@ check_exiftool, check_onnx_cuda, general, get_onnx_device, get_version, plot_all_bounding_boxes, plot_box, plot_from_file, softmax, visualization,) -from animl.video_processing import (extract_frames, get_frame_as_image,) +from animl.video_processing import (extract_frames, get_frame_as_image, + get_images, get_videos,) __all__ = ['IMAGE_EXTENSIONS', 'MD_COLORS', 'MD_LABELS', 'MEGADETECTORv5_SIZE', - 'MODEL_TYPES', 'ManifestGenerator', 'SDZWA_CLASSIFIER_SIZE', - 'VALID_EXTENSIONS', 'VIDEO_EXTENSIONS', 'WorkingDirectory', - 'active_times', 'animlr', 'build_file_manifest', 'check_exiftool', - 'check_file', 'check_onnx_cuda', 'class_list_to_dict', - 'classification', 'classify', 'compute_batched_distance_matrix', - 'compute_distance_matrix', 'cosine_distance', 'detect', 'detection', - 'distance', 'euclidean_squared_distance', 'export', - 'export_camptrapdp', 'export_camtrapR', 'export_coco', - 'export_folders', 'export_megadetector', 'export_timelapse', - 'export_yolo', 'extract_frames', 'extract_miew_embeddings', - 'file_management', 'from_config', 'from_paths', 'general', - 'generator', 'get_animals', 'get_empty', 'get_frame_as_image', - 'get_onnx_device', 'get_version', 'inference', 'load_class_list', + 'MIEWID_SIZE', 'MODEL_TYPES', 'ManifestGenerator', + 'SDZWA_CLASSIFIER_SIZE', 'VALID_EXTENSIONS', 'VIDEO_EXTENSIONS', + 'WorkingDirectory', 'active_times', 'animlr', 'build_file_manifest', + 'check_exiftool', 'check_file', 'check_onnx_cuda', + 'class_list_to_dict', 'classification', 'classify', + 'compute_batched_distance_matrix', 'compute_distance_matrix', + 'cosine_distance', 'detect', 'detection', 'distance', + 'euclidean_squared_distance', 'export', 'export_camptrapdp', + 'export_camtrapR', 'export_coco', 'export_folders', + 'export_megadetector', 'export_timelapse', 'export_yolo', + 'extract_frames', 'extract_miew_embeddings', 'file_management', + 'from_config', 'from_paths', 'general', 'generator', 'get_animals', + 'get_empty', 'get_frame_as_image', 'get_images', 'get_onnx_device', + 'get_version', 'get_videos', 'inference', 'load_class_list', 'load_classifier', 'load_data', 'load_detector', 'load_json', 'load_miew', 'load_yaml', 'manifest_dataloader', 'parse_detections', 'pipeline', 'plot_all_bounding_boxes', 'plot_box', 'plot_from_file', diff --git a/src/animl/detection.py b/src/animl/detection.py index 6f288172..f01a5496 100644 --- a/src/animl/detection.py +++ b/src/animl/detection.py @@ -8,7 +8,6 @@ from shutil import copyfile from typing import Optional import time -from animl.utils.visualization import MD_LABELS import numpy as np import pandas as pd from pathlib import Path @@ -17,7 +16,7 @@ from animl import file_management from animl.generator import manifest_dataloader -from animl.utils.general import _normalize_boxes, _xyxy2xywh, _scale_letterbox, get_onnx_device +from animl.utils.general import _normalize_boxes, _xyxy2xywh, _scale_letterbox, get_onnx_device, _laplacian_variance from animl.utils.visualization import MD_LABELS @@ -49,6 +48,7 @@ def detect(detector, letterbox: bool = True, category_map: Optional[dict] = MD_LABELS, confidence_threshold: float = 0.1, + calculate_clarity: bool = False, file_col: str = 'filepath', checkpoint_path: Optional[str] = None, checkpoint_frequency: int = -1) -> list[dict]: @@ -64,6 +64,7 @@ def detect(detector, letterbox (bool): if True, resize and pad image to keep aspect ratio, else resize without padding category_map (dict): mapping of category IDs to human-readable labels confidence_threshold (float): only detections above this threshold are returned + calculate_clarity (bool): if True, calculate image clarity using Laplacian variance file_col (str): column name containing file paths device (str): specify to run on cpu or gpu checkpoint_path (str): path to checkpoint file @@ -152,6 +153,11 @@ def detect(detector, for _, batch in tqdm(enumerate(dataloader), total=len(manifest)): count += 1 + # handle bad batches (eg. empty images, corrupted files) + if batch is None: + print(f"Warning: batch {count} is None, skipping.") + continue + # ONNX Runtime inference input_name = detector.get_inputs()[0].name outputs = detector.run(None, {input_name: batch[0]})[0] @@ -160,7 +166,8 @@ def detect(detector, letterbox, confidence_threshold=confidence_threshold, category_map=category_map, - model_type=detector.model_type) + model_type=detector.model_type, + calculate_clarity=calculate_clarity) # Process outputs to match expected format results.extend(outputs) @@ -181,7 +188,8 @@ def _convert_detections(predictions: list, letterbox: bool, confidence_threshold: float = 0.1, model_type: str = 'megadetector', - category_map: dict = MD_LABELS) -> pd.DataFrame: + category_map: dict = MD_LABELS, + calculate_clarity=False) -> pd.DataFrame: # Converts output into nested list with categories, conf, and bboxes in expected format for parsing function. # Supports YOLOv5/MDv5, YOLOv6+, and ONNX models with either relative or absolute bounding box outputs. # If letterbox=True, rescales bboxes back to original image size. @@ -192,7 +200,7 @@ def _convert_detections(predictions: list, image_frames = batch_from_dataloader[2] image_sizes = batch_from_dataloader[3] - # if no category map provided, default to MD_LABELS + # if no category map provided, default to MD_LABELS if category_map is None: print("No category map provided, defaulting to MD_LABELS. ", "This may lead to incorrect category labels if using a custom model.") @@ -246,6 +254,15 @@ def _convert_detections(predictions: list, 'bbox_y': float(round(bbox[1], 4)), 'bbox_w': float(round(bbox[2], 4)), 'bbox_h': float(round(bbox[3], 4))} + + if calculate_clarity: + # calculate image clarity using Laplacian variance + clarity = _laplacian_variance(image_tensors[i].transpose(1, 2, 0)) + detection['clarity'] = min(clarity / 500, 1.0) + detection['score'] = (0.4 * detection['conf'] + + 0.3 * detection['bbox_w'] * detection['bbox_h'] + + 0.3 * detection['clarity']) + detections.append(detection) data = {'filepath': str(image_paths[i]), @@ -261,7 +278,8 @@ def parse_detections(detections: list[dict], manifest: Optional[pd.DataFrame] = None, out_file: Optional[str] = None, threshold: float = 0, - file_col: str = "filepath"): + file_col: str = "filepath", + score: bool = False) -> pd.DataFrame: """ Converts listed output from detector to DataFrame. @@ -271,6 +289,7 @@ def parse_detections(detections: list[dict], out_file (str): path to save dataframe threshold (float): parse only detections above given confidence threshold file_col (str): if manifest, merge results onto file_col + score (bool): if True, calculate a score for each detection based on confidence, bbox size, and clarity Returns: df (pd.DataFrame): formatted md outputs, one row per detection @@ -307,6 +326,9 @@ def parse_detections(detections: list[dict], 'category': frame['category'] if 'category' in frame else None, 'category_label': frame['category_label'] if 'category_label' in frame else 'empty', 'conf': None, 'bbox_x': None, 'bbox_y': None, 'bbox_w': None, 'bbox_h': None} + if score: + data['score'] = None + data['clarity'] = None lst.append(data) else: @@ -322,6 +344,10 @@ def parse_detections(detections: list[dict], 'bbox_y': np.clip(detection['bbox_y'], 0, 1), 'bbox_w': np.clip(detection['bbox_w'], 0, 1), 'bbox_h': np.clip(detection['bbox_h'], 0, 1)} + + if score: + data['score'] = detection.get('score', None) + data['clarity'] = detection.get('clarity', None) lst.append(data) df = pd.DataFrame(lst) @@ -362,7 +388,6 @@ def _save_detection_checkpoint(checkpoint_path: str, results: dict) -> None: Path(checkpoint_tmp_path).unlink() - def get_animals(manifest: pd.DataFrame): """ Pulls MD animal detections for classification diff --git a/src/animl/export.py b/src/animl/export.py index 0ecdefb2..f8539d2e 100644 --- a/src/animl/export.py +++ b/src/animl/export.py @@ -5,7 +5,6 @@ @ Kyra Swanson 2023 """ -import json import os import pandas as pd from typing import Optional, Union diff --git a/src/animl/generator.py b/src/animl/generator.py index 6bcc00f4..68e10f2d 100644 --- a/src/animl/generator.py +++ b/src/animl/generator.py @@ -124,8 +124,8 @@ def __getitem__(self, idx: int) -> Optional[Tuple[np.ndarray, str, int, np.ndarr # Normalize if isinstance(self.normalize, dict): img_arr = self.Normalize(img_arr, - mean=self.normalize.get("mean", [0.485, 0.456, 0.406]), - std=self.normalize.get("std", [0.229, 0.224, 0.225])) + mean=self.normalize.get("mean", [0.485, 0.456, 0.406]), + std=self.normalize.get("std", [0.229, 0.224, 0.225])) elif self.normalize is False: # unnormalize back to [0,255] if needed img_arr = img_arr * 255.0 @@ -150,7 +150,7 @@ def extract_frames(self, idx: int, filepath: str) -> Optional[Image.Image]: cap.release() cv2.destroyAllWindows() return img - + def _pil_to_numpy_array(self, img: Image.Image) -> np.ndarray: """ Convert PIL RGB image to numpy array with shape (C, H, W), dtype float32 scaled to [0,1]. @@ -163,7 +163,7 @@ def _pil_to_numpy_array(self, img: Image.Image) -> np.ndarray: arr = arr / 255.0 return arr - def Letterbox(self, + def Letterbox(self, resize_height: int, resize_width: int, image: Image.Image) -> Image.Image: @@ -196,8 +196,7 @@ def Letterbox(self, return padded.resize((target_w, target_h), Image.BILINEAR) - - def Normalize(self, + def Normalize(self, img: np.ndarray, mean: Sequence[float], std: Sequence[float], diff --git a/src/animl/pipeline.py b/src/animl/pipeline.py index f9a1e869..40c8fb6d 100644 --- a/src/animl/pipeline.py +++ b/src/animl/pipeline.py @@ -4,7 +4,6 @@ @ Kyra Swanson 2023 """ import pandas as pd -from pathlib import Path from animl import (classification, detection, export, file_management, video_processing) from animl.utils import visualization @@ -172,8 +171,6 @@ def from_config(config: str): if (file_management.check_file(working_dir.detections, output_type="Detections")): detections = file_management.load_data(working_dir.detections) else: - - detector = detection.load_detector(cfg['detector_file'], model_type=cfg.get('detector_type', 'megadetector'), device=device) categories = cfg.get('detector_class_list', None) if categories is None: diff --git a/src/animl/utils/general.py b/src/animl/utils/general.py index 3b787c8e..22a7f717 100755 --- a/src/animl/utils/general.py +++ b/src/animl/utils/general.py @@ -2,6 +2,7 @@ General utils """ +import cv2 import numpy as np import onnxruntime as ort @@ -27,12 +28,12 @@ def get_onnx_device(user_set=None, quiet=False): if 'CUDAExecutionProvider' in providers: # user selects cuda device and is available - if user_set == 'cpu': + if user_set in ['cpu', 'CPUExecutionProvider']: if not quiet: print('CUDA is available but set to cpu by user.') providers = ['CPUExecutionProvider'] # user selects cuda device and is available - elif user_set in ['cuda', 'cuda:0', 'cuda:1', 'cuda:2', 'cuda:3']: + elif user_set in ['CUDAExecutionProvider', 'cuda', 'cuda:0', 'cuda:1', 'cuda:2', 'cuda:3']: device_number = int(user_set.split(':')[-1]) if ':' in user_set else 0 providers = [('CUDAExecutionProvider', {'device_id': device_number}), 'CPUExecutionProvider'] if not quiet: @@ -58,9 +59,29 @@ def get_onnx_device(user_set=None, quiet=False): return providers # ============================================================================== -# COORDINATE CONVERSION +# FRAME SELECTION # ============================================================================== +def _laplacian_variance(image): + """Calculate Laplacian variance for sharpness""" + # Convert to grayscale + if len(image.shape) == 3: + gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) + else: + gray = image + + # Scale float [0,1] to uint8 [0,255] + if gray.dtype == np.float32 or gray.dtype == np.float64: + gray = (gray * 255).astype(np.uint8) + + # Compute Laplacian variance + laplacian = cv2.Laplacian(gray, cv2.CV_64F) + + return laplacian.var() + +# ============================================================================== +# COORDINATE CONVERSION +# ============================================================================== def _xywh2xyxy(bbox): """ diff --git a/src/animl/video_processing.py b/src/animl/video_processing.py index 64ca55d3..4de742f8 100644 --- a/src/animl/video_processing.py +++ b/src/animl/video_processing.py @@ -4,6 +4,8 @@ """ import os import cv2 +import subprocess +import re from tqdm import tqdm import multiprocessing as mp import pandas as pd @@ -14,6 +16,35 @@ from animl import file_management +def get_images(files, file_col: str = "filepath"): + """ + Get images from a DataFrame of files + Args: + files (pd.DataFrame): DataFrame containing file paths to videos and images. + file_col (str): Column name in the DataFrame that contains the file paths (default is "filepath"). + Returns: + pd.DataFrame: A DataFrame containing only the image files from the input DataFrame. + """ + images = files[files[file_col].apply( + lambda x: Path(x).suffix.lower()).isin(file_management.IMAGE_EXTENSIONS)] + images = images.assign(frame=0) + return images + + +def get_videos(files, file_col: str = "filepath"): + """ + Get videos from a DataFrame of files + Args: + files (pd.DataFrame): DataFrame containing file paths to videos and images. + file_col (str): Column name in the DataFrame that contains the file paths (default is "filepath"). + Returns: + pd.DataFrame: A DataFrame containing only the video files from the input DataFrame. + """ + videos = files[files[file_col].apply( + lambda x: Path(x).suffix.lower()).isin(file_management.VIDEO_EXTENSIONS)] + return videos + + def extract_frames(files, frames: int = 5, fps: Optional[int] = None, @@ -54,12 +85,8 @@ def extract_frames(files, if (fps is None) and (frames is None): raise AssertionError("Either fps or frames need to be defined.") - images = files[files[file_col].apply( - lambda x: Path(x).suffix.lower()).isin(file_management.IMAGE_EXTENSIONS)] - images = images.assign(frame=0) - - videos = files[files[file_col].apply( - lambda x: Path(x).suffix.lower()).isin(file_management.VIDEO_EXTENSIONS)] + images = get_images(files, file_col=file_col) + videos = get_videos(files, file_col=file_col) if not videos.empty: video_frames = [] @@ -120,9 +147,6 @@ def _count_frames(filepath, frames=5, fps=None) -> int: # print(f"Video file {filepath} has 0 frames, skipping.") return None - cap.release() - cv2.destroyAllWindows() - frames_saved = [] frame_capture = 0 @@ -130,21 +154,16 @@ def _count_frames(filepath, frames=5, fps=None) -> int: if fps is not None: video_fps = cap.get(cv2.CAP_PROP_FPS) if video_fps == 0: - # try to calculate fps from duration - duration = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000 # Sometimes unreliable - if duration > 0: - video_fps = frame_count / duration - else: - print(f"Could not determine video FPS, defaulting to {frames} frames uniformly sampled.") - increment = int(frame_count / frames) - while len(frames_saved) < frames: - frames_saved.append([str(filepath), frame_capture]) - frame_capture += increment - return frames_saved - - frames = int(frame_count / video_fps * fps) - sampled_times = [i / fps for i in range(frames)] + # Attempt to get FPS using ffmpeg if OpenCV fails + video_fps = get_fps_from_ffmpeg(filepath) + if video_fps is None: + print("Could not determine video FPS, defaulting to 30 FPS.") + video_fps = 30 # Default to 30 if unable to determine + + n_frames = int(frame_count / video_fps) * fps + sampled_times = [i / fps for i in range(n_frames)] frames_saved = [min(int(round(t * video_fps)), frame_count-1) for t in sampled_times] + frames_saved = [[str(filepath), frame] for frame in frames_saved] # select set number of frames else: @@ -153,6 +172,9 @@ def _count_frames(filepath, frames=5, fps=None) -> int: frames_saved.append([str(filepath), frame_capture]) frame_capture += increment + cap.release() + cv2.destroyAllWindows() + return frames_saved @@ -176,3 +198,29 @@ def get_frame_as_image(video_path, frame=0): if ret: rgb_frame = cv2.cvtColor(still, cv2.COLOR_BGR2RGB) return rgb_frame + + +def get_fps_from_ffmpeg(video_path): + """Extract FPS from ffmpeg output""" + try: + result = subprocess.run( + ['ffmpeg', '-i', video_path], + capture_output=True, + text=True, + timeout=10 + ) + # Search for fps value in output + # Pattern: "X fps" where X is a number (can be decimal) + match = re.search(r'(\d+\.?\d*)\s+fps', result.stderr) + + if match: + fps = float(match.group(1)) + print(f"FPS: {fps}") + return fps + else: + print("Could not find fps in ffmpeg output") + return None + + except Exception as e: + print(f"Error: {e}") + return None diff --git a/tests/buow_test.py b/tests/buow_test.py deleted file mode 100644 index 33eaf9f8..00000000 --- a/tests/buow_test.py +++ /dev/null @@ -1,56 +0,0 @@ -""" -Test custom detector - -""" -import unittest -import time -from pathlib import Path -import pandas as pd - -from animl import file_management, detection, visualization, MEGADETECTORv5_SIZE - - -@unittest.skip -def buow_test(): - start_time = time.time() - # get files - image_dir = Path.cwd() / 'examples' / 'BUOW' - detector = Path.cwd() / 'models/mani_buow_2025.pt' - files = file_management.build_file_manifest(image_dir, exif=False) - - activetimes = file_management.active_times(files, depth=1) - print(activetimes) - - detector = detection.load_detector(detector, "YOLO") - - md_results = detection.detect(detector, - files, - MEGADETECTORv5_SIZE, - MEGADETECTORv5_SIZE, - letterbox=False, - file_col="filepath", - batch_size=4, - num_workers=4, - confidence_threshold=0.1) - - detections = detection.parse_detections(md_results, manifest=files) - - gt_path = Path.cwd() / 'tests' / 'GroundTruth' / 'buow' / 'buow_detections.csv' - gt_manifest = pd.read_csv(gt_path) - - visualization.plot_all_bounding_boxes(detections, 'buow_boxes/', file_col='filepath', min_conf=0.1, - label_col='category', show_confidence=True, - detector_labels={'1': 'adult', '2': 'juvenile'}) - - try: - detections.equals(gt_manifest) - print("BUOW Detection Test Passed!") - except ValueError: - print("filepath columns do not match. Test Failure :(") - print(detections.compare(gt_manifest)) - exit(1) - - print(f"Test completed in {time.time() - start_time:.2f} seconds") - - -buow_test() diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..9998d67b --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +import sys +import types +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +from PIL import Image + + +def _install_optional_stubs() -> None: + try: + import onnxruntime # noqa: F401 + except Exception: + ort = types.ModuleType("onnxruntime") + + class _FakeSession: + def __init__(self, *args, **kwargs): + self.model_type = "mock" + + def get_inputs(self): + return [types.SimpleNamespace(name="input")] + + def run(self, *_args, **_kwargs): + return [np.zeros((1, 1), dtype=np.float32)] + + def get_modelmeta(self): + return types.SimpleNamespace(custom_metadata_map={}) + + ort.InferenceSession = _FakeSession + ort.get_available_providers = lambda: ["CPUExecutionProvider"] + sys.modules["onnxruntime"] = ort + + try: + import exiftool # noqa: F401 + except Exception: + exiftool = types.ModuleType("exiftool") + + class _FakeExifToolHelper: + version = "0.0" + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def get_metadata(self, _filepath): + return [{}] + + exiftool.ExifToolHelper = _FakeExifToolHelper + sys.modules["exiftool"] = exiftool + + +_install_optional_stubs() + + +@pytest.fixture +def image_dir(tmp_path: Path) -> Path: + root = tmp_path / "images" + (root / "station_a" / "cam_1").mkdir(parents=True) + (root / "station_b" / "cam_2").mkdir(parents=True) + + for idx, rel in enumerate([ + Path("station_a/cam_1/a.jpg"), + Path("station_a/cam_1/b.png"), + Path("station_b/cam_2/c.jpeg"), + ]): + out = root / rel + out.parent.mkdir(parents=True, exist_ok=True) + Image.new("RGB", (32 + idx, 24 + idx), color=(idx * 10, 20, 30)).save(out) + + # Supported extension placeholder; exif=False tests do not open it. + (root / "station_b" / "cam_2" / "clip.mp4").write_bytes(b"fake-video") + # Unsupported file + (root / "station_b" / "cam_2" / "notes.txt").write_text("ignore me") + return root + + +@pytest.fixture +def required_results_manifest(tmp_path: Path) -> pd.DataFrame: + img = tmp_path / "img.jpg" + Image.new("RGB", (100, 50), color=(255, 255, 255)).save(img) + return pd.DataFrame( + [ + { + "filepath": str(img), + "filename": img.name, + "filemodifydate": "2026-01-01 00:00:00", + "frame": 0, + "max_detection_conf": 0.9, + "category": 1, + "category_label": "animal", + "conf": 0.9, + "bbox_x": 0.1, + "bbox_y": 0.2, + "bbox_w": 0.3, + "bbox_h": 0.4, + "prediction": "jaguar", + "confidence": 0.8, + "width": 100, + "height": 50, + "datetime": "2026-01-01 00:00:00", + "station": "station_a", + "extension": ".jpg", + } + ] + ) diff --git a/tests/main_test.py b/tests/main_test.py deleted file mode 100644 index f83d070f..00000000 --- a/tests/main_test.py +++ /dev/null @@ -1,44 +0,0 @@ -''' - Main script. - - Runs full animl workflow on a given directory. - User must provide MegaDetector, Classifier, and Class list files, - otherwise will pull MDv5 and the CTL Southwest v2 models by default. - - Usage example - > python -m animl /home/usr/animl-py/examples/Southwest/ - - OR - - > python -m animl /image/dir megadetector.pt classifier.h5 class_file.csv - - Paths to model files must be edited to local machine. - - @ Kyra Swanson, 2023 -''' -import unittest -import time -import shutil -from pathlib import Path - -import animl - - -@unittest.skip -def main_test(): - start_time = time.time() - - image_dir = Path.cwd() / 'examples' / 'Southwest' - workingdir = Path.cwd() / 'examples' / 'Southwest' / 'Animl-Directory' - shutil.rmtree(workingdir, ignore_errors=True) - - megadetector = Path.cwd() / 'models/md_v1000.0.0-sorrel.onnx' - classifier_file = Path.cwd() / 'models/sdzwa_southwest_v3.onnx' - - animl.from_paths(image_dir, megadetector, classifier_file, - sort=True, visualize=True, sequence=False) - - print(f"Pipeline took {time.time() - start_time:.2f} seconds") - - -main_test() diff --git a/tests/reid_test.py b/tests/reid_test.py deleted file mode 100644 index 3a4f3749..00000000 --- a/tests/reid_test.py +++ /dev/null @@ -1,24 +0,0 @@ -import unittest -from pathlib import Path - -import animl - - -# @unittest.skip -def reid_test(): - manifest_path = Path.cwd() / 'examples' / 'Jaguar' - miew_path = Path.cwd() / 'models/miewid_v3.onnx' - manifest = animl.build_file_manifest(manifest_path) - - miew = animl.load_miew(miew_path) - embeddings = animl.extract_miew_embeddings(miew, - manifest, - file_col="filepath") - - print(embeddings.shape) - - e2 = animl.compute_distance_matrix(embeddings, embeddings, metric='euclidean') - cos = animl.compute_distance_matrix(embeddings, embeddings, metric='cosine') - - -reid_test() diff --git a/tests/test_export.py b/tests/test_export.py new file mode 100644 index 00000000..b0e2fae3 --- /dev/null +++ b/tests/test_export.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import pandas as pd +import pytest + +from animl import export + + +def test_export_folders_creates_links(tmp_path: Path, required_results_manifest: pd.DataFrame): + out_dir = tmp_path / "sorted" + out_dir.mkdir() + manifest = required_results_manifest.copy() + out = export.export_folders(manifest, out_dir, label_col="prediction", copy=False) + assert "link" in out.columns + link = Path(out.loc[0, "link"]) + assert link.exists() + + +def test_export_folders_copy_mode(tmp_path: Path, required_results_manifest: pd.DataFrame): + out_dir = tmp_path / "sorted_copy" + out_dir.mkdir() + manifest = required_results_manifest.copy() + out = export.export_folders(manifest, out_dir, label_col="prediction", copy=True) + link = Path(out.loc[0, "link"]) + assert link.exists() + + +def test_export_folders_requires_label_col(tmp_path: Path, required_results_manifest: pd.DataFrame): + out_dir = tmp_path / "sorted_missing" + out_dir.mkdir() + manifest = required_results_manifest.drop(columns=["prediction"]) + with pytest.raises(AssertionError): + export.export_folders(manifest, out_dir, label_col="prediction") + + +def test_remove_link_deletes_file_and_column(tmp_path: Path, required_results_manifest: pd.DataFrame): + out_dir = tmp_path / "sorted_remove" + out_dir.mkdir() + manifest = export.export_folders(required_results_manifest.copy(), out_dir, label_col="prediction") + link_path = Path(manifest.loc[0, "link"]) + assert link_path.exists() + out = export.remove_link(manifest) + assert "link" not in out.columns + assert not link_path.exists() + + +def test_remove_link_requires_column(required_results_manifest: pd.DataFrame): + with pytest.raises(AssertionError): + export.remove_link(required_results_manifest.copy(), link_col="link") + + +def test_update_labels_from_folders_reads_label_from_parent_dir(tmp_path: Path, required_results_manifest: pd.DataFrame): + out_dir = tmp_path / "review" + out_dir.mkdir() + sorted_manifest = export.export_folders(required_results_manifest.copy(), out_dir, label_col="prediction") + updated = export.update_labels_from_folders(sorted_manifest, out_dir) + assert "label" in updated.columns + assert updated.loc[0, "label"] == "jaguar" + + +def test_export_coco_writes_expected_shape(tmp_path: Path, required_results_manifest: pd.DataFrame): + out_file = tmp_path / "out_coco.json" + class_dict = {7: "jaguar"} + coco = export.export_coco(required_results_manifest.copy(), class_dict, out_file) + assert out_file.exists() + assert set(coco.keys()) == {"info", "licenses", "images", "annotations", "categories"} + assert len(coco["images"]) == 1 + assert len(coco["categories"]) == 1 + assert coco["categories"][0]["id"] == 7 + + +def test_export_coco_skips_nan_bbox(tmp_path: Path, required_results_manifest: pd.DataFrame): + out_file = tmp_path / "out_coco_nan.json" + manifest = required_results_manifest.copy() + manifest.loc[0, "bbox_x"] = float("nan") + coco = export.export_coco(manifest, {1: "jaguar"}, out_file) + assert len(coco["annotations"]) == 0 + + +@pytest.mark.parametrize( + "missing_col", + [ + "filepath", + "filename", + "filemodifydate", + "frame", + "max_detection_conf", + "category", + "conf", + "bbox_x", + "bbox_y", + "bbox_w", + "bbox_h", + "prediction", + "confidence", + ], +) +def test_export_coco_requires_all_columns(tmp_path: Path, required_results_manifest: pd.DataFrame, missing_col: str): + out_file = tmp_path / "out.json" + manifest = required_results_manifest.drop(columns=[missing_col]) + with pytest.raises(AssertionError): + export.export_coco(manifest, {1: "jaguar"}, out_file) + + +def test_export_timelapse_animals_only(tmp_path: Path, required_results_manifest: pd.DataFrame): + out_dir = tmp_path / "timelapse1" + csv_loc = export.export_timelapse(required_results_manifest.copy(), out_dir, only_animal=True) + assert csv_loc.exists() + assert (out_dir / "animals.csv").exists() + assert (out_dir / "manifest.csv").exists() + + +def test_export_timelapse_with_non_animals(tmp_path: Path, required_results_manifest: pd.DataFrame): + out_dir = tmp_path / "timelapse2" + manifest = pd.concat( + [ + required_results_manifest.copy(), + required_results_manifest.assign(category=2, category_label="human", prediction="human"), + ], + ignore_index=True, + ) + csv_loc = export.export_timelapse(manifest, out_dir, only_animal=False) + assert csv_loc.exists() + assert (out_dir / "non-animals.csv").exists() + + +def test_export_megadetector_writes_expected_json(tmp_path: Path, required_results_manifest: pd.DataFrame): + out_file = tmp_path / "md.json" + md = export.export_megadetector(required_results_manifest.copy(), out_file=out_file, prompt=False) + assert out_file.exists() + assert "images" in md + assert "detection_categories" in md + assert md["images"][0]["detections"][0]["category"] == 1 + + +def test_export_megadetector_skips_empty_rows(tmp_path: Path, required_results_manifest: pd.DataFrame): + out_file = tmp_path / "md_empty.json" + manifest = required_results_manifest.copy() + manifest.loc[0, "category"] = 0 + md = export.export_megadetector(manifest, out_file=out_file, prompt=False) + assert md["images"][0]["detections"] == [] + + +def test_export_megadetector_requires_columns(tmp_path: Path, required_results_manifest: pd.DataFrame): + out_file = tmp_path / "bad.json" + manifest = required_results_manifest.drop(columns=["bbox_w"]) + with pytest.raises(ValueError): + export.export_megadetector(manifest, out_file=out_file, prompt=False) + + +def test_export_megadetector_output_is_valid_json(tmp_path: Path, required_results_manifest: pd.DataFrame): + out_file = tmp_path / "md2.json" + export.export_megadetector(required_results_manifest.copy(), out_file=out_file, prompt=False) + loaded = json.loads(out_file.read_text()) + assert loaded["info"]["format_version"] == "3.0" diff --git a/tests/test_file_management.py b/tests/test_file_management.py new file mode 100644 index 00000000..2df9dc77 --- /dev/null +++ b/tests/test_file_management.py @@ -0,0 +1,267 @@ +from __future__ import annotations + +from pathlib import Path + +import pandas as pd +import pytest + +from animl import file_management +from animl.generator import ManifestGenerator, manifest_dataloader + + +def test_valid_extensions_contains_expected_types(): + assert ".jpg" in file_management.IMAGE_EXTENSIONS + assert ".mp4" in file_management.VIDEO_EXTENSIONS + assert file_management.IMAGE_EXTENSIONS.issubset(file_management.VALID_EXTENSIONS) + assert file_management.VIDEO_EXTENSIONS.issubset(file_management.VALID_EXTENSIONS) + + +def test_build_file_manifest_finds_supported_files(image_dir: Path): + manifest = file_management.build_file_manifest(image_dir, exif=False) + assert not manifest.empty + assert set(manifest["extension"]).issubset(file_management.VALID_EXTENSIONS) + assert "notes.txt" not in manifest["filename"].tolist() + + +def test_build_file_manifest_nonrecursive_depth0(image_dir: Path): + top_image = image_dir / "top.jpg" + top_image.write_bytes(b"fake") + manifest = file_management.build_file_manifest(image_dir, exif=False, recursive=False) + assert len(manifest) == 1 + assert manifest.iloc[0]["filename"] == "top.jpg" + + +def test_build_file_manifest_missing_dir_raises(tmp_path: Path): + with pytest.raises(FileNotFoundError): + file_management.build_file_manifest(tmp_path / "missing", exif=False) + + +def test_build_file_manifest_station_camera_columns(image_dir: Path): + manifest = file_management.build_file_manifest( + image_dir, + exif=False, + station_depth=1, + camera_depth=2, + ) + assert "station" in manifest.columns + assert "camera" in manifest.columns + assert set(manifest["station"].dropna()) == {"station_a", "station_b"} + + +@pytest.mark.parametrize("depth_key", ["station_depth", "camera_depth"]) +def test_build_file_manifest_depth_validation_for_nonrecursive(image_dir: Path, depth_key: str): + (image_dir / "top.jpg").write_text("x") + kwargs = {"exif": False, "recursive": False, depth_key: 1} + with pytest.raises(ValueError): + file_management.build_file_manifest(image_dir, **kwargs) + + +def test_build_file_manifest_empty_dir_returns_empty_dataframe(tmp_path: Path): + empty = tmp_path / "empty" + empty.mkdir() + manifest = file_management.build_file_manifest(empty, exif=False) + assert isinstance(manifest, pd.DataFrame) + assert manifest.empty + + +def test_working_directory_creates_expected_paths(tmp_path: Path): + wd = file_management.WorkingDirectory(tmp_path) + assert wd.basedir.is_dir() + assert wd.filemanifest.name == "FileManifest.csv" + assert wd.results.name == "Results.csv" + + +def test_working_directory_activate_dirs(tmp_path: Path): + wd = file_management.WorkingDirectory(tmp_path) + wd.activate_linkdir() + wd.activate_visdir() + assert wd.linkdir.is_dir() + assert wd.visdir.is_dir() + + +def test_working_directory_missing_root_raises(tmp_path: Path): + with pytest.raises(FileNotFoundError): + file_management.WorkingDirectory(tmp_path / "nope") + + +def test_save_and_load_data_roundtrip(tmp_path: Path): + df = pd.DataFrame({"a": [1, 2], "b": ["x", "y"]}) + out = tmp_path / "out.csv" + file_management.save_data(df, out, prompt=False) + loaded = file_management.load_data(out) + pd.testing.assert_frame_equal(loaded, df) + + +def test_load_data_requires_csv(tmp_path: Path): + file = tmp_path / "data.txt" + file.write_text("x") + with pytest.raises(AssertionError): + file_management.load_data(file) + + +def test_save_and_load_json_roundtrip(tmp_path: Path): + payload = {"a": 1, "b": [1, 2]} + out = tmp_path / "out.json" + file_management.save_json(payload, out, prompt=False) + loaded = file_management.load_json(out) + assert loaded == payload + + +def test_load_json_requires_json(tmp_path: Path): + file = tmp_path / "x.csv" + file.write_text("a,b\n1,2") + with pytest.raises(AssertionError): + file_management.load_json(file) + + +@pytest.mark.parametrize("suffix", [".yaml", ".yml"]) +def test_save_and_load_yaml_roundtrip(tmp_path: Path, suffix: str): + payload = {"name": "animl", "n": 2} + out = tmp_path / f"out{suffix}" + file_management.save_yaml(payload, out, prompt=False) + loaded = file_management.load_yaml(out) + assert loaded == payload + + +def test_load_yaml_requires_yaml(tmp_path: Path): + file = tmp_path / "x.txt" + file.write_text("name: animl") + with pytest.raises(AssertionError): + file_management.load_yaml(file) + + +@pytest.mark.parametrize( + "response,expected", + [ + ("y", True), + ("n", False), + ("invalid", False), + ], +) +def test_check_file_prompt_responses(monkeypatch, tmp_path: Path, response: str, expected: bool): + f = tmp_path / "file.csv" + f.write_text("a,b\n1,2") + monkeypatch.setattr("builtins.input", lambda _prompt: response) + assert file_management.check_file(f, output_type="Manifest") is expected + + +def test_check_file_false_when_missing(tmp_path: Path): + assert file_management.check_file(tmp_path / "missing.csv") is False + + +def test_class_list_to_dict_happy_path(): + df = pd.DataFrame({"id": [1, 2], "class": ["jaguar", "ocelot"]}) + out = file_management.class_list_to_dict(df) + assert out == {1: "jaguar", 2: "ocelot"} + + +def test_class_list_to_dict_missing_columns_raises(): + df = pd.DataFrame({"x": [1]}) + with pytest.raises(ValueError): + file_management.class_list_to_dict(df) + + +def test_active_times_groups_by_camera(tmp_path: Path): + p1 = tmp_path / "c1" / "a.jpg" + p1.parent.mkdir(parents=True) + p1.write_text("x") + p2 = tmp_path / "c1" / "b.jpg" + p2.write_text("x") + manifest = pd.DataFrame( + { + "filepath": [str(p1), str(p2)], + "datetime": ["2026-01-01 00:00:00", "2026-01-01 00:01:00"], + "camera": ["c1", "c1"], + } + ) + times = file_management.active_times(manifest) + assert ("datetime", "min") in times.columns + assert ("datetime", "max") in times.columns + + +def test_active_times_adds_timestamp_when_missing(tmp_path: Path): + p = tmp_path / "camx" / "a.jpg" + p.parent.mkdir(parents=True) + p.write_text("x") + manifest = pd.DataFrame({"filepath": [str(p)]}) + times = file_management.active_times(manifest, camera_depth=0) + assert len(times) == 1 + + +def test_active_times_validates_manifest_type(): + with pytest.raises(ValueError): + file_management.active_times(["not", "df"]) + + +def test_sequence_calculation_assigns_sequences(): + manifest = pd.DataFrame( + { + "filepath": ["a.jpg", "b.jpg", "c.jpg"], + "station": ["s1", "s1", "s1"], + "datetime": ["2026-01-01 00:00:00", "2026-01-01 00:00:30", "2026-01-01 00:03:00"], + } + ) + out = file_management.sequence_calculation(manifest, station_col="station", maxdiff=60) + assert out["sequence"].tolist() == [0.0, 0.0, 1.0] + + +@pytest.mark.parametrize( + "station_col,maxdiff,exc", + [ + ("", 60, Exception), + ("station", -1, Exception), + ], +) +def test_sequence_calculation_input_validation(station_col, maxdiff, exc): + manifest = pd.DataFrame({"filepath": ["a.jpg"], "station": ["s"], "datetime": ["2026-01-01 00:00:00"]}) + with pytest.raises(exc): + file_management.sequence_calculation(manifest, station_col=station_col, maxdiff=maxdiff) + + +def test_manifest_generator_len_and_item(tmp_path: Path): + img = tmp_path / "img.jpg" + from PIL import Image + + Image.new("RGB", (20, 10), color=(100, 50, 10)).save(img) + df = pd.DataFrame({"filepath": [str(img)]}) + gen = ManifestGenerator(df, crop=False, resize_height=16, resize_width=16) + assert len(gen) == 1 + item = gen[0] + assert item is not None + img_arr, path, frame, hw = item + assert img_arr.shape == (3, 16, 16) + assert path == str(img) + assert frame == 0 + assert hw.tolist() == [10, 20] + + +def test_manifest_generator_requires_bbox_when_crop_true(tmp_path: Path): + img = tmp_path / "img.jpg" + from PIL import Image + + Image.new("RGB", (20, 10), color=(0, 0, 0)).save(img) + df = pd.DataFrame({"filepath": [str(img)]}) + with pytest.raises(ValueError): + ManifestGenerator(df, crop=True) + + +def test_manifest_generator_invalid_crop_coord(tmp_path: Path): + img = tmp_path / "img.jpg" + from PIL import Image + + Image.new("RGB", (20, 10), color=(0, 0, 0)).save(img) + df = pd.DataFrame({"filepath": [str(img)], "bbox_x": [0], "bbox_y": [0], "bbox_w": [1], "bbox_h": [1]}) + with pytest.raises(ValueError): + ManifestGenerator(df, crop=True, crop_coord="bad") + + +def test_manifest_dataloader_yields_batches(tmp_path: Path): + img = tmp_path / "img.jpg" + from PIL import Image + + Image.new("RGB", (20, 10), color=(0, 0, 0)).save(img) + df = pd.DataFrame({"filepath": [str(img)]}) + loader = manifest_dataloader(df, crop=False, resize_height=8, resize_width=8) + batch = next(loader) + assert batch[0].shape == (1, 3, 8, 8) + assert batch[1] == [str(img)] diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py new file mode 100644 index 00000000..a5ec088d --- /dev/null +++ b/tests/test_pipeline.py @@ -0,0 +1,296 @@ +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pandas as pd + +from animl import pipeline + + +class FakeWorkingDir: + def __init__(self, root: Path): + self.root = root + self.basedir = root / "Animl-Directory" + self.basedir.mkdir(exist_ok=True) + self.linkdir = self.basedir / "Sorted" + self.visdir = self.basedir / "Plots" + self.filemanifest = self.basedir / "FileManifest.csv" + self.imageframes = self.basedir / "ImageFrames.csv" + self.results = self.basedir / "Results.csv" + self.predictions = self.basedir / "Predictions.csv" + self.detections = self.basedir / "Detections.csv" + self.mdraw = self.basedir / "MD_Raw.json" + + def activate_linkdir(self): + self.linkdir.mkdir(exist_ok=True) + + def activate_visdir(self): + self.visdir.mkdir(exist_ok=True) + + +def _base_frames_df(tmp_path: Path): + f = tmp_path / "img.jpg" + f.write_bytes(b"x") + return pd.DataFrame( + { + "filepath": [str(f)], + "filename": [f.name], + "extension": [".jpg"], + "datetime": ["2026-01-01 00:00:00"], + "station": ["s1"], + "frame": [0], + } + ) + + +def _base_detections_df(frames_df: pd.DataFrame): + return frames_df.assign( + max_detection_conf=0.9, + category=1, + category_label="animal", + conf=0.9, + bbox_x=0.1, + bbox_y=0.2, + bbox_w=0.3, + bbox_h=0.4, + ) + + +def test_from_paths_detect_only_flow(monkeypatch, tmp_path: Path): + calls = {"save": 0, "classify": 0} + wd = FakeWorkingDir(tmp_path) + frames = _base_frames_df(tmp_path) + detections = _base_detections_df(frames) + + monkeypatch.setattr(pipeline.file_management, "WorkingDirectory", lambda _p: wd) + monkeypatch.setattr(pipeline.file_management, "build_file_manifest", lambda *_a, **_k: frames) + monkeypatch.setattr(pipeline.video_processing, "extract_frames", lambda *_a, **_k: frames) + monkeypatch.setattr(pipeline, "get_onnx_device", lambda **_k: ["CPUExecutionProvider"]) + monkeypatch.setattr(pipeline.file_management, "check_file", lambda *_a, **_k: False) + monkeypatch.setattr(pipeline.detection, "load_detector", lambda *_a, **_k: object()) + monkeypatch.setattr(pipeline.detection, "detect", lambda *_a, **_k: [{"filepath": frames.iloc[0]["filepath"], "detections": []}]) + monkeypatch.setattr(pipeline.detection, "parse_detections", lambda *_a, **_k: detections) + monkeypatch.setattr(pipeline.classification, "classify", lambda *_a, **_k: calls.__setitem__("classify", calls["classify"] + 1)) + monkeypatch.setattr(pipeline.file_management, "save_data", lambda *_a, **_k: calls.__setitem__("save", calls["save"] + 1)) + + out = pipeline.from_paths("/tmp/images", "det.onnx", "cls.onnx", detect_only=True) + assert len(out) == 1 + assert calls["classify"] == 0 + assert calls["save"] == 1 + + +def test_from_paths_classification_single(monkeypatch, tmp_path: Path): + wd = FakeWorkingDir(tmp_path) + frames = _base_frames_df(tmp_path) + detections = _base_detections_df(frames) + animals = detections.copy() + empty = pd.DataFrame(columns=detections.columns) + + monkeypatch.setattr(pipeline.file_management, "WorkingDirectory", lambda _p: wd) + monkeypatch.setattr(pipeline.file_management, "build_file_manifest", lambda *_a, **_k: frames) + monkeypatch.setattr(pipeline.video_processing, "extract_frames", lambda *_a, **_k: frames) + monkeypatch.setattr(pipeline, "get_onnx_device", lambda **_k: ["CPUExecutionProvider"]) + monkeypatch.setattr(pipeline.file_management, "check_file", lambda *_a, **_k: False) + monkeypatch.setattr(pipeline.detection, "load_detector", lambda *_a, **_k: object()) + monkeypatch.setattr(pipeline.detection, "detect", lambda *_a, **_k: [{"filepath": frames.iloc[0]["filepath"], "detections": []}]) + monkeypatch.setattr(pipeline.detection, "parse_detections", lambda *_a, **_k: detections) + monkeypatch.setattr(pipeline.detection, "get_animals", lambda *_a, **_k: animals) + monkeypatch.setattr(pipeline.detection, "get_empty", lambda *_a, **_k: empty) + monkeypatch.setattr(pipeline.classification, "load_classifier", lambda *_a, **_k: (object(), pd.DataFrame({"class": ["jaguar"]}))) + monkeypatch.setattr(pipeline.classification, "classify", lambda *_a, **_k: np.array([[0.8]])) + monkeypatch.setattr( + pipeline.classification, + "single_classification", + lambda *_a, **_k: animals.assign(prediction="jaguar", confidence=0.72), + ) + monkeypatch.setattr(pipeline.file_management, "save_data", lambda *_a, **_k: None) + + out = pipeline.from_paths("/tmp/images", "det.onnx", "cls.onnx", sequence=False) + assert "prediction" in out.columns + assert out.loc[0, "prediction"] == "jaguar" + + +def test_from_paths_classification_sequence(monkeypatch, tmp_path: Path): + wd = FakeWorkingDir(tmp_path) + frames = _base_frames_df(tmp_path) + detections = _base_detections_df(frames) + + monkeypatch.setattr(pipeline.file_management, "WorkingDirectory", lambda _p: wd) + monkeypatch.setattr(pipeline.file_management, "build_file_manifest", lambda *_a, **_k: frames) + monkeypatch.setattr(pipeline.video_processing, "extract_frames", lambda *_a, **_k: frames) + monkeypatch.setattr(pipeline, "get_onnx_device", lambda **_k: ["CPUExecutionProvider"]) + monkeypatch.setattr(pipeline.file_management, "check_file", lambda *_a, **_k: False) + monkeypatch.setattr(pipeline.detection, "load_detector", lambda *_a, **_k: object()) + monkeypatch.setattr(pipeline.detection, "detect", lambda *_a, **_k: [{"filepath": frames.iloc[0]["filepath"], "detections": []}]) + monkeypatch.setattr(pipeline.detection, "parse_detections", lambda *_a, **_k: detections) + monkeypatch.setattr(pipeline.detection, "get_animals", lambda *_a, **_k: detections) + monkeypatch.setattr(pipeline.detection, "get_empty", lambda *_a, **_k: pd.DataFrame()) + monkeypatch.setattr(pipeline.classification, "load_classifier", lambda *_a, **_k: (object(), pd.DataFrame({"class": ["jaguar"]}))) + monkeypatch.setattr(pipeline.classification, "classify", lambda *_a, **_k: np.array([[0.8]])) + monkeypatch.setattr( + pipeline.classification, + "sequence_classification", + lambda *_a, **_k: detections.assign(prediction="jaguar", confidence=0.72, sequence=0), + ) + monkeypatch.setattr(pipeline.file_management, "save_data", lambda *_a, **_k: None) + + out = pipeline.from_paths("/tmp/images", "det.onnx", "cls.onnx", sequence=True) + assert "sequence" in out.columns + + +def test_from_config_detect_only_with_sort_and_visualize(monkeypatch, tmp_path: Path): + wd = FakeWorkingDir(tmp_path) + frames = _base_frames_df(tmp_path) + detections = _base_detections_df(frames) + cfg = { + "image_dir": "/tmp/images", + "detector_file": "det.onnx", + "detect_only": True, + "sort": True, + "visualize": True, + "copy": False, + } + calls = {"sorted": 0, "viz": 0} + + monkeypatch.setattr(pipeline.file_management, "load_yaml", lambda *_a, **_k: cfg) + monkeypatch.setattr(pipeline.file_management, "WorkingDirectory", lambda _p: wd) + monkeypatch.setattr(pipeline.file_management, "build_file_manifest", lambda *_a, **_k: frames) + monkeypatch.setattr(pipeline.video_processing, "extract_frames", lambda *_a, **_k: frames) + monkeypatch.setattr(pipeline, "get_onnx_device", lambda **_k: ["CPUExecutionProvider"]) + monkeypatch.setattr(pipeline.file_management, "check_file", lambda *_a, **_k: False) + monkeypatch.setattr(pipeline.detection, "load_detector", lambda *_a, **_k: object()) + monkeypatch.setattr(pipeline.detection, "detect", lambda *_a, **_k: [{"filepath": frames.iloc[0]["filepath"], "detections": []}]) + monkeypatch.setattr(pipeline.detection, "parse_detections", lambda *_a, **_k: detections) + monkeypatch.setattr( + pipeline.export, + "export_folders", + lambda manifest, *_a, **_k: calls.__setitem__("sorted", calls["sorted"] + 1) or manifest, + ) + monkeypatch.setattr( + pipeline.visualization, + "plot_all_bounding_boxes", + lambda *_a, **_k: calls.__setitem__("viz", calls["viz"] + 1), + ) + monkeypatch.setattr(pipeline.file_management, "save_data", lambda *_a, **_k: None) + + pipeline.from_config("config.yaml") + assert calls["sorted"] == 1 + assert calls["viz"] == 1 + + +def test_from_config_uses_existing_detections_if_present(monkeypatch, tmp_path: Path): + wd = FakeWorkingDir(tmp_path) + frames = _base_frames_df(tmp_path) + detections = _base_detections_df(frames) + cfg = {"image_dir": "/tmp/images", "detector_file": "det.onnx", "classifier_file": "cls.onnx"} + calls = {"load_detector": 0, "load_data": 0} + + monkeypatch.setattr(pipeline.file_management, "load_yaml", lambda *_a, **_k: cfg) + monkeypatch.setattr(pipeline.file_management, "WorkingDirectory", lambda _p: wd) + monkeypatch.setattr(pipeline.file_management, "build_file_manifest", lambda *_a, **_k: frames) + monkeypatch.setattr(pipeline.video_processing, "extract_frames", lambda *_a, **_k: frames) + monkeypatch.setattr(pipeline, "get_onnx_device", lambda **_k: ["CPUExecutionProvider"]) + monkeypatch.setattr(pipeline.file_management, "check_file", lambda *_a, **_k: True) + monkeypatch.setattr( + pipeline.file_management, + "load_data", + lambda *_a, **_k: calls.__setitem__("load_data", calls["load_data"] + 1) or detections, + ) + monkeypatch.setattr( + pipeline.detection, + "load_detector", + lambda *_a, **_k: calls.__setitem__("load_detector", calls["load_detector"] + 1) or object(), + ) + monkeypatch.setattr(pipeline.detection, "get_animals", lambda *_a, **_k: detections) + monkeypatch.setattr(pipeline.detection, "get_empty", lambda *_a, **_k: pd.DataFrame()) + monkeypatch.setattr(pipeline.classification, "load_classifier", lambda *_a, **_k: (object(), pd.DataFrame({"class": ["jaguar"]}))) + monkeypatch.setattr(pipeline.classification, "classify", lambda *_a, **_k: np.array([[0.8]])) + monkeypatch.setattr(pipeline.classification, "single_classification", lambda *_a, **_k: detections.assign(prediction="jaguar", confidence=0.72)) + monkeypatch.setattr(pipeline.file_management, "save_data", lambda *_a, **_k: None) + + pipeline.from_config("config.yaml") + assert calls["load_data"] == 1 + assert calls["load_detector"] == 0 + + +def test_from_config_uses_custom_detector_category_map(monkeypatch, tmp_path: Path): + wd = FakeWorkingDir(tmp_path) + frames = _base_frames_df(tmp_path) + detections = _base_detections_df(frames) + class_list_path = tmp_path / "detector_classes.csv" + pd.DataFrame({"id": [1], "class": ["animal"]}).to_csv(class_list_path, index=False) + cfg = { + "image_dir": "/tmp/images", + "detector_file": "det.onnx", + "classifier_file": "cls.onnx", + "detector_class_list": str(class_list_path), + } + detect_args = {} + + monkeypatch.setattr(pipeline.file_management, "load_yaml", lambda *_a, **_k: cfg) + monkeypatch.setattr(pipeline.file_management, "WorkingDirectory", lambda _p: wd) + monkeypatch.setattr(pipeline.file_management, "build_file_manifest", lambda *_a, **_k: frames) + monkeypatch.setattr(pipeline.video_processing, "extract_frames", lambda *_a, **_k: frames) + monkeypatch.setattr(pipeline, "get_onnx_device", lambda **_k: ["CPUExecutionProvider"]) + monkeypatch.setattr(pipeline.file_management, "check_file", lambda *_a, **_k: False) + monkeypatch.setattr(pipeline.detection, "load_detector", lambda *_a, **_k: object()) + + def _detect(_detector, _frames, **kwargs): + detect_args.update(kwargs) + return [{"filepath": frames.iloc[0]["filepath"], "detections": []}] + + monkeypatch.setattr(pipeline.detection, "detect", _detect) + monkeypatch.setattr(pipeline.detection, "parse_detections", lambda *_a, **_k: detections) + monkeypatch.setattr(pipeline.detection, "get_animals", lambda *_a, **_k: detections) + monkeypatch.setattr(pipeline.detection, "get_empty", lambda *_a, **_k: pd.DataFrame()) + monkeypatch.setattr(pipeline.classification, "load_classifier", lambda *_a, **_k: (object(), pd.DataFrame({"class": ["jaguar"]}))) + monkeypatch.setattr(pipeline.classification, "classify", lambda *_a, **_k: np.array([[0.8]])) + monkeypatch.setattr(pipeline.classification, "single_classification", lambda *_a, **_k: detections.assign(prediction="jaguar", confidence=0.72)) + monkeypatch.setattr(pipeline.file_management, "save_data", lambda *_a, **_k: None) + + pipeline.from_config("config.yaml") + assert detect_args["category_map"] == {1: "animal"} + + +def test_from_config_sequence_path_when_station_present(monkeypatch, tmp_path: Path): + wd = FakeWorkingDir(tmp_path) + frames = _base_frames_df(tmp_path) + detections = _base_detections_df(frames) + cfg = { + "image_dir": "/tmp/images", + "detector_file": "det.onnx", + "classifier_file": "cls.onnx", + "sequence": True, + "empty_class": "empty", + } + calls = {"seq": 0, "single": 0} + + monkeypatch.setattr(pipeline.file_management, "load_yaml", lambda *_a, **_k: cfg) + monkeypatch.setattr(pipeline.file_management, "WorkingDirectory", lambda _p: wd) + monkeypatch.setattr(pipeline.file_management, "build_file_manifest", lambda *_a, **_k: frames) + monkeypatch.setattr(pipeline.video_processing, "extract_frames", lambda *_a, **_k: frames) + monkeypatch.setattr(pipeline, "get_onnx_device", lambda **_k: ["CPUExecutionProvider"]) + monkeypatch.setattr(pipeline.file_management, "check_file", lambda *_a, **_k: False) + monkeypatch.setattr(pipeline.detection, "load_detector", lambda *_a, **_k: object()) + monkeypatch.setattr(pipeline.detection, "detect", lambda *_a, **_k: [{"filepath": frames.iloc[0]["filepath"], "detections": []}]) + monkeypatch.setattr(pipeline.detection, "parse_detections", lambda *_a, **_k: detections) + monkeypatch.setattr(pipeline.detection, "get_animals", lambda *_a, **_k: detections) + monkeypatch.setattr(pipeline.detection, "get_empty", lambda *_a, **_k: pd.DataFrame()) + monkeypatch.setattr(pipeline.classification, "load_classifier", lambda *_a, **_k: (object(), pd.DataFrame({"class": ["jaguar"]}))) + monkeypatch.setattr(pipeline.classification, "classify", lambda *_a, **_k: np.array([[0.8]])) + monkeypatch.setattr( + pipeline.classification, + "sequence_classification", + lambda *_a, **_k: calls.__setitem__("seq", calls["seq"] + 1) or detections.assign(prediction="jaguar", confidence=0.72), + ) + monkeypatch.setattr( + pipeline.classification, + "single_classification", + lambda *_a, **_k: calls.__setitem__("single", calls["single"] + 1) or detections.assign(prediction="jaguar", confidence=0.72), + ) + monkeypatch.setattr(pipeline.file_management, "save_data", lambda *_a, **_k: None) + + pipeline.from_config("config.yaml") + assert calls["seq"] == 1 + assert calls["single"] == 0 diff --git a/tests/test_reid_distance.py b/tests/test_reid_distance.py new file mode 100644 index 00000000..3e0ce59d --- /dev/null +++ b/tests/test_reid_distance.py @@ -0,0 +1,160 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from animl.reid.distance import ( + compute_batched_distance_matrix, + compute_distance_matrix, + cosine_distance, + euclidean_squared_distance, + remove_diagonal, +) + + +def test_remove_diagonal_shape_and_values(): + mat = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + out = remove_diagonal(mat) + assert out.shape == (3, 2) + assert out.tolist() == [[2, 3], [4, 6], [7, 8]] + + +@pytest.mark.parametrize( + "bad_input", + [ + np.array([1, 2, 3]), + np.array([[[1]]]), + ], +) +def test_remove_diagonal_requires_2d_square_array(bad_input): + with pytest.raises(ValueError): + remove_diagonal(bad_input) + + +def test_remove_diagonal_requires_square_matrix(): + with pytest.raises(ValueError): + remove_diagonal(np.zeros((2, 3))) + + +@pytest.mark.parametrize( + "x,y,expected", + [ + ( + np.array([[0.0, 0.0], [1.0, 0.0]]), + np.array([[0.0, 0.0], [0.0, 1.0]]), + np.array([[0.0, 1.0], [1.0, 2.0]]), + ), + ( + np.array([[1.0, 1.0]]), + np.array([[2.0, 1.0], [1.0, 3.0]]), + np.array([[1.0, 4.0]]), + ), + ], +) +def test_euclidean_squared_distance_values(x, y, expected): + out = euclidean_squared_distance(x, y) + assert np.allclose(out, expected) + + +@pytest.mark.parametrize( + "x,y", + [ + (np.array([1, 2]), np.array([[1, 2]])), + (np.array([[1, 2]]), np.array([1, 2])), + (np.array([[1, 2, 3]]), np.array([[1, 2]])), + ], +) +def test_euclidean_squared_distance_validates_shapes(x, y): + with pytest.raises(ValueError): + euclidean_squared_distance(x, y) + + +def test_cosine_distance_identity_and_orthogonal_cases(): + x = np.array([[1.0, 0.0], [0.0, 1.0]]) + out = cosine_distance(x, x) + assert np.allclose(np.diag(out), 0.0) + assert np.isclose(out[0, 1], 1.0) + + +def test_cosine_distance_handles_zero_vectors(): + x = np.array([[0.0, 0.0]]) + y = np.array([[1.0, 0.0]]) + out = cosine_distance(x, y) + assert np.isfinite(out).all() + + +@pytest.mark.parametrize( + "x,y", + [ + (np.array([1, 2]), np.array([[1, 2]])), + (np.array([[1, 2]]), np.array([1, 2])), + (np.array([[1, 2, 3]]), np.array([[1, 2]])), + ], +) +def test_cosine_distance_validates_shapes(x, y): + with pytest.raises(ValueError): + cosine_distance(x, y) + + +@pytest.mark.parametrize("metric", ["euclidean", "cosine"]) +def test_compute_distance_matrix_dispatches(metric): + x = np.array([[1.0, 0.0], [0.0, 1.0]]) + y = np.array([[1.0, 0.0]]) + out = compute_distance_matrix(x, y, metric=metric) + assert out.shape == (2, 1) + + +def test_compute_distance_matrix_unknown_metric_raises(): + with pytest.raises(ValueError): + compute_distance_matrix(np.array([[1.0, 2.0]]), np.array([[1.0, 2.0]]), metric="manhattan") + + +@pytest.mark.parametrize( + "x,y", + [ + (np.array([1, 2]), np.array([[1, 2]])), + (np.array([[1, 2]]), np.array([1, 2])), + (np.array([[1, 2, 3]]), np.array([[1, 2]])), + ], +) +def test_compute_distance_matrix_validates_shapes(x, y): + with pytest.raises(ValueError): + compute_distance_matrix(x, y) + + +def test_compute_batched_distance_matrix_matches_non_batched_euclidean(): + x = np.array([[0.0, 0.0], [1.0, 0.0], [1.0, 1.0]]) + y = np.array([[0.0, 0.0], [0.0, 1.0]]) + full = compute_distance_matrix(x, y, metric="euclidean") + batched = compute_batched_distance_matrix(x, y, metric="euclidean", batch_size=2) + assert np.allclose(full, batched) + + +def test_compute_batched_distance_matrix_matches_non_batched_cosine(): + x = np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]) + y = np.array([[1.0, 0.0], [1.0, 1.0]]) + full = compute_distance_matrix(x, y, metric="cosine") + batched = compute_batched_distance_matrix(x, y, metric="cosine", batch_size=1) + assert np.allclose(full, batched) + + +@pytest.mark.parametrize( + "batch_size", + [0, -1], +) +def test_compute_batched_distance_matrix_requires_positive_batch_size(batch_size): + with pytest.raises(ValueError): + compute_batched_distance_matrix(np.array([[1.0, 2.0]]), np.array([[1.0, 2.0]]), batch_size=batch_size) + + +@pytest.mark.parametrize( + "x,y", + [ + (np.array([1, 2]), np.array([[1, 2]])), + (np.array([[1, 2]]), np.array([1, 2])), + (np.array([[1, 2, 3]]), np.array([[1, 2]])), + ], +) +def test_compute_batched_distance_matrix_validates_shapes(x, y): + with pytest.raises(ValueError): + compute_batched_distance_matrix(x, y) diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 00000000..c49911dc --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,211 @@ +from __future__ import annotations + +import sys +import types + +import numpy as np +import pytest + +from animl import __version__ +from animl.utils import animlr +from animl.utils import general + + +@pytest.mark.parametrize( + "values", + [ + np.array([[0.0, 0.0]]), + np.array([[1.0, 2.0]]), + np.array([[-1.0, -2.0, -3.0]]), + np.array([[10.0, 10.0, 10.0]]), + np.array([[0.1, 0.2, 0.3, 0.4]]), + np.array([[5.0, 1.0, -5.0]]), + np.array([[100.0, 99.0]]), + np.array([[-100.0, -99.0]]), + np.array([[2.5, 2.5, 0.0]]), + np.array([[9.0, 1.0, 1.0, 1.0]]), + ], +) +def test_softmax_rows_sum_to_one(values): + out = general.softmax(values) + assert np.allclose(out.sum(axis=1), 1.0) + + +@pytest.mark.parametrize( + "values,expected_argmax", + [ + (np.array([[1.0, 2.0, 3.0]]), 2), + (np.array([[3.0, 2.0, 1.0]]), 0), + (np.array([[-3.0, -2.0, -1.0]]), 2), + (np.array([[0.0, 5.0, 1.0]]), 1), + (np.array([[8.0, 8.0, 1.0]]), 0), + ], +) +def test_softmax_argmax_matches_input(values, expected_argmax): + out = general.softmax(values) + assert int(np.argmax(out[0])) == expected_argmax + + +def test_get_version_matches_package_version(): + assert animlr.get_version() == __version__ + + +@pytest.mark.parametrize( + "name,value", + [ + ("MEGADETECTORv5_SIZE", general.MEGADETECTORv5_SIZE), + ("SDZWA_CLASSIFIER_SIZE", general.SDZWA_CLASSIFIER_SIZE), + ], +) +def test_constants_are_positive_ints(name, value): + assert isinstance(value, int), f"{name} should be int" + assert value > 0, f"{name} should be positive" + + +@pytest.mark.parametrize("model", ["megadetector", "yolo", "miewid", "classifier"]) +def test_model_types_contains_expected_values(model): + assert model in general.MODEL_TYPES + + +@pytest.mark.parametrize( + "providers,user_set,quiet,expected", + [ + (["CUDAExecutionProvider", "CPUExecutionProvider"], "cpu", False, ["CPUExecutionProvider"]), + (["CUDAExecutionProvider", "CPUExecutionProvider"], "cpu", True, ["CUDAExecutionProvider", "CPUExecutionProvider"]), + ( + ["CUDAExecutionProvider", "CPUExecutionProvider"], + "cuda:1", + True, + [("CUDAExecutionProvider", {"device_id": 1}), "CPUExecutionProvider"], + ), + ( + ["CUDAExecutionProvider", "CPUExecutionProvider"], + "cuda", + True, + [("CUDAExecutionProvider", {"device_id": 0}), "CPUExecutionProvider"], + ), + (["CUDAExecutionProvider", "CPUExecutionProvider"], None, True, ["CUDAExecutionProvider", "CPUExecutionProvider"]), + (["CUDAExecutionProvider", "CPUExecutionProvider"], "weird", True, ["CUDAExecutionProvider", "CPUExecutionProvider"]), + (["CPUExecutionProvider"], "cuda:0", True, ["CPUExecutionProvider"]), + (["CPUExecutionProvider"], None, True, ["CPUExecutionProvider"]), + ], +) +def test_get_onnx_device_variants(monkeypatch, providers, user_set, quiet, expected): + monkeypatch.setattr(general.ort, "get_available_providers", lambda: providers) + result = general.get_onnx_device(user_set=user_set, quiet=quiet) + assert result == expected + + +def test_check_onnx_cuda_true(monkeypatch): + fake_ort = types.SimpleNamespace(get_available_providers=lambda: ["CUDAExecutionProvider", "CPUExecutionProvider"]) + monkeypatch.setitem(sys.modules, "onnxruntime", fake_ort) + assert animlr.check_onnx_cuda() is True + + +def test_check_onnx_cuda_false(monkeypatch): + fake_ort = types.SimpleNamespace(get_available_providers=lambda: ["CPUExecutionProvider"]) + monkeypatch.setitem(sys.modules, "onnxruntime", fake_ort) + assert animlr.check_onnx_cuda() is False + + +def test_check_exiftool_success(monkeypatch): + class OkExif: + version = "12.0" + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + monkeypatch.setitem(sys.modules, "exiftool", types.SimpleNamespace(ExifToolHelper=OkExif)) + assert animlr.check_exiftool() == "12.0" + + +def test_check_exiftool_failure(monkeypatch): + class BadExif: + def __enter__(self): + raise RuntimeError("boom") + + def __exit__(self, exc_type, exc, tb): + return False + + monkeypatch.setitem(sys.modules, "exiftool", types.SimpleNamespace(ExifToolHelper=BadExif)) + assert animlr.check_exiftool() is False + + +@pytest.mark.parametrize( + "bbox,expected", + [ + (np.array([0.1, 0.2, 0.3, 0.4]), np.array([0.1, 0.2, 0.4, 0.6])), + (np.array([0.0, 0.0, 1.0, 1.0]), np.array([0.0, 0.0, 1.0, 1.0])), + (np.array([0.2, 0.2, 0.1, 0.1]), np.array([0.2, 0.2, 0.3, 0.3])), + ], +) +def test_xywh2xyxy(bbox, expected): + out = general._xywh2xyxy(bbox) + assert np.allclose(out, expected) + + +@pytest.mark.parametrize( + "bbox,expected", + [ + (np.array([0.1, 0.2, 0.4, 0.6]), np.array([0.1, 0.2, 0.3, 0.4])), + (np.array([0.0, 0.0, 1.0, 1.0]), np.array([0.0, 0.0, 1.0, 1.0])), + (np.array([0.2, 0.2, 0.3, 0.3]), np.array([0.2, 0.2, 0.1, 0.1])), + ], +) +def test_xyxy2xywh(bbox, expected): + out = general._xyxy2xywh(bbox) + assert np.allclose(out, expected) + + +@pytest.mark.parametrize( + "bbox,expected", + [ + (np.array([0.1, 0.2, 0.3, 0.4]), np.array([0.25, 0.4, 0.3, 0.4])), + (np.array([0.0, 0.0, 1.0, 1.0]), np.array([0.5, 0.5, 1.0, 1.0])), + ], +) +def test_xywh_to_xywhc(bbox, expected): + out = general._xywh_to_xywhc(bbox) + assert np.allclose(out, expected) + + +@pytest.mark.parametrize( + "bbox,width,height,expected", + [ + ([0.1, 0.2, 0.3, 0.4], 100, 50, [10, 10, 40, 30]), + ([0.0, 0.0, 1.0, 1.0], 10, 10, [0, 0, 10, 10]), + ([0.5, 0.5, 0.2, 0.2], 200, 100, [100, 50, 140, 70]), + ], +) +def test_xywh_to_absxyxy(bbox, width, height, expected): + assert general._xywh_to_absxyxy(bbox, width, height) == expected + + +@pytest.mark.parametrize( + "bbox,image_sizes,expected", + [ + (np.array([10, 10, 50, 50], dtype=np.float32), (100, 100), np.array([0.1, 0.1, 0.5, 0.5])), + (np.array([-5, -5, 120, 120], dtype=np.float32), (100, 100), np.array([0.0, 0.0, 1.0, 1.0])), + ], +) +def test_normalize_boxes_clips(bbox, image_sizes, expected): + out = general._normalize_boxes(bbox, image_sizes) + assert np.allclose(out, expected) + + +@pytest.mark.parametrize( + "bbox,resized,original", + [ + (np.array([0.1, 0.1, 0.5, 0.5]), (640, 640), (480, 640)), + (np.array([0.0, 0.0, 1.0, 1.0]), (640, 640), (320, 320)), + (np.array([0.3, 0.2, 0.2, 0.2]), (1280, 1280), (720, 1280)), + ], +) +def test_scale_letterbox_returns_valid_normalized_bbox(bbox, resized, original): + out = general._scale_letterbox(bbox, resized, original) + assert out.shape == (4,) + assert np.all(out >= 0) + assert np.all(out <= 1)