def test_chicken(self): chicken_img = Image.open( pathlib.Path(__file__).parent / 'efficientnet/Chicken.jpg') model = EfficientNet(number=0) model.load_from_pretrained() label = _infer(model, chicken_img) self.assertEqual(label, "hen", f"Expected hen but got {label} for number=0")
# run the net out = model.forward(Tensor(img)).cpu() # if you want to look at the outputs """ import matplotlib.pyplot as plt plt.plot(out.data[0]) plt.show() """ return out, retimg if __name__ == "__main__": # instantiate my net model = EfficientNet(int(os.getenv("NUM", "0"))) model.load_from_pretrained() # category labels import ast lbls = fetch("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt") lbls = ast.literal_eval(lbls.decode('utf-8')) # load image and preprocess from PIL import Image url = sys.argv[1] if url == 'webcam': import cv2 cap = cv2.VideoCapture(0) cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) while 1: _ = cap.grab() # discard one frame to circumvent capture buffering