from mobile_cv.model_zoo.models.fbnet_v2 import fbnet
from mobile_cv.model_zoo.models.preprocess import get_preprocess
# Download an example image from the pytorch website
"https://github.com/pytorch/hub/raw/master/dog.jpg",
local_filename, headers = urllib.request.urlretrieve(url, filename)
input_image = Image.open(local_filename)
# fbnet models, supported models could be found in
# mobile_cv/model_zoo/models/model_info/fbnet_v2/*.json
model_name = "dmasking_l3"
model = fbnet(model_name, pretrained=True)
preprocess = get_preprocess(model.arch_def.get("input_size", 224))
input_image = _get_input()
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0)
output = model(input_batch)
output_softmax = torch.nn.functional.softmax(output[0], dim=0)
print(output_softmax.max(0))
if __name__ == "__main__":