def predict(image_path, checkpoint_path, save_path): model = Unet() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model.load_state_dict(torch.load(checkpoint_path, map_location=device)) image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE) size = image.shape image = Compose([ ToTensor(), Resize((512, 512)), ])(image) image = image.unsqueeze(0) mask = model(image)[0] mask[mask < 0.5] = 0 mask[mask > 0.5] = 255 mask = Resize(size)(mask) mask = mask.detach().numpy() cv2.imwrite('result.png', mask[0]) pass
from mirror import mirror from mirror.visualisations import * from PIL import Image from torchvision.models import resnet101, resnet18, vgg16 from torchvision.transforms import ToTensor, Resize, Compose # create a model model = vgg16(pretrained=True) cat = Image.open("./cat.jpg") # resize the image and make it a tensor input = Compose([Resize((224, 224)), ToTensor()])(cat) # add 1 dim for batch input = input.unsqueeze(0) # call mirror with the input and the model mirror(input, model, visualisations=[DeepDream])