init commit

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.st filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,77 @@
+import gradio as gr
+from PIL import Image
+from torchvision import transforms
+import torch
+import random
+import os
+from models.structure.model import SketchKeras
+from safetensors.torch import load_model
+import cv2
+import numpy as np
+
+
+path_to_weights = os.path.join(
+    os.path.dirname(__file__), "models/weights/sketch_keras.st"
+)
+model = SketchKeras()
+load_model(model, path_to_weights)
+model.eval()
+
+
+def preprocess(img):
+    h, w, c = img.shape
+    blurred = cv2.GaussianBlur(img, (0, 0), 3)
+    highpass = img.astype(int) - blurred.astype(int)
+    highpass = highpass.astype(float) / 128.0
+    highpass /= np.max(highpass)
+
+    ret = np.zeros((512, 512, 3), dtype=float)
+    ret[0:h, 0:w, 0:c] = highpass
+    return ret
+
+
+def postprocess(pred, thresh=0.18, smooth=False):
+    assert thresh <= 1.0 and thresh >= 0.0
+
+    pred = np.amax(pred, 0)
+    pred[pred < thresh] = 0
+    pred = 1 - pred
+    pred *= 255
+    pred = np.clip(pred, 0, 255).astype(np.uint8)
+    if smooth:
+        pred = cv2.medianBlur(pred, 3)
+    return pred
+
+
+def output_sketch(img):
+    # resize
+    height, width = float(img.shape[0]), float(img.shape[1])
+    if width > height:
+        new_width, new_height = (512, int(512 / width * height))
+    else:
+        new_width, new_height = (int(512 / height * width), 512)
+    img = cv2.resize(img, (new_width, new_height))
+
+    img = preprocess(img)
+    x = img.reshape(1, *img.shape).transpose(3, 0, 1, 2)
+    x = torch.tensor(x).float()
+
+    with torch.no_grad():
+        pred = model(x)
+
+    pred = pred.squeeze()
+
+    # postprocess
+    output = pred.cpu().detach().numpy()
+    output = postprocess(output, thresh=0.1, smooth=False)
+    output = output[:new_height, :new_width]
+
+    return output
+
+
+gr.Interface(
+    title="Turn Any Image Into a Sketch",
+    fn=output_sketch,
+    inputs=gr.Image(type="numpy"),
+    outputs=gr.Image(type="numpy"),
+).launch()

flagged/log.csv

@@ -0,0 +1,2 @@
+img,output,flag,username,timestamp
+"{""path"":""flagged/img/11968894a849e0c35c0b/a1555132334_10.jpg"",""url"":""http://localhost:7860/file=/private/var/folders/cm/2c5qsgtd2z9_c__1cmj825dm0000gn/T/gradio/2ec5829897639b70466458f600fc46c698660291/a1555132334_10.jpg"",""size"":424307,""orig_name"":""a1555132334_10.jpg"",""mime_type"":""""}","{""path"":""flagged/output/0d8f13ee9e4c0ae5521d/image.png"",""url"":null,""size"":null,""orig_name"":""image.png"",""mime_type"":null}",,,2023-12-08 15:35:22.482224

models/structure/model.py

@@ -0,0 +1,129 @@
+import torch
+import torch.nn as nn
+
+
+class SketchKeras(nn.Module):
+    def __init__(self):
+        super(SketchKeras, self).__init__()
+
+        self.downblock_1 = nn.Sequential(
+            nn.ReflectionPad2d((1, 1, 1, 1)),
+            nn.Conv2d(1, 32, kernel_size=3, stride=1),
+            nn.BatchNorm2d(32, eps=1e-3, momentum=0),
+            nn.ReLU(),
+        )
+        self.downblock_2 = nn.Sequential(
+            nn.ReflectionPad2d((1, 1, 1, 1)),
+            nn.Conv2d(32, 64, kernel_size=4, stride=2),
+            nn.BatchNorm2d(64, eps=1e-3, momentum=0),
+            nn.ReLU(),
+            nn.ReflectionPad2d((1, 1, 1, 1)),
+            nn.Conv2d(64, 64, kernel_size=3, stride=1),
+            nn.BatchNorm2d(64, eps=1e-3, momentum=0),
+            nn.ReLU(),
+        )
+        self.downblock_3 = nn.Sequential(
+            nn.ReflectionPad2d((1, 1, 1, 1)),
+            nn.Conv2d(64, 128, kernel_size=4, stride=2),
+            nn.BatchNorm2d(128, eps=1e-3, momentum=0),
+            nn.ReLU(),
+            nn.ReflectionPad2d((1, 1, 1, 1)),
+            nn.Conv2d(128, 128, kernel_size=3, stride=1),
+            nn.BatchNorm2d(128, eps=1e-3, momentum=0),
+            nn.ReLU(),
+        )
+        self.downblock_4 = nn.Sequential(
+            nn.ReflectionPad2d((1, 1, 1, 1)),
+            nn.Conv2d(128, 256, kernel_size=4, stride=2),
+            nn.BatchNorm2d(256, eps=1e-3, momentum=0),
+            nn.ReLU(),
+            nn.ReflectionPad2d((1, 1, 1, 1)),
+            nn.Conv2d(256, 256, kernel_size=3, stride=1),
+            nn.BatchNorm2d(256, eps=1e-3, momentum=0),
+            nn.ReLU(),
+        )
+        self.downblock_5 = nn.Sequential(
+            nn.ReflectionPad2d((1, 1, 1, 1)),
+            nn.Conv2d(256, 512, kernel_size=4, stride=2),
+            nn.BatchNorm2d(512, eps=1e-3, momentum=0),
+            nn.ReLU(),
+        )
+        self.downblock_6 = nn.Sequential(
+            nn.ReflectionPad2d((1, 1, 1, 1)),
+            nn.Conv2d(512, 512, kernel_size=3, stride=1),
+            nn.BatchNorm2d(512, eps=1e-3, momentum=0),
+            nn.ReLU(),
+        )
+
+        self.upblock_1 = nn.Sequential(
+            nn.Upsample((64, 64)),
+            nn.ReflectionPad2d((1, 2, 1, 2)),
+            nn.Conv2d(1024, 512, kernel_size=4, stride=1),
+            nn.BatchNorm2d(512, eps=1e-3, momentum=0),
+            nn.ReLU(),
+            nn.ReflectionPad2d((1, 1, 1, 1)),
+            nn.Conv2d(512, 256, kernel_size=3, stride=1),
+            nn.BatchNorm2d(256, eps=1e-3, momentum=0),
+            nn.ReLU(),
+        )
+
+        self.upblock_2 = nn.Sequential(
+            nn.Upsample((128, 128)),
+            nn.ReflectionPad2d((1, 2, 1, 2)),
+            nn.Conv2d(512, 256, kernel_size=4, stride=1),
+            nn.BatchNorm2d(256, eps=1e-3, momentum=0),
+            nn.ReLU(),
+            nn.ReflectionPad2d((1, 1, 1, 1)),
+            nn.Conv2d(256, 128, kernel_size=3, stride=1),
+            nn.BatchNorm2d(128, eps=1e-3, momentum=0),
+            nn.ReLU(),
+        )
+
+        self.upblock_3 = nn.Sequential(
+            nn.Upsample((256, 256)),
+            nn.ReflectionPad2d((1, 2, 1, 2)),
+            nn.Conv2d(256, 128, kernel_size=4, stride=1),
+            nn.BatchNorm2d(128, eps=1e-3, momentum=0),
+            nn.ReLU(),
+            nn.ReflectionPad2d((1, 1, 1, 1)),
+            nn.Conv2d(128, 64, kernel_size=3, stride=1),
+            nn.BatchNorm2d(64, eps=1e-3, momentum=0),
+            nn.ReLU(),
+        )
+
+        self.upblock_4 = nn.Sequential(
+            nn.Upsample((512, 512)),
+            nn.ReflectionPad2d((1, 2, 1, 2)),
+            nn.Conv2d(128, 64, kernel_size=4, stride=1),
+            nn.BatchNorm2d(64, eps=1e-3, momentum=0),
+            nn.ReLU(),
+            nn.ReflectionPad2d((1, 1, 1, 1)),
+            nn.Conv2d(64, 32, kernel_size=3, stride=1),
+            nn.BatchNorm2d(32, eps=1e-3, momentum=0),
+            nn.ReLU(),
+        )
+
+        self.last_pad = nn.ReflectionPad2d((1, 1, 1, 1))
+        self.last_conv = nn.Conv2d(64, 1, kernel_size=3, stride=1)
+
+    def forward(self, x):
+        d1 = self.downblock_1(x)
+        d2 = self.downblock_2(d1)
+        d3 = self.downblock_3(d2)
+        d4 = self.downblock_4(d3)
+        d5 = self.downblock_5(d4)
+        d6 = self.downblock_6(d5)
+
+        u1 = torch.cat((d5, d6), dim=1)
+        u1 = self.upblock_1(u1)
+        u2 = torch.cat((d4, u1), dim=1)
+        u2 = self.upblock_2(u2)
+        u3 = torch.cat((d3, u2), dim=1)
+        u3 = self.upblock_3(u3)
+        u4 = torch.cat((d2, u3), dim=1)
+        u4 = self.upblock_4(u4)
+        u5 = torch.cat((d1, u4), dim=1)
+
+        out = self.last_conv(self.last_pad(u5))
+
+        return out

models/weights/sketch_keras.st

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

requirements.txt

@@ -0,0 +1,11 @@
+einops 
+datasets 
+tqdm 
+accelerate 
+torchinfo 
+diffusers
+transformers
+pathlib
+safetensors
+torchvision
+Pillow