PlanOpticon
| 287a3bb… | leo | 1 | """Tests for the frame extractor module.""" |
| 829e24a… | leo | 2 | |
| 287a3bb… | leo | 3 | import os |
| 287a3bb… | leo | 4 | import tempfile |
| 287a3bb… | leo | 5 | |
| 287a3bb… | leo | 6 | import numpy as np |
| 287a3bb… | leo | 7 | import pytest |
| 287a3bb… | leo | 8 | |
| 287a3bb… | leo | 9 | from video_processor.extractors.frame_extractor import ( |
| 287a3bb… | leo | 10 | calculate_frame_difference, |
| 287a3bb… | leo | 11 | is_gpu_available, |
| 829e24a… | leo | 12 | save_frames, |
| 287a3bb… | leo | 13 | ) |
| 829e24a… | leo | 14 | |
| 287a3bb… | leo | 15 | |
| 287a3bb… | leo | 16 | # Create dummy test frames |
| 287a3bb… | leo | 17 | @pytest.fixture |
| 287a3bb… | leo | 18 | def dummy_frames(): |
| 287a3bb… | leo | 19 | # Create a list of dummy frames with different content |
| 287a3bb… | leo | 20 | frames = [] |
| 287a3bb… | leo | 21 | for i in range(3): |
| 287a3bb… | leo | 22 | # Create frame with different intensity for each |
| 287a3bb… | leo | 23 | frame = np.ones((100, 100, 3), dtype=np.uint8) * (i * 50) |
| 287a3bb… | leo | 24 | frames.append(frame) |
| 287a3bb… | leo | 25 | return frames |
| 287a3bb… | leo | 26 | |
| 829e24a… | leo | 27 | |
| 287a3bb… | leo | 28 | def test_calculate_frame_difference(): |
| 287a3bb… | leo | 29 | """Test frame difference calculation.""" |
| 287a3bb… | leo | 30 | # Create two frames with some difference |
| 287a3bb… | leo | 31 | frame1 = np.zeros((100, 100, 3), dtype=np.uint8) |
| 287a3bb… | leo | 32 | frame2 = np.ones((100, 100, 3), dtype=np.uint8) * 128 # 50% intensity |
| 829e24a… | leo | 33 | |
| 287a3bb… | leo | 34 | # Calculate difference |
| 287a3bb… | leo | 35 | diff = calculate_frame_difference(frame1, frame2) |
| 829e24a… | leo | 36 | |
| 287a3bb… | leo | 37 | # Expected difference is around 128/255 = 0.5 |
| 287a3bb… | leo | 38 | assert 0.45 <= diff <= 0.55 |
| 829e24a… | leo | 39 | |
| 287a3bb… | leo | 40 | # Test identical frames |
| 287a3bb… | leo | 41 | diff_identical = calculate_frame_difference(frame1, frame1.copy()) |
| 287a3bb… | leo | 42 | assert diff_identical < 0.001 # Should be very close to 0 |
| 829e24a… | leo | 43 | |
| 287a3bb… | leo | 44 | |
| 287a3bb… | leo | 45 | def test_is_gpu_available(): |
| 287a3bb… | leo | 46 | """Test GPU availability check.""" |
| 287a3bb… | leo | 47 | # This just tests that the function runs without error |
| 287a3bb… | leo | 48 | # We don't assert the result because it depends on the system |
| 287a3bb… | leo | 49 | result = is_gpu_available() |
| 287a3bb… | leo | 50 | assert isinstance(result, bool) |
| 287a3bb… | leo | 51 | |
| 829e24a… | leo | 52 | |
| 287a3bb… | leo | 53 | def test_save_frames(dummy_frames): |
| 287a3bb… | leo | 54 | """Test saving frames to disk.""" |
| 287a3bb… | leo | 55 | with tempfile.TemporaryDirectory() as temp_dir: |
| 287a3bb… | leo | 56 | # Save frames |
| 287a3bb… | leo | 57 | paths = save_frames(dummy_frames, temp_dir, "test_frame") |
| 829e24a… | leo | 58 | |
| 287a3bb… | leo | 59 | # Check that we got the correct number of paths |
| 287a3bb… | leo | 60 | assert len(paths) == len(dummy_frames) |
| 829e24a… | leo | 61 | |
| 287a3bb… | leo | 62 | # Check that files were created |
| 287a3bb… | leo | 63 | for path in paths: |
| 287a3bb… | leo | 64 | assert os.path.exists(path) |
| 829e24a… | leo | 65 | assert os.path.getsize(path) > 0 # Files should have content |