Spaces:
Sleeping
Sleeping
File size: 1,245 Bytes
91f4aea |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
import torch
import torchvision.transforms as transforms
from PIL import Image
import sys
import numpy as np
from scipy.special import betainc
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
msg_decoder_path = sys.argv[3]
img_path = sys.argv[1]
key = int(sys.argv[2])
transform_imnet = transforms.Compose([
transforms.ToTensor(),
# transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
transforms.Normalize(mean=[0.5, 0.5, 0.5],std=[0.5, 0.5, 0.5])
])
img = Image.open(sys.argv[1]).convert("RGB").resize((256, 256), Image.BICUBIC)
img = transform_imnet(img).unsqueeze(0).to(device)
print("img.min", img.min())
print("img.max", img.max())
print("img.shape", img.shape)
msg_decoder = torch.jit.load(msg_decoder_path).to(device)
msg_decoder.eval()
with torch.no_grad():
dec = msg_decoder(img)[0].cpu().numpy()
#print("dec = ", dec)
print("dec = ", dec.shape)
msg = np.random.default_rng(seed=key).standard_normal(256)
msg = msg / np.sqrt(np.dot(msg, msg))
print("dec.dec", dec.dot(dec))
print("msg.msg", msg.dot(msg))
print("dec.msg", dec.dot(msg))
cos_angle = dec.dot(msg)
pfa = betainc((256 - 1) * 0.5, 0.5, 1 - cos_angle*cos_angle)
print("pfa = ", pfa)
|