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
# ----------------------------------------------------------------------