92 lines
2.3 KiB
Python
92 lines
2.3 KiB
Python
import torch
|
||
import numpy as np
|
||
from PIL import Image
|
||
import io
|
||
|
||
def _check_dependencies():
|
||
missing = []
|
||
|
||
try:
|
||
import PIL
|
||
from PIL import Image
|
||
except Exception:
|
||
missing.append("Pillow>=10.0,<11.0")
|
||
|
||
try:
|
||
import numpy
|
||
except Exception:
|
||
missing.append("numpy>=1.24,<2.0")
|
||
|
||
if missing:
|
||
raise RuntimeError(
|
||
"Remove Metadata (Extra Stealth) - Missing dependencies:\n"
|
||
+ "\n".join(f"- {m}" for m in missing)
|
||
+ "\n\nInstall via:\n"
|
||
+ "pip install " + " ".join(missing)
|
||
)
|
||
|
||
def _verify_stealth(pil_img):
|
||
leaks = []
|
||
|
||
if pil_img.info:
|
||
leaks.append("Residual PNG metadata")
|
||
|
||
try:
|
||
exif = pil_img.getexif()
|
||
if exif and len(exif) > 0:
|
||
leaks.append("EXIF data present")
|
||
except Exception:
|
||
pass
|
||
|
||
if leaks:
|
||
print(
|
||
"[Remove Metadata – Stealth Warning]\n"
|
||
+ "\n".join(f"- {l}" for l in leaks)
|
||
)
|
||
|
||
class RemoveMetadataStealth:
|
||
def __init__(self):
|
||
_check_dependencies()
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE",),
|
||
"remove_ai_metadata": ("BOOLEAN", {"default": True}),
|
||
"remove_media_metadata": ("BOOLEAN", {"default": True}),
|
||
"remove_privacy_metadata": ("BOOLEAN", {"default": True}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
FUNCTION = "sanitize"
|
||
CATEGORY = "image/utils"
|
||
|
||
def sanitize(
|
||
self,
|
||
image,
|
||
remove_ai_metadata=True,
|
||
remove_media_metadata=True,
|
||
remove_privacy_metadata=True,
|
||
):
|
||
output_images = []
|
||
|
||
for img in image:
|
||
img_np = (img.cpu().numpy() * 255).astype(np.uint8)
|
||
pil_img = Image.fromarray(img_np)
|
||
buffer = io.BytesIO()
|
||
pil_img.save(
|
||
buffer,
|
||
format="PNG",
|
||
optimize=True
|
||
)
|
||
buffer.seek(0)
|
||
clean_img = Image.open(buffer)
|
||
clean_img.info.clear()
|
||
clean_img = clean_img.copy()
|
||
_verify_stealth(clean_img)
|
||
clean_np = np.asarray(clean_img).astype(np.float32) / 255.0
|
||
output_images.append(torch.from_numpy(clean_np))
|
||
return (torch.stack(output_images),)
|