Streamlit upload

app.py

@@ -0,0 +1,111 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import streamlit as st
+import torch
+from huggingface_hub import PyTorchModelHubMixin
+from PIL import Image
+from torchvision import transforms
+from torchvision.transforms.functional import to_pil_image
+
+from model import ICN
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def mask_processing(x):
+    if x > 90:
+        return 140
+    elif x < 80:
+        return 0
+    else:
+        return 255
+
+
+def grid_to_heatmap(grid, size=1024):
+    mask = to_pil_image(grid.view(7, 7))
+    mask = mask.resize((size, size), Image.BICUBIC)
+    mask = Image.eval(mask, mask_processing)
+
+    colormap = plt.get_cmap("Wistia")
+    heatmap = np.array(colormap(mask))
+    heatmap = (heatmap * 255).astype(np.uint8)
+    heatmap = Image.fromarray(heatmap)
+
+    return heatmap, mask
+
+
+def summary_image(img, fake, prediction):
+    prediction -= prediction.min()
+    prediction = prediction / prediction.max()
+
+    size = 1024
+
+    img1 = img.resize((size, size))
+    img2 = fake.resize((size, size))
+
+    heatmap, mask = grid_to_heatmap(prediction)
+    img1.paste(heatmap, (0, 0), mask)
+    img2.paste(heatmap, (0, 0), mask)
+
+    return img1, img2
+
+
+@st.cache_resource
+def load_model():
+    model = torch.jit.load("traced_model.pt")
+    model.eval().to(device)
+    return model
+
+
+model = ICN.from_pretrained("AlexBlck/image-comparator").eval().to(device)
+
+# model = load_model()
+
+st.title("Image Comparator Network")
+
+st.write("## Upload a pair of images")
+cols = st.columns(2)
+with cols[0]:
+    im1 = st.file_uploader("Image 1", type=["jpg", "png"])
+with cols[1]:
+    im2 = st.file_uploader("Image 2", type=["jpg", "png"])
+
+if not (im1 and im2):
+    st.stop()
+
+btn = st.button("Run")
+if not btn:
+    st.stop()
+
+im1 = Image.open(im1).convert("RGB")
+im2 = Image.open(im2).convert("RGB")
+
+tr = transforms.Compose(
+    [
+        transforms.Resize(size=(224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ]
+)
+
+img = torch.vstack((tr(im1), tr(im2))).unsqueeze(0)
+
+heatmap, cl = model(img.to(device))
+confs = torch.softmax(cl, dim=1)
+pred = torch.argmax(confs, dim=1).item()
+
+if pred == 0:
+    st.success("No Manipulation Detected")
+    heatmap *= 0
+elif pred == 1:
+    st.warning("Manipulation Detected!")
+else:
+    st.error("Images are not related.")
+    heatmap *= 0
+
+img1, img2 = summary_image(im1, im2, heatmap[0])
+cols = st.columns(2)
+with cols[0]:
+    st.image(img1)
+with cols[1]:
+    st.image(img2)

model.py

@@ -0,0 +1,60 @@
+import torch
+import torch.nn.functional as F
+from huggingface_hub import PyTorchModelHubMixin
+from torch import nn
+from torchvision import models
+
+
+class ICN(nn.Module, PyTorchModelHubMixin):
+    def __init__(self):
+        super().__init__()
+
+        cnn = models.resnet50(pretrained=False)
+        self.cnn_head = nn.Sequential(
+            *list(cnn.children())[:4],
+            *list(list(list(cnn.children())[4].children())[0].children())[:4],
+        )
+        self.cnn_tail = nn.Sequential(
+            *list(list(cnn.children())[4].children()
+                  )[1:], *list(cnn.children())[5:-2]
+        )
+
+        self.conv1 = nn.Conv2d(128, 256, 3, padding=1)
+        self.bn1 = nn.BatchNorm2d(num_features=256)
+
+        self.fc1 = nn.Linear(2048 * 7 * 7, 256)
+        self.fc2 = nn.Linear(256, 7 * 7)
+
+        self.cls_fc = nn.Linear(256, 3)
+
+        self.criterion = nn.CrossEntropyLoss()
+
+    def forward(self, x):
+        # Input: [-1, 6, 224, 224]
+        real = x[:, :3, :, :]
+        fake = x[:, 3:, :, :]
+
+        # Push both images through pretrained backbone
+        real_features = F.relu(self.cnn_head(real))  # [-1, 64, 56, 56]
+        fake_features = F.relu(self.cnn_head(fake))  # [-1, 64, 56, 56]
+
+        # [-1, 128, 56, 56]
+        combined = torch.cat((real_features, fake_features), 1)
+
+        x = self.conv1(combined)  # [-1, 256, 56, 56]
+        x = self.bn1(x)
+        x = F.relu(x)
+
+        x = self.cnn_tail(x)
+        x = x.view(-1, 2048 * 7 * 7)
+
+        # Final feature [-1, 256]
+        d = F.relu(self.fc1(x))
+
+        # Heatmap [-1, 49]
+        grid = self.fc2(d)
+
+        # Classifier [-1, 1]
+        cl = self.cls_fc(d)
+
+        return grid, cl