Deci-AI
/
super-gradients
connected to https://github.com/Deci-AI/super-gradients.git


  
1

	
2

	
3

	
4

	
5

	
6

	
7

	
8

	
9

	
10

	
11

	
12

	
13

	
14

	
15

	
16

	
17

	
18

	
19

	
20

	
21

	
22

	
23

	
24

	
25

	
26

	
27

	
28

	
29

	
30

	
31

	
32

	
33

	
34

	
35

	
36

	
37

	
38

	
39

	
40

	
41

	
42

	
43

	
44

	
45

	
46

	
            import unittest

import numpy as np
import torch

from super_gradients.training.utils.media.image import load_images


class TrainingParamsTest(unittest.TestCase):
    def test_load_images(self):

        # list - numpy
        list_images = [np.zeros((3, 100, 100)) for _ in range(15)]
        loaded_images = load_images(list_images)
        self.assertEqual(len(loaded_images), 15)
        for image in loaded_images:
            self.assertIsInstance(image, np.ndarray)
            self.assertEqual(image.shape, (100, 100, 3))

        # numpy - batch
        np_images = np.zeros((15, 3, 100, 100))
        loaded_images = load_images(np_images)
        self.assertEqual(len(loaded_images), 15)
        for image in loaded_images:
            self.assertIsInstance(image, np.ndarray)
            self.assertEqual(image.shape, (100, 100, 3))

        # list - torcj
        list_images = [torch.zeros((3, 100, 100)) for _ in range(15)]
        loaded_images = load_images(list_images)
        self.assertEqual(len(loaded_images), 15)
        for image in loaded_images:
            self.assertIsInstance(image, np.ndarray)
            self.assertEqual(image.shape, (100, 100, 3))

        # torch - batch
        torch_images = torch.zeros((15, 3, 100, 100))
        loaded_images = load_images(torch_images)
        self.assertEqual(len(loaded_images), 15)
        for image in loaded_images:
            self.assertIsInstance(image, np.ndarray)
            self.assertEqual(image.shape, (100, 100, 3))


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