mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 05:23:41 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
115 lines
3.3 KiB
Python
115 lines
3.3 KiB
Python
import ast, pathlib, unittest
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from tinygrad import Tensor, Context
|
|
from tinygrad.helpers import getenv
|
|
from test.helpers import slow
|
|
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 Context(TRAINING=0):
|
|
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
|
|
|
|
@slow
|
|
def test_chicken(self):
|
|
labels = _infer(self.model, chicken_img)
|
|
self.assertEqual(_LABELS[labels[0]], "hen")
|
|
|
|
@slow
|
|
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.concatenate([chicken_img, car_img], axis=0))
|
|
self.assertEqual(_LABELS[labels[0]], "hen")
|
|
self.assertEqual(_LABELS[labels[1]], "sports car, sport car")
|
|
|
|
@unittest.skip("these pretrained models are no longer available")
|
|
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()
|