def test_model_to_learner(tmp): model = models.resnet18 # Test if the function loads an ImageNet model (ResNet) trainer learn = model_to_learner(model(pretrained=True)) assert len(learn.data.classes) == 1000 # Check Image net classes assert isinstance(learn.model, models.ResNet) # Test if model can predict very simple image IM_URL = "https://cvbp.blob.core.windows.net/public/images/cvbp_cup.jpg" imagefile = os.path.join(tmp, "cvbp_cup.jpg") urllib.request.urlretrieve(IM_URL, imagefile) category, ind, predict_output = learn.predict( open_image(imagefile, convert_mode="RGB")) assert learn.data.classes[ind] == str(category) == "coffee_mug" # Test if .predict() yield the same output when use .get_preds() one_data = ( ImageList.from_folder(tmp).split_none().label_const( ) # cannot use label_empty because of fastai bug: # https://github.com/fastai/fastai/issues/1908 .transform( tfms=None, size=IMAGENET_IM_SIZE).databunch(bs=1).normalize(imagenet_stats)) learn.data.train_dl = one_data.train_dl get_preds_output = learn.get_preds(ds_type=DatasetType.Train) assert np.all( np.isclose( np.array(get_preds_output[0].tolist() [0]), # Note, get_preds() produces a batch (list) output np.array(predict_output.tolist()), rtol=1e-05, atol=1e-08, ))
imfile = os.path.join(tempdir, "temp.jpg") urllib.request.urlretrieve(path, imfile) else: imfile = os.path.join(get_cmd_cwd(), path) try: im = open_image(imfile, convert_mode='RGB') except: sys.stderr.write(f"'{imfile}' may not be an image file and will be skipped.\n") continue # Select the pre-built model. for m in modeln: if m == "densenet201": model = model_to_learner(models.densenet201(pretrained=True), IMAGENET_IM_SIZE) elif m == "resnet152": model = model_to_learner(models.resnet152(pretrained=True), IMAGENET_IM_SIZE) elif m == "alexnet": model = model_to_learner(models.alexnet(pretrained=True), IMAGENET_IM_SIZE) elif m == "densenet121": model = model_to_learner(models.densenet121(pretrained=True), IMAGENET_IM_SIZE) elif m == "densenet161": model = model_to_learner(models.densenet161(pretrained=True), IMAGENET_IM_SIZE) elif m == "densenet169": model = model_to_learner(models.densenet169(pretrained=True), IMAGENET_IM_SIZE) elif m == "densenet201": model = model_to_learner(models.densenet201(pretrained=True), IMAGENET_IM_SIZE) elif m == "resnet101": model = model_to_learner(models.resnet101(pretrained=True), IMAGENET_IM_SIZE) elif m == "resnet152":
def classify_frame(capture, learner, label): """Use the learner to predict the class label. """ _, frame = capture.read() # Capture frame-by-frame _, ind, prob = learner.predict(Image(utils.cv2torch(frame))) utils.put_text(frame, f"{label[ind]} ({prob[ind]:.2f})") return utils.cv2matplotlib(frame) labels = imagenet_labels() # Load model labels # Load ResNet model # * https://download.pytorch.org/models/resnet18-5c106cde.pth -> ~/.cache/torch/checkpoints/resnet18-5c106cde.pth learn = model_to_learner(models.resnet18(pretrained=True), IMAGENET_IM_SIZE) #learn = model_to_learner(models.resnet152(pretrained=True), IMAGENET_IM_SIZE) #learn = model_to_learner(models.xresnet152(pretrained=True), IMAGENET_IM_SIZE) # Want to load from local copy rather than from ~/.torch? Maybe #learn = load_learner(file="resnet18-5c106cde.pth") #model = untar_data("resnet18-5c106cde.pth") #learn = load_learner(model) func = partial(classify_frame, learner=learn, label=labels) # ---------------------------------------------------------------------- # Run webcam to show processed results # ----------------------------------------------------------------------