mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-10 07:28:15 -05:00
113 lines
3.3 KiB
Python
113 lines
3.3 KiB
Python
import ast, pathlib, unittest
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from tinygrad import Tensor
|
|
from tinygrad.helpers import getenv, CI
|
|
from extra.models.efficientnet import EfficientNet
|
|
from extra.models.vit import ViT
|
|
from extra.models.resnet import ResNet50
|
|
|
|
def _load_labels():
|
|
labels_filename = pathlib.Path(__file__).parent / 'efficientnet/imagenet1000_clsidx_to_labels.txt'
|
|
return ast.literal_eval(labels_filename.read_text())
|
|
|
|
_LABELS = _load_labels()
|
|
|
|
def preprocess(img, new=False):
|
|
# preprocess image
|
|
aspect_ratio = img.size[0] / img.size[1]
|
|
img = img.resize((int(224*max(aspect_ratio,1.0)), int(224*max(1.0/aspect_ratio,1.0))))
|
|
|
|
img = np.array(img)
|
|
y0, x0 =(np.asarray(img.shape)[:2] - 224) // 2
|
|
img = img[y0: y0 + 224, x0: x0 + 224]
|
|
|
|
# low level preprocess
|
|
if new:
|
|
img = img.astype(np.float32)
|
|
img -= [127.0, 127.0, 127.0]
|
|
img /= [128.0, 128.0, 128.0]
|
|
img = img[None]
|
|
else:
|
|
img = np.moveaxis(img, [2, 0, 1], [0, 1, 2])
|
|
img = img.astype(np.float32)[:3].reshape(1, 3, 224, 224)
|
|
img /= 255.0
|
|
img -= np.array([0.485, 0.456, 0.406]).reshape((1, -1, 1, 1))
|
|
img /= np.array([0.229, 0.224, 0.225]).reshape((1, -1, 1, 1))
|
|
return img
|
|
|
|
def _infer(model: EfficientNet, img):
|
|
with Tensor.train(False):
|
|
out = model.forward(Tensor(img)).argmax(axis=-1)
|
|
return out.tolist()
|
|
|
|
chicken_img = preprocess(Image.open(pathlib.Path(__file__).parent / 'efficientnet/Chicken.jpg'))
|
|
car_img = preprocess(Image.open(pathlib.Path(__file__).parent / 'efficientnet/car.jpg'))
|
|
|
|
class TestEfficientNet(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.model = EfficientNet(number=getenv("NUM"))
|
|
cls.model.load_from_pretrained()
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
del cls.model
|
|
|
|
@unittest.skipIf(CI, "covered by test_chicken_car")
|
|
def test_chicken(self):
|
|
labels = _infer(self.model, chicken_img)
|
|
self.assertEqual(_LABELS[labels[0]], "hen")
|
|
|
|
@unittest.skipIf(CI, "covered by test_chicken_car")
|
|
def test_car(self):
|
|
labels = _infer(self.model, car_img)
|
|
self.assertEqual(_LABELS[labels[0]], "sports car, sport car")
|
|
|
|
def test_chicken_car(self):
|
|
labels = _infer(self.model, np.concat([chicken_img, car_img], axis=0))
|
|
self.assertEqual(_LABELS[labels[0]], "hen")
|
|
self.assertEqual(_LABELS[labels[1]], "sports car, sport car")
|
|
|
|
class TestViT(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.model = ViT()
|
|
cls.model.load_from_pretrained()
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
del cls.model
|
|
|
|
def test_chicken(self):
|
|
labels = _infer(self.model, chicken_img)
|
|
self.assertEqual(_LABELS[labels[0]], "cock")
|
|
|
|
def test_car(self):
|
|
labels = _infer(self.model, car_img)
|
|
self.assertEqual(_LABELS[labels[0]], "racer, race car, racing car")
|
|
|
|
class TestResNet(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.model = ResNet50()
|
|
cls.model.load_from_pretrained()
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
del cls.model
|
|
|
|
def test_chicken(self):
|
|
labels = _infer(self.model, chicken_img)
|
|
# NOTE: logits for these two are close
|
|
self.assertIn(_LABELS[labels[0]], ("hen", "cock"))
|
|
|
|
def test_car(self):
|
|
labels = _infer(self.model, car_img)
|
|
self.assertEqual(_LABELS[labels[0]], "sports car, sport car")
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|