change files

app.py

@@ -0,0 +1,142 @@
+import gradio as gr
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from PIL import Image
+import torch.nn.functional as F
+
+device = torch.device("cpu")
+
+class VGGBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, batch_norm=False):
+        super().__init__()
+        conv2_params = {'kernel_size': (3, 3),
+                        'stride'     : (1, 1),
+                        'padding'   : 1}
+
+        noop = lambda x : x
+        self._batch_norm = batch_norm
+
+        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=out_channels , **conv2_params)
+        self.bn1 = nn.BatchNorm2d(out_channels) if batch_norm else noop
+
+        self.conv2 = nn.Conv2d(in_channels=out_channels, out_channels=out_channels, **conv2_params)
+        self.bn2 = nn.BatchNorm2d(out_channels) if batch_norm else noop
+
+        self.max_pooling = nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2))
+
+    @property
+    def batch_norm(self):
+        return self._batch_norm
+
+    def forward(self,x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = F.relu(x)
+
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = F.relu(x)
+
+        x = self.max_pooling(x)
+
+        return x
+
+
+class VGG16(nn.Module):
+    def __init__(self, input_size, num_classes=10, batch_norm=False):
+        super(VGG16, self).__init__()
+
+        self.in_channels, self.in_width, self.in_height = input_size
+
+        self.block_1 = VGGBlock(self.in_channels, 64, batch_norm=batch_norm)
+        self.block_2 = VGGBlock(64, 128, batch_norm=batch_norm)
+        self.block_3 = VGGBlock(128, 256, batch_norm=batch_norm)
+        self.block_4 = VGGBlock(256,512, batch_norm=batch_norm)
+
+        self.classifier = nn.Sequential(
+                nn.Linear(2048, 4096),
+                nn.ReLU(True),
+                nn.Dropout(p=0.65),
+                nn.Linear(4096, 4096),
+                nn.ReLU(True),
+                nn.Dropout(p=0.65),
+                nn.Linear(4096, num_classes)
+            )
+
+    @property
+    def input_size(self):
+          return self.in_channels, self.in_width, self.in_height
+
+    def forward(self, x):
+        x = self.block_1(x)
+        x = self.block_2(x)
+        x = self.block_3(x)
+        x = self.block_4(x)
+        x = x.view(x.size(0), -1)
+        x = self.classifier(x)
+
+        return x
+
+
+model = VGG16((1,32,32), batch_norm=True)
+model.to(device)
+# Load the saved checkpoint
+model.load_state_dict(torch.load('model.pth', map_location=device))
+
+
+label_map = {
+    0: 'T-shirt/top',
+    1: 'Trouser',
+    2: 'Pullover',
+    3: 'Dress',
+    4: 'Coat',
+    5: 'Sandal',
+    6: 'Shirt',
+    7: 'Sneaker',
+    8: 'FLAG{3883}',
+    9: 'Ankle boot'
+}
+
+def predict_from_local_image(image: str):
+    # Define the transformation to match the model's input requirements
+
+    transform = transforms.Compose([
+        transforms.Resize((32, 32)),  # Resize to the input size of the model
+        transforms.ToTensor(),        # Convert the image to a tensor
+    ])
+
+    # Load the image
+    image = Image.open(image).convert('L')  # Convert numpy array to PIL image and then to grayscale if necessary
+    image = transform(image).unsqueeze(0)        # Add batch dimension
+
+    # Move the image to the specified device
+    image = image.to(device)
+
+    # Set the model to evaluation mode
+    model.eval()
+
+    # Make a prediction
+    with torch.no_grad():
+        output = model(image)
+        _, predicted_label = torch.max(output, 1)
+        confidence = torch.nn.functional.softmax(output, dim=1)[0] * 100
+
+    # Get the predicted class label and confidence
+    predicted_class = label_map[predicted_label.item()]
+    predicted_confidence = confidence[predicted_label.item()].item()
+
+    return predicted_class, predicted_confidence
+
+
+# Gradio interface
+iface = gr.Interface(
+    fn=predict_from_local_image,  # Function to call for prediction
+    inputs=gr.Image(type='filepath', label="Upload an image"),  # Input: .pt file upload
+    outputs=gr.Textbox(label="Predicted Class"),  # Output: Text showing predicted class
+    title="Vault Challenge 4 - DeepFool",  # Title of the interface
+    description="Upload an image, and the model will predict the class. Try to fool the model into predicting the FLAG using DeepFool! Tips: apply DeepFool attack on the image to make the model predict it as a BAG. Note that you should save the adverserial image as a .pt file and upload it to the model to get the FLAG."
+)
+
+# Launch the Gradio interface
+iface.launch()

model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7d2ddbca9b99982d2cbecbbd61174e4dd3f0a06f92ae3f30bbb45003bb66d1ee
+size 119648747

requirements.ttx

@@ -0,0 +1,3 @@
+torch
+torchvision
+Pillow