file uploaded

app.py

@@ -0,0 +1,72 @@
+import torch
+from torchvision import transforms, models
+from art.attacks.evasion import FastGradientMethod
+from art.estimators.classification import PyTorchClassifier
+from PIL import Image
+import numpy as np
+import io
+import base64
+from blind_watermark import WaterMark
+
+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(inputs: dict):
+    input_image = inputs["inputs"]
+    eps_value = inputs.get("eps", 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,
+    )
+
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                             std=[0.229, 0.224, 0.225])
+    ])
+
+    img = Image.open(io.BytesIO(base64.b64decode(input_image))).convert('RGB')
+    img_tensor = transform(img).unsqueeze(0).to(device)
+
+    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])
+    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))
+
+    wm_text = "123"
+    bwm = WaterMark(password_img=123, password_wm=456)
+
+    img_bytes = io.BytesIO()
+    adv_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 {"image": result_image}

requirements.txt

@@ -0,0 +1,6 @@
+torch
+torchvision
+numpy
+pillow
+blind-watermark
+art