def get_validation_augmentation():
    """Add paddings to make image shape divisible by 32"""
    test_transform = [
        albu.Resize(320, 640)
    ]
    return albu.Compose(test_transform)

image_path = '002f507.jpg'
img = cv2.imread(image_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
augmented = get_validation_augmentation()(image=img,mask=None)


img = transforms.ToTensor()(augmented['image'])[None]

mask = model(img)

if mask.max() > 0.2:
  print('YES')