Coverage for tests/test_wrapper_yolo.py: 100%

1#!/usr/bin/env python3 

2 

3import sys 

4import gc 

5import pytest 

6 

7from unittest.mock import patch 

8from unittest.mock import MagicMock 

9from tests.torch_mock import TorchMock 

10 

11from PIL import Image 

12 

13mock_torch = TorchMock() 

14 

15sys.path.append("yolo") 

16 

17with patch.dict(sys.modules, { 

18 'ultralytics': MagicMock(), 

19 'ultralytics.YOLO': MagicMock(), 

20 'cv2': MagicMock(), 

21 'nvidia_smi': MagicMock(), 

22 'imageio': MagicMock(), 

23 'torch': mock_torch, 

24}): 

25 from image_utils import img_to_base64 

26 from yolo.wrapper_yolo import ImageCharacterExtractor 

27 from yolo.wrapper_yolo import zoom_image 

28 from yolo.wrapper_yolo import take_top_characters 

29 

30 

31@pytest.mark.asyncio 

32async def test_wrapper_yolo() -> None: 

33 model = ImageCharacterExtractor() 

34 assert model is not None 

35 assert model.model_name == "yolo" 

36 

37 model.init() 

38 health = model.get_health() 

39 assert health is not None 

40 timestamps = model.get_timestamps() 

41 assert timestamps is not None 

42 

43 with pytest.raises(ValueError): 

44 await model.get_rest_args({}) 

45 img = Image.new("RGB", (40, 30)) 

46 img_base64 = img_to_base64(img) 

47 await model.get_rest_args({ 

48 "img": img_base64 

49 }) 

50 

51 await model.warmup() 

52 

53 await model.generate(img=img) 

54 

55 del model 

56 del img 

57 del img_base64 

58 gc.collect() 

59 

60 

61@pytest.mark.asyncio 

62async def test_zoom() -> None: 

63 image = zoom_image( 

64 image=Image.new("RGB", (100, 100)), 

65 x=50, y=50, 

66 w=10, h=10, 

67 zoom_factor=2.0) 

68 assert image is not None 

69 assert image.size == (50, 50) 

70 

71 image = zoom_image( 

72 image=Image.new("RGB", (100, 100)), 

73 x=0, y=0, 

74 w=20, h=20, 

75 zoom_factor=1.2) 

76 assert image is not None 

77 assert image.size == (83, 83) 

78 

79 

80@pytest.mark.asyncio 

81async def test_take_top_characters() -> None: 

82 person_zoom_images = [ 

83 (0, 0.7, Image.new("RGB", (32, 32), color=(255, 0, 0))), 

84 (1, 0.9, Image.new("RGB", (64, 48), color=(0, 255, 0))), 

85 (2, 0.85, Image.new("RGB", (32, 16), color=(0, 0, 255))) 

86 ] 

87 

88 top_characters = take_top_characters( 

89 person_zoom_images, 

90 num_characters=2) 

91 assert len(top_characters) == 2 

92 assert top_characters[0].size == (64, 48) 

93 assert top_characters[1].size == (32, 16) 

94 

95 top_characters = take_top_characters( 

96 person_zoom_images, 

97 num_characters=1) 

98 assert len(top_characters) == 1 

99 assert top_characters[0].size == (64, 48)