security-ai-edge/tests/test_preprocessor.py

"""
预处理模块单元测试
"""

import unittest
from unittest.mock import MagicMock, patch
import sys
import os
import numpy as np

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


class TestROICropper(unittest.TestCase):
    """测试ROI裁剪器"""
    
    def setUp(self):
        """设置测试环境"""
        from core.preprocessor import ROICropper
        self.cropper = ROICropper()
        
        self.test_image = np.zeros((480, 640, 3), dtype=np.uint8)
        self.test_image[100:200, 200:400] = 255
    
    def test_crop_rectangle(self):
        """测试矩形裁剪"""
        from config.config_models import ROIInfo, ROIType, AlgorithmType
        from core.preprocessor import ROICropper
        
        roi = ROIInfo(
            roi_id="roi001",
            camera_id="cam001",
            roi_type=ROIType.RECTANGLE,
            coordinates=[[200, 100], [400, 200]],
            algorithm_type=AlgorithmType.LEAVE_POST,
        )
        
        cropper = ROICropper()
        result = cropper.crop(self.test_image, roi)
        
        self.assertIsNotNone(result)
        self.assertEqual(result.shape[1], 200)
        self.assertEqual(result.shape[0], 100)
    
    def test_crop_polygon(self):
        """测试多边形裁剪"""
        from config.config_models import ROIInfo, ROIType, AlgorithmType
        from core.preprocessor import ROICropper
        
        roi = ROIInfo(
            roi_id="roi001",
            camera_id="cam001",
            roi_type=ROIType.POLYGON,
            coordinates=[[200, 100], [400, 100], [400, 200], [200, 200]],
            algorithm_type=AlgorithmType.LEAVE_POST,
        )
        
        cropper = ROICropper()
        result = cropper.crop(self.test_image, roi)
        
        self.assertIsNotNone(result)
    
    def test_create_mask(self):
        """测试创建掩码"""
        from config.config_models import ROIInfo, ROIType, AlgorithmType
        from core.preprocessor import ROICropper
        
        roi = ROIInfo(
            roi_id="roi001",
            camera_id="cam001",
            roi_type=ROIType.RECTANGLE,
            coordinates=[[100, 100], [200, 200]],
            algorithm_type=AlgorithmType.LEAVE_POST,
        )
        
        cropper = ROICropper()
        mask = cropper.create_mask((480, 640), roi)
        
        self.assertEqual(mask.shape, (480, 640))
        self.assertTrue(mask[150, 150] > 0)


class TestLetterboxPreprocessor(unittest.TestCase):
    """测试Letterbox预处理器"""
    
    def test_preprocess_16_9(self):
        """测试16:9图像预处理"""
        from core.preprocessor import LetterboxPreprocessor
        
        preprocessor = LetterboxPreprocessor(target_size=(480, 480))
        
        image = np.zeros((720, 1280, 3), dtype=np.uint8)
        image[:, :] = [100, 120, 140]
        
        result, scale_info = preprocessor.preprocess(image)
        
        self.assertEqual(result.shape, (480, 480, 3))
        self.assertEqual(len(scale_info), 4)
    
    def test_preprocess_square(self):
        """测试正方形图像预处理"""
        from core.preprocessor import LetterboxPreprocessor
        
        preprocessor = LetterboxPreprocessor(target_size=(480, 480))
        
        image = np.zeros((640, 640, 3), dtype=np.uint8)
        
        result, scale_info = preprocessor.preprocess(image)
        
        self.assertEqual(result.shape, (480, 480, 3))
    
    def test_revert_coordinates(self):
        """测试坐标还原"""
        from core.preprocessor import LetterboxPreprocessor
        
        preprocessor = LetterboxPreprocessor(target_size=(480, 480))
        
        scale = 0.5
        pad_x = 60
        pad_y = 60
        
        scale_info = (scale, pad_x, pad_y, scale)
        
        box = [100, 100, 200, 200]
        reverted = preprocessor.revert_coordinates(box, scale_info)
        
        self.assertEqual(len(reverted), 4)
        self.assertGreater(reverted[0], 0)


class TestBatchPreprocessor(unittest.TestCase):
    """测试Batch预处理器"""
    
    def test_preprocess_batch(self):
        """测试批次预处理"""
        from core.preprocessor import BatchPreprocessor
        
        preprocessor = BatchPreprocessor(
            target_size=(480, 480),
            max_batch_size=4,
            fp16_mode=True
        )
        
        images = [
            np.zeros((640, 640, 3), dtype=np.uint8)
            for _ in range(2)
        ]
        
        result, scale_info_list = preprocessor.preprocess_batch(images)
        
        self.assertEqual(result.shape[0], 2)
        self.assertEqual(len(scale_info_list), 2)
    
    def test_memory_allocation(self):
        """测试内存分配"""
        from core.preprocessor import BatchPreprocessor
        
        preprocessor = BatchPreprocessor(
            target_size=(480, 480),
            max_batch_size=4,
            fp16_mode=True
        )
        
        mem = preprocessor.allocate_batch_memory(2)
        
        self.assertEqual(mem.shape[0], 2)
        self.assertEqual(mem.dtype, np.float16)


class TestImagePreprocessor(unittest.TestCase):
    """测试图像预处理主类"""
    
    def test_preprocess_single(self):
        """测试单张图像预处理"""
        from core.preprocessor import ImagePreprocessor
        
        preprocessor = ImagePreprocessor()
        
        image = np.zeros((720, 1280, 3), dtype=np.uint8)
        
        result, scale_info = preprocessor.preprocess_single(image)
        
        self.assertEqual(result.shape, (480, 480, 3))
    
    def test_preprocess_batch(self):
        """测试批次预处理"""
        from core.preprocessor import ImagePreprocessor
        
        preprocessor = ImagePreprocessor()
        
        images = [
            np.zeros((720, 1280, 3), dtype=np.uint8)
            for _ in range(4)
        ]
        
        result, scale_info_list = preprocessor.preprocess_batch(images)
        
        self.assertEqual(result.shape[0], 4)
    
    def test_get_statistics(self):
        """测试获取统计"""
        from core.preprocessor import ImagePreprocessor
        
        preprocessor = ImagePreprocessor()
        stats = preprocessor.get_statistics()
        
        self.assertIn("config", stats)
        self.assertIn("memory", stats)


if __name__ == "__main__":
    unittest.main()