requirement

app.py

@@ -1,5 +1,7 @@
 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
 from art.estimators.classification import PyTorchClassifier
@@ -8,44 +10,53 @@ import numpy as np
 import io
 import base64
 from blind_watermark import WaterMark
+import matplotlib.pyplot as plt
 
-def load_model():
-    model = models.resnet50(pretrained=False)
-    num_ftrs = model.fc.in_features
-    model.fc = torch.nn.Linear(num_ftrs, 10)
-    
-    model.load_state_dict(torch.load("model.pt", map_location=torch.device('cpu')))
-    model.eval()
-    return model
-
-def process_image(input_image, eps_value=0.3):
-    model = load_model()
-    device = torch.device("cpu")
-    model = model.to(device)
-
-    criterion = torch.nn.CrossEntropyLoss()
-    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
-
-    classifier = PyTorchClassifier(
-        model=model,
-        loss=criterion,
-        optimizer=optimizer,
-        input_shape=(3, 64, 64),
-        nb_classes=10,
-    )
+# Pretrained ResNet50 모델 불러오기 (ImageNet 사전 훈련)
+model = models.resnet50(pretrained=True)
+
+# CIFAR-10에 맞춰 마지막 분류 레이어 수정
+num_ftrs = model.fc.in_features
+model.fc = nn.Linear(num_ftrs, 10)
+
+# 모델을 GPU로 이동
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
 
+# 손실 함수와 옵티마이저 설정
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=0.001)
+
+# PyTorchClassifier 생성
+classifier = PyTorchClassifier(
+    model=model,
+    loss=criterion,
+    optimizer=optimizer,
+    input_shape=(3, 64, 64),
+    nb_classes=10,
+)
+
+# 이미지 전처리 함수
+def preprocess_image(image):
     transform = transforms.Compose([
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.485, 0.456, 0.406],
                              std=[0.229, 0.224, 0.225])
     ])
+    return transform(image).unsqueeze(0).to(device)
 
-    img_tensor = transform(input_image).unsqueeze(0).to(device)
+# FGSM 공격 적용 및 이미지 처리 함수
+def generate_adversarial_image(image, eps_value):
+    img_tensor = preprocess_image(image)
 
+    # FGSM 공격 설정
     attack = FastGradientMethod(estimator=classifier, eps=eps_value)
+
+    # 적대적 예제 생성
     adv_img_tensor = attack.generate(x=img_tensor.cpu().numpy())
     adv_img_tensor = torch.tensor(adv_img_tensor).to(device)
 
+    # 적대적 이미지 변환
     adv_img_np = adv_img_tensor.squeeze(0).cpu().numpy()
     mean = np.array([0.485, 0.456, 0.406])
     std = np.array([0.229, 0.224, 0.225])
@@ -53,28 +64,39 @@ def process_image(input_image, eps_value=0.3):
     adv_img_np = np.clip(adv_img_np, 0, 1)
     adv_img_np = adv_img_np.transpose(1, 2, 0)
 
+    # PIL 이미지로 변환
     adv_image_pil = Image.fromarray((adv_img_np * 255).astype(np.uint8))
+    
+    return adv_image_pil
 
+# 워터마크 삽입 함수
+def apply_watermark(image_pil):
     wm_text = "123"
     bwm = WaterMark(password_img=123, password_wm=456)
-
+    
+    # 이미지에 워터마크 삽입
     img_bytes = io.BytesIO()
-    adv_image_pil.save(img_bytes, format='PNG')
+    image_pil.save(img_bytes, format='PNG')
     bwm.read_img(img_bytes)
     bwm.read_wm(wm_text, mode='str')
     
+    # 워터마크 삽입
     bwm.embed(img_bytes)
-    result_image = base64.b64encode(img_bytes.getvalue()).decode('utf-8')
+    return img_bytes
 
-    return result_image
+# Gradio 함수 (이미지 처리 및 적대적 예제 생성)
+def process_image(image, eps_value):
+    # 적대적 이미지 생성
+    adv_image = generate_adversarial_image(image, eps_value)
 
-# Gradio 인터페이스 정의
-def inference(image, eps_value):
-    return process_image(image, eps_value)
+    # 워터마크 적용
+    result_image = apply_watermark(adv_image)
+
+    return result_image.getvalue()
 
-# 최신 Gradio에서 inputs 대신 components 사용
+# Gradio 인터페이스 정의
 gr.Interface(
-    fn=inference,
+    fn=process_image,
     inputs=[gr.components.Image(type="pil"), gr.components.Slider(0.1, 1.0, step=0.1, value=0.3, label="Epsilon")],
     outputs="image"
 ).launch()