Spaces:
Runtime error
Runtime error
File size: 3,080 Bytes
ad552d8 |
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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
import os
import io
import numpy as np
import onnxruntime as ort
from PIL import Image
import dotenv
dotenv.load_dotenv()
GT_MESSAGE = os.environ["GT_MESSAGE"]
QUALITY_COEFFICIENTS = {
"psnr": -0.0022186489180419534,
"ssim": -0.11337077856710862,
"nmi": -0.09878221979274945,
"lpips": 0.3412626374646173,
}
QUALITY_OFFSETS = {
"psnr": 43.54757854447622,
"ssim": 0.984229018845295,
"nmi": 1.7536553655336136,
"lpips": 0.014247652621287854,
}
def compute_performance(image):
session_options = ort.SessionOptions()
session_options.intra_op_num_threads = 1
session_options.inter_op_num_threads = 1
session_options.log_severity_level = 3
model = ort.InferenceSession(
"./kit/models/stable_signature.onnx",
sess_options=session_options,
)
inputs = np.stack(
[
(
(
np.array(
image,
dtype=np.float32,
)
/ 255.0
- [0.485, 0.456, 0.406]
)
/ [0.229, 0.224, 0.225]
)
.transpose((2, 0, 1))
.astype(np.float32)
],
axis=0,
)
outputs = model.run(
None,
{
"image": inputs,
},
)
decoded = (outputs[0] > 0).astype(int)[0]
gt_message = np.array([int(bit) for bit in GT_MESSAGE])
return 1 - np.mean(gt_message != decoded)
from .metrics import (
compute_image_distance_repeated,
load_perceptual_models,
compute_perceptual_metric_repeated,
load_aesthetics_and_artifacts_models,
compute_aesthetics_and_artifacts_scores,
)
def compute_quality(attacked_image, clean_image, quiet=True):
# Compress the image
buffer = io.BytesIO()
attacked_image.save(buffer, format="JPEG", quality=95)
buffer.seek(0)
# Update attacked_image with the compressed version
attacked_image = Image.open(buffer)
modes = ["psnr", "ssim", "nmi", "lpips"]
results = {}
for mode in modes:
if mode in ["psnr", "ssim", "nmi"]:
metrics = compute_image_distance_repeated(
[clean_image],
[attacked_image],
metric_name=mode,
num_workers=1,
verbose=not quiet,
)
results[mode] = metrics
elif mode == "lpips":
model = load_perceptual_models(
mode,
mode="alex",
device="cpu",
)
metrics = compute_perceptual_metric_repeated(
[clean_image],
[attacked_image],
metric_name=mode,
mode="alex",
model=model,
device="cpu",
)
results[mode] = metrics
normalized_quality = 0
for key, value in results.items():
normalized_quality += (value[0] - QUALITY_OFFSETS[key]) * QUALITY_COEFFICIENTS[
key
]
return normalized_quality
|