├── Informative_Drawings_Line_Art_Generator_ONNX_conversion.ipynb
├── LICENSE
├── README.md
├── github-pages-coop-coep-workaround.js
└── index.html


/Informative_Drawings_Line_Art_Generator_ONNX_conversion.ipynb:
--------------------------------------------------------------------------------
  1 | {
  2 |   "nbformat": 4,
  3 |   "nbformat_minor": 0,
  4 |   "metadata": {
  5 |     "colab": {
  6 |       "name": "Informative Drawings Line Art Generator - ONNX conversion.ipynb",
  7 |       "provenance": [],
  8 |       "collapsed_sections": []
  9 |     },
 10 |     "kernelspec": {
 11 |       "name": "python3",
 12 |       "display_name": "Python 3"
 13 |     },
 14 |     "language_info": {
 15 |       "name": "python"
 16 |     }
 17 |   },
 18 |   "cells": [
 19 |     {
 20 |       "cell_type": "code",
 21 |       "source": [
 22 |         "# https://huggingface.co/spaces/carolineec/informativedrawings"
 23 |       ],
 24 |       "metadata": {
 25 |         "id": "LD-wo3YBz44e"
 26 |       },
 27 |       "execution_count": 1,
 28 |       "outputs": []
 29 |     },
 30 |     {
 31 |       "cell_type": "code",
 32 |       "source": [
 33 |         "!wget -O model.pth https://cdn-lfs.huggingface.co/spaces/carolineec/informativedrawings/c686ced2a666b4850b4bb6ccf0748031c3eda9f822de73a34b8979970d90f0c6"
 34 |       ],
 35 |       "metadata": {
 36 |         "id": "0WJz5Lp3ziyv"
 37 |       },
 38 |       "execution_count": null,
 39 |       "outputs": []
 40 |     },
 41 |     {
 42 |       "cell_type": "code",
 43 |       "source": [
 44 |         "!wget https://huggingface.co/spaces/carolineec/informativedrawings/resolve/main/cat.png\n",
 45 |         "!wget https://huggingface.co/spaces/carolineec/informativedrawings/resolve/main/lizard.png"
 46 |       ],
 47 |       "metadata": {
 48 |         "id": "a1gAtW2xzzXt"
 49 |       },
 50 |       "execution_count": null,
 51 |       "outputs": []
 52 |     },
 53 |     {
 54 |       "cell_type": "code",
 55 |       "execution_count": null,
 56 |       "metadata": {
 57 |         "id": "-toqZE9bzPOF"
 58 |       },
 59 |       "outputs": [],
 60 |       "source": [
 61 |         "import numpy as np\n",
 62 |         "import torch\n",
 63 |         "import torch.nn as nn\n",
 64 |         "from PIL import Image\n",
 65 |         "import torchvision.transforms as transforms\n",
 66 |         "\n",
 67 |         "norm_layer = nn.InstanceNorm2d\n",
 68 |         "\n",
 69 |         "class ResidualBlock(nn.Module):\n",
 70 |         "    def __init__(self, in_features):\n",
 71 |         "        super(ResidualBlock, self).__init__()\n",
 72 |         "\n",
 73 |         "        conv_block = [  nn.ReflectionPad2d(1),\n",
 74 |         "                        nn.Conv2d(in_features, in_features, 3),\n",
 75 |         "                        norm_layer(in_features),\n",
 76 |         "                        nn.ReLU(inplace=True),\n",
 77 |         "                        nn.ReflectionPad2d(1),\n",
 78 |         "                        nn.Conv2d(in_features, in_features, 3),\n",
 79 |         "                        norm_layer(in_features)\n",
 80 |         "                        ]\n",
 81 |         "\n",
 82 |         "        self.conv_block = nn.Sequential(*conv_block)\n",
 83 |         "\n",
 84 |         "    def forward(self, x):\n",
 85 |         "        return x + self.conv_block(x)\n",
 86 |         "\n",
 87 |         "\n",
 88 |         "class Generator(nn.Module):\n",
 89 |         "    def __init__(self, input_nc, output_nc, n_residual_blocks=9, sigmoid=True):\n",
 90 |         "        super(Generator, self).__init__()\n",
 91 |         "\n",
 92 |         "        # Initial convolution block\n",
 93 |         "        model0 = [   nn.ReflectionPad2d(3),\n",
 94 |         "                    nn.Conv2d(input_nc, 64, 7),\n",
 95 |         "                    norm_layer(64),\n",
 96 |         "                    nn.ReLU(inplace=True) ]\n",
 97 |         "        self.model0 = nn.Sequential(*model0)\n",
 98 |         "\n",
 99 |         "        # Downsampling\n",
100 |         "        model1 = []\n",
101 |         "        in_features = 64\n",
102 |         "        out_features = in_features*2\n",
103 |         "        for _ in range(2):\n",
104 |         "            model1 += [  nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),\n",
105 |         "                        norm_layer(out_features),\n",
106 |         "                        nn.ReLU(inplace=True) ]\n",
107 |         "            in_features = out_features\n",
108 |         "            out_features = in_features*2\n",
109 |         "        self.model1 = nn.Sequential(*model1)\n",
110 |         "\n",
111 |         "        model2 = []\n",
112 |         "        # Residual blocks\n",
113 |         "        for _ in range(n_residual_blocks):\n",
114 |         "            model2 += [ResidualBlock(in_features)]\n",
115 |         "        self.model2 = nn.Sequential(*model2)\n",
116 |         "\n",
117 |         "        # Upsampling\n",
118 |         "        model3 = []\n",
119 |         "        out_features = in_features//2\n",
120 |         "        for _ in range(2):\n",
121 |         "            model3 += [  nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),\n",
122 |         "                        norm_layer(out_features),\n",
123 |         "                        nn.ReLU(inplace=True) ]\n",
124 |         "            in_features = out_features\n",
125 |         "            out_features = in_features//2\n",
126 |         "        self.model3 = nn.Sequential(*model3)\n",
127 |         "\n",
128 |         "        # Output layer\n",
129 |         "        model4 = [  nn.ReflectionPad2d(3),\n",
130 |         "                        nn.Conv2d(64, output_nc, 7)]\n",
131 |         "        if sigmoid:\n",
132 |         "            model4 += [nn.Sigmoid()]\n",
133 |         "\n",
134 |         "        self.model4 = nn.Sequential(*model4)\n",
135 |         "\n",
136 |         "    def forward(self, x, cond=None):\n",
137 |         "        out = self.model0(x)\n",
138 |         "        out = self.model1(out)\n",
139 |         "        out = self.model2(out)\n",
140 |         "        out = self.model3(out)\n",
141 |         "        out = self.model4(out)\n",
142 |         "\n",
143 |         "        return out\n",
144 |         "\n",
145 |         "model = Generator(3, 1, 3)\n",
146 |         "model.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))\n",
147 |         "model.eval()\n",
148 |         "\n",
149 |         "def predict(input_img):\n",
150 |         "    input_img = Image.open(input_img)\n",
151 |         "    transform = transforms.Compose([transforms.Resize(256, Image.BICUBIC), transforms.ToTensor()])\n",
152 |         "    input_img = transform(input_img)\n",
153 |         "    input_img = torch.unsqueeze(input_img, 0)\n",
154 |         "\n",
155 |         "    with torch.no_grad():\n",
156 |         "        drawing = model(input_img)[0].detach()                  \n",
157 |         "        torch.onnx.export(model, input_img, f=\"model.onnx\", export_params=True, opset_version=12, do_constant_folding=True, input_names=['input'], output_names=['output'], dynamic_axes={'input':{0:'batch', 2:'height', 3:'width'}, 'output':{0:'batch', 2:'height', 3:'width'}}) \n",
158 |         "    \n",
159 |         "    drawing = transforms.ToPILImage()(drawing)\n",
160 |         "    return drawing\n",
161 |         "\n",
162 |         "predict(\"lizard.png\")"
163 |       ]
164 |     }
165 |   ]
166 | }


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2022 josephrocca
 4 | 
 5 | Permission is hereby granted, free of charge, to any person obtaining a copy
 6 | of this software and associated documentation files (the "Software"), to deal
 7 | in the Software without restriction, including without limitation the rights
 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 | 
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 | 
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # Image to line art / sketch in JavaScript with ONNX
 2 | 
 3 | **Try it here**: https://josephrocca.github.io/image-to-line-art-js
 4 | 
 5 | This is a simple demo of a JavaScript port of this model:
 6 | 
 7 |  * https://github.com/carolineec/informative-drawings
 8 |  * https://huggingface.co/spaces/carolineec/informativedrawings
 9 | 
10 | ### Notes:
11 |  * ONNX model file (about 17mb) is served from this Hugging Face repo: https://huggingface.co/rocca/informative-drawings-line-art-onnx
12 |  * The notebook I used to convert the model to ONNX is here: https://colab.research.google.com/josephrocca/image-to-line-art-js/blob/main/Informative_Drawings_Line_Art_Generator_ONNX_conversion.ipynb
13 | 


--------------------------------------------------------------------------------
/github-pages-coop-coep-workaround.js:
--------------------------------------------------------------------------------
 1 | // NOTE: This file creates a service worker that cross-origin-isolates the page (read more here: https://web.dev/coop-coep/) which allows us to use wasm threads.
 2 | // Normally you would set the COOP and COEP headers on the server to do this, but Github Pages doesn't allow this, so this is a hack to do that.
 3 | 
 4 | /* Edited version of: coi-serviceworker v0.1.6 - Guido Zuidhof, licensed under MIT */
 5 | // From here: https://github.com/gzuidhof/coi-serviceworker
 6 | if(typeof window === 'undefined') {
 7 |   self.addEventListener("install", () => self.skipWaiting());
 8 |   self.addEventListener("activate", e => e.waitUntil(self.clients.claim()));
 9 | 
10 |   async function handleFetch(request) {
11 |     if(request.cache === "only-if-cached" && request.mode !== "same-origin") {
12 |       return;
13 |     }
14 |     
15 |     if(request.mode === "no-cors") { // We need to set `credentials` to "omit" for no-cors requests, per this comment: https://bugs.chromium.org/p/chromium/issues/detail?id=1309901#c7
16 |       request = new Request(request.url, {
17 |         cache: request.cache,
18 |         credentials: "omit",
19 |         headers: request.headers,
20 |         integrity: request.integrity,
21 |         destination: request.destination,
22 |         keepalive: request.keepalive,
23 |         method: request.method,
24 |         mode: request.mode,
25 |         redirect: request.redirect,
26 |         referrer: request.referrer,
27 |         referrerPolicy: request.referrerPolicy,
28 |         signal: request.signal,
29 |       });
30 |     }
31 |     
32 |     let r = await fetch(request).catch(e => console.error(e));
33 |     
34 |     if(r.status === 0) {
35 |       return r;
36 |     }
37 | 
38 |     const headers = new Headers(r.headers);
39 |     headers.set("Cross-Origin-Embedder-Policy", "credentialless"); // or: require-corp
40 |     headers.set("Cross-Origin-Opener-Policy", "same-origin");
41 |     
42 |     return new Response(r.body, { status: r.status, statusText: r.statusText, headers });
43 |   }
44 | 
45 |   self.addEventListener("fetch", function(e) {
46 |     e.respondWith(handleFetch(e.request)); // respondWith must be executed synchonously (but can be passed a Promise)
47 |   });
48 |   
49 | } else {
50 |   (async function() {
51 |     if(window.crossOriginIsolated !== false) return;
52 | 
53 |     let registration = await navigator.serviceWorker.register(window.document.currentScript.src).catch(e => console.error("COOP/COEP Service Worker failed to register:", e));
54 |     if(registration) {
55 |       console.log("COOP/COEP Service Worker registered", registration.scope);
56 | 
57 |       registration.addEventListener("updatefound", () => {
58 |         console.log("Reloading page to make use of updated COOP/COEP Service Worker.");
59 |         window.location.reload();
60 |       });
61 | 
62 |       // If the registration is active, but it's not controlling the page
63 |       if(registration.active && !navigator.serviceWorker.controller) {
64 |         console.log("Reloading page to make use of COOP/COEP Service Worker.");
65 |         window.location.reload();
66 |       }
67 |     }
68 |   })();
69 | }
70 | 
71 | // Code to deregister:
72 | // let registrations = await navigator.serviceWorker.getRegistrations();
73 | // for(let registration of registrations) {
74 | //   await registration.unregister();
75 | // }
76 | 


--------------------------------------------------------------------------------
/index.html:
--------------------------------------------------------------------------------
  1 | <!DOCTYPE html>
  2 | <html>
  3 |   <head>
  4 |     <title>Convert Image to Line Art / Sketch (with JS)</title>
  5 |     <script src="github-pages-coop-coep-workaround.js"></script>
  6 |   </head>
  7 |   <body>
  8 |     <script src="https://cdn.jsdelivr.net/npm/onnxruntime-web@1.10.0/dist/ort.js"></script>
  9 |     
 10 |     <div>
 11 |       Select one or more images: <input type="file" accept="image/*" multiple id="fileInputEl">
 12 |       <br>
 13 |       Then click here: <button id="startBtn" onclick="main()" disabled>loading model...</button>
 14 |       <br>
 15 |       Save all the resulting line art images: <button id="downloadBtn" disabled onclick="downloadOutputImages()">save</button>
 16 |     </div>
 17 |     <p><a href="https://github.com/josephrocca/image-to-line-art-js">github repo</a> - <a href="https://huggingface.co/rocca/informative-drawings-line-art-onnx">huggingface repo</a></p>
 18 |     
 19 |     <script>
 20 |       if(self.crossOriginIsolated) { // needs to be cross-origin-isolated to use wasm threads. you need to serve this html file with these two headers: https://web.dev/coop-coep/
 21 |         ort.env.wasm.numThreads = navigator.hardwareConcurrency
 22 |       }
 23 |       ort.env.wasm.proxy = true;
 24 |       
 25 |       if(!window.OffscreenCanvas) alert("Your browser doesn't support OffscreenCanvas - a browser feature that was standardized way back in 2018. Please use a modern browser like Chrome, Edge or Brave.");
 26 |       
 27 |       let onnxSession;
 28 |       (async function() {
 29 |         console.log("Downloading model... (see network tab for progress)");
 30 |         onnxSession = await ort.InferenceSession.create('https://huggingface.co/rocca/informative-drawings-line-art-onnx/resolve/main/model.onnx', { executionProviders: ["wasm"] });
 31 |         console.log("Model loaded.");
 32 |         startBtn.disabled = false;
 33 |         startBtn.textContent = "start";
 34 |       })();
 35 |       
 36 |       async function main() {
 37 |         if(!fileInputEl.files[0]) return alert("Please select some images first.");
 38 |         
 39 |         startBtn.disabled = true;
 40 |         startBtn.innerHTML = "processing...";
 41 |         
 42 |         let inputBlob;
 43 |         for(let inputBlob of fileInputEl.files) {
 44 |           let {data, width, height} = await blobToLinearRGBArray(inputBlob);
 45 | 
 46 |           const feeds = {'input': new ort.Tensor('float32', data, [1, 3, height, width])};
 47 | 
 48 |           console.log("Running inference...");
 49 |           let t = Date.now();
 50 |           const results = await onnxSession.run(feeds);
 51 |           console.log(`Finished in ${Date.now()-t}ms`);
 52 | 
 53 |           const out = results["output"]; // greyscale data tensor
 54 |           console.log(`results:`, out);
 55 | 
 56 |           let blob = await linearGreyscaleArrayToBlob(out.data, {width:out.dims[3], height:out.dims[2]});
 57 |           let blobUrl = URL.createObjectURL(blob);
 58 |           let imgEl = document.createElement("img");
 59 |           imgEl.dataset.name = inputBlob.name;
 60 |           imgEl.src = blobUrl;
 61 |           document.body.appendChild(imgEl);
 62 |         }
 63 |         
 64 |         startBtn.disabled = false;
 65 |         startBtn.innerHTML = "start";
 66 |         downloadBtn.disabled = false;
 67 |       }
 68 |       
 69 |       async function getBlobImageDims(blob) {
 70 |         let img = await createImageBitmap(blob);
 71 |         return {width:img.width, height:img.height};
 72 |       }
 73 |       
 74 |       async function linearGreyscaleArrayToBlob(linearArr, dims) {
 75 |         let dataArray = [];
 76 |         for(let i = 0; i < linearArr.length; i++) {
 77 |           dataArray.push(linearArr[i]*255); // R
 78 |           dataArray.push(linearArr[i]*255); // G
 79 |           dataArray.push(linearArr[i]*255); // B
 80 |           dataArray.push(1*255);            // A
 81 |         }
 82 |         //let imageData = ctx.createImageData(256, 256);
 83 |         let imageData = new ImageData(new Uint8ClampedArray(dataArray), dims.width, dims.height);
 84 |         let canvas = new OffscreenCanvas(dims.width, dims.height);
 85 |         let ctx = canvas.getContext("2d");
 86 |         ctx.putImageData(imageData, 0, 0);
 87 |         return canvas.convertToBlob({type:"image/jpg"});
 88 |       }
 89 |       
 90 |       // this function resizes to 256px width
 91 |       async function blobToLinearRGBArray(blob) {
 92 |         let img = await createImageBitmap(blob);
 93 |         let canvas = new OffscreenCanvas(img.width, img.height);
 94 |         let ctx = canvas.getContext("2d");
 95 |         ctx.drawImage(img, 0, 0, canvas.width, canvas.height);
 96 |         let imageData = ctx.getImageData(0, 0, canvas.width, canvas.height);
 97 | 
 98 |         let rgbData = [[], [], []]; // [r, g, b]
 99 |         // remove alpha and put into correct shape:
100 |         let d = imageData.data;
101 |         for(let i = 0; i < d.length; i += 4) { 
102 |           let x = (i/4) % canvas.width;
103 |           let y = Math.floor((i/4) / canvas.width)
104 |           if(!rgbData[0][y]) rgbData[0][y] = [];
105 |           if(!rgbData[1][y]) rgbData[1][y] = [];
106 |           if(!rgbData[2][y]) rgbData[2][y] = [];
107 |           rgbData[0][y][x] = d[i+0]/255;
108 |           rgbData[1][y][x] = d[i+1]/255;
109 |           rgbData[2][y][x] = d[i+2]/255;
110 |         }
111 |         rgbData = Float32Array.from(rgbData.flat().flat());
112 |         return {data:rgbData, width:img.width, height:img.height};
113 |       }
114 |       
115 |       async function downloadOutputImages() {
116 |         let directoryHandle = await window.showDirectoryPicker();
117 |         for(let img of [...document.querySelectorAll("img")]) {
118 |           let url = img.src;
119 |           let name = img.dataset.name+".lines.jpg";
120 |           let fileHandle = await directoryHandle.getFileHandle(name, {create:true});
121 |           let response = await fetch(url);
122 |           await response.body.pipeTo(await fileHandle.createWritable());
123 |         }
124 |       }
125 |     </script>
126 |   </body>
127 | </html>
128 | 


--------------------------------------------------------------------------------