app.py

import gradio as gr
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, models
from art.attacks.evasion import (
    FastGradientMethod, CarliniL2Method, DeepFool, AutoAttack,
    ProjectedGradientDescent, BasicIterativeMethod, SpatialTransformation,
    MomentumIterativeMethod, SaliencyMapMethod, NewtonFool
)
from art.estimators.classification import PyTorchClassifier
from PIL import Image, ImageOps, ImageEnhance, ImageDraw
import numpy as np
import os
from blind_watermark import WaterMark
from torchvision.models import resnet50, vgg16, ResNet50_Weights, VGG16_Weights
import tempfile
import cv2
import dlib

resnet_model = resnet50(weights=ResNet50_Weights.DEFAULT)
num_ftrs_resnet = resnet_model.fc.in_features
resnet_model.fc = nn.Linear(num_ftrs_resnet, 10)
resnet_model = resnet_model.to("cuda" if torch.cuda.is_available() else "cpu")

vgg_model = vgg16(weights=VGG16_Weights.DEFAULT)
num_ftrs_vgg = vgg_model.classifier[6].in_features
vgg_model.classifier[6] = nn.Linear(num_ftrs_vgg, 10)
vgg_model = vgg_model.to("cuda" if torch.cuda.is_available() else "cpu")

criterion = nn.CrossEntropyLoss()
optimizer_resnet = optim.Adam(resnet_model.parameters(), lr=0.001)
optimizer_vgg = optim.Adam(vgg_model.parameters(), lr=0.001)

resnet_classifier = PyTorchClassifier(
    model=resnet_model,
    loss=criterion,
    optimizer=optimizer_resnet,
    input_shape=(3, 224, 224),
    nb_classes=10,
)

vgg_classifier = PyTorchClassifier(
    model=vgg_model,
    loss=criterion,
    optimizer=optimizer_vgg,
    input_shape=(3, 224, 224),
    nb_classes=10,
)

models_dict = {
    "ResNet50": resnet_classifier,
    "VGG16": vgg_classifier
}

face_detector = dlib.get_frontal_face_detector()
landmark_predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")

def detect_face_landmarks(image):
    gray_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
    faces = face_detector(gray_image)

    landmarks = []
    for face in faces:
        shape = landmark_predictor(gray_image, face)

        landmarks.extend([(shape.part(i).x, shape.part(i).y) for i in range(36, 48)])  # 눈
        landmarks.extend([(shape.part(i).x, shape.part(i).y) for i in range(27, 36)])  # 코
        landmarks.extend([(shape.part(i).x, shape.part(i).y) for i in range(48, 68)])  # 입
    return landmarks

def apply_focus_mask(image, landmarks):
    mask = Image.new("L", image.size, 0)
    draw = ImageDraw.Draw(mask)
    for (x, y) in landmarks:
        draw.ellipse((x-10, y-10, x+10, y+10), fill=255)
    return mask

def preprocess_image_with_landmark_focus(image, downscale_factor=0.5):
    landmarks = detect_face_landmarks(image)
    original_size = image.size
    low_res_size = (int(original_size[0] * downscale_factor), int(original_size[1] * downscale_factor))


    mask = apply_focus_mask(image, landmarks)
    low_res_image = image.resize(low_res_size, Image.BILINEAR).resize(original_size, Image.BILINEAR)
    masked_image = Image.composite(low_res_image, image, mask)

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    return transform(masked_image).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu")

def postprocess_image(tensor, original_size):
    adv_img_np = tensor.squeeze(0).cpu().numpy()
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    adv_img_np = (adv_img_np * std[:, None, None]) + mean[:, None, None]
    adv_img_np = np.clip(adv_img_np, 0, 1)
    adv_img_np = adv_img_np.transpose(1, 2, 0)
    adv_image_pil = Image.fromarray((adv_img_np * 255).astype(np.uint8)).resize(original_size)
    return adv_image_pil

def generate_adversarial_image(image, model_name, attack_types, eps_value):
    original_size = image.size

    img_tensor = preprocess_image_with_landmark_focus(image)

    classifier = models_dict[model_name]

    try:
        for attack_type in attack_types:
            if attack_type == "FGSM":
                attack = FastGradientMethod(estimator=classifier, eps=eps_value)
            elif attack_type == "C&W":
                attack = CarliniL2Method(classifier=classifier, confidence=0.05)
            elif attack_type == "DeepFool":
                attack = DeepFool(classifier=classifier, max_iter=20)
            elif attack_type == "AutoAttack":
                attack = AutoAttack(estimator=classifier, eps=eps_value, batch_size=1)
            elif attack_type == "PGD":
                attack = ProjectedGradientDescent(estimator=classifier, eps=eps_value, eps_step=eps_value / 10, max_iter=40)
            elif attack_type == "BIM":
                attack = BasicIterativeMethod(estimator=classifier, eps=eps_value, eps_step=eps_value / 10, max_iter=10)
            elif attack_type == "STA":
                attack = SpatialTransformation(estimator=classifier, max_translation=0.2)
            elif attack_type == "MIM":
                attack = MomentumIterativeMethod(estimator=classifier, eps=eps_value, eps_step=eps_value / 10, max_iter=10)
            elif attack_type == "JSMA":
                attack = SaliencyMapMethod(classifier=classifier, theta=0.1, gamma=0.1)
            elif attack_type == "NewtonFool":
                attack = NewtonFool(classifier=classifier, max_iter=20)

            adv_img_np = attack.generate(x=img_tensor.cpu().numpy())
            img_tensor = torch.tensor(adv_img_np).to("cuda" if torch.cuda.is_available() else "cpu")
    except Exception as e:
        print(f"Error in adversarial generation: {e}")
        return image

    adv_image_pil = postprocess_image(img_tensor, original_size)
    return adv_image_pil

def apply_watermark(image_pil, wm_text="텍스트 삽입", password_img=0, password_wm=0):
    try:
        bwm = WaterMark(password_img=password_img, password_wm=password_wm)
        temp_image_path = tempfile.mktemp(suffix=".png")
        image_pil.save(temp_image_path, format="PNG")

        bwm.read_img(temp_image_path)
        bwm.read_wm(wm_text, mode='str')
        output_path = tempfile.mktemp(suffix=".png")
        bwm.embed(output_path)

        result_image = Image.open(output_path).convert("RGB")
        os.remove(temp_image_path)
        os.remove(output_path)

        return result_image
    except Exception as e:
        print(f"Error in apply_watermark: {str(e)}")
        return image_pil

def extract_watermark(image_pil, password_img=0, password_wm=0):
    bwm = WaterMark(password_img=password_img, password_wm=password_wm)
    temp_image_path = tempfile.mktemp(suffix=".png")
    image_pil.save(temp_image_path, format="PNG")

    extracted_wm_text = bwm.extract(temp_image_path, wm_shape=(32, 32), mode='str')
    os.remove(temp_image_path)
    return extracted_wm_text

def process_image(image, model_name, attack_types, eps_value, wm_text, password_img, password_wm):
    try:
        adv_image = generate_adversarial_image(image, model_name, attack_types, eps_value)
    except Exception as e:
        error_message = f"Error in adversarial generation: {str(e)}"
        return image, error_message, None, None, None

    try:
        watermarked_image = apply_watermark(adv_image, wm_text, int(password_img), int(password_wm))
    except Exception as e:
        error_message = f"Error in watermarking: {str(e)}"
        return image, adv_image, error_message, None, None

    try:
        extracted_wm_text = extract_watermark(watermarked_image, int(password_img), int(password_wm))
    except Exception as e:
        error_message = f"Error in watermark extraction: {str(e)}"
        return image, adv_image, watermarked_image, error_message, None

    output_path = tempfile.mktemp(suffix=".png")
    watermarked_image.save(output_path, format="PNG")
    return image, adv_image, watermarked_image, extracted_wm_text, output_path

# Gradio 인터페이스
interface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="pil", label="이미지를 업로드하세요"),
        gr.Dropdown(choices=["ResNet50", "VGG16"], label="모델 선택"),
        gr.CheckboxGroup(
            choices=["FGSM", "C&W", "DeepFool", "AutoAttack", "PGD", "BIM", "STA", "MIM", "JSMA", "NewtonFool"],
            label="공격 유형 선택",
            value=["PGD"]  # 기본값으로 PGD 선택
        ),
        gr.Slider(
            minimum=0.001,
            maximum=0.9,
            step=0.001,
            value=0.01,  # 기본값 EPS 0.01 설정
            label="Epsilon 값 설정 (노이즈 강도)"
        ),
        gr.Textbox(label="워터마크 텍스트 입력", value="텍스트 삽입"),
        gr.Number(label="이미지 비밀번호", value=0),
        gr.Number(label="워터마크 비밀번호", value=0)
    ],
    outputs=[
        gr.Image(type="numpy", label="원본 이미지"),
        gr.Image(type="numpy", label="적대적 이미지 생성 단계"),
        gr.Image(type="numpy", label="워터마크가 삽입된 최종 이미지"),
        gr.Textbox(label="추출된 워터마크 텍스트"),
        gr.File(label="PNG로 다운로드")
    ]
)

interface.launch(debug=True, share=True)