feat: ai

view/src/ai.ts

@@ -1,19 +1,18 @@
 import '@tensorflow/tfjs-backend-webgl';
+import "@tensorflow/tfjs-backend-cpu"
 import * as tfconv from '@tensorflow/tfjs-converter';
 import * as deeplab from '@tensorflow-models/deeplab';
+import { getLabels, getColormap, getURL, SemanticSegmentation, toSegmentationImage } from '@tensorflow-models/deeplab';
 
+const base = 'cityscapes';        // set to your preferred model, out of `pascal`,
 const createModel = async () => {
-
-  console.log("Ezzzz")
-  const base = 'pascal';        // set to your preferred model, out of `pascal`,
   // `cityscapes` and `ade20k`
   const quantizationBytes = 2;  // either 1, 2 or 4
   // use the getURL utility function to get the URL to the pre-trained weights
-  const modelUrl = deeplab.getURL(base, quantizationBytes);
+  const modelUrl = getURL(base, quantizationBytes);
   const rawModel = await tfconv.loadGraphModel(modelUrl);
-  const modelName = 'pascal';  // set to your preferred model, out of `pascal`,
-  // `cityscapes` and `ade20k`
-  return new deeplab.SemanticSegmentation(rawModel);
+  const modelName = 'pascal';  // set to your preferred model, out of `pascal`, `cityscapes` and `ade20k`
+  return new SemanticSegmentation(rawModel, modelName);
 };
 
 const model = createModel();
@@ -21,16 +20,15 @@ model.then(() => console.log(`Loaded the model successfully!`));
 
 self.addEventListener('message', async (e) => {
 
-  const { pixels, width, height } = e.data;
-
-  console.log(pixels, width, height)
+  const { pixels, width, height, i } = e.data;
 
   var array = new Uint8ClampedArray(pixels);
 
   var image = new ImageData(array, width, height);
 
-  console.log("model", await (await model).segment(image))
+  const result = await (await model).segment(image);
 
-  console.log(e);
+  //@ts-ignore
+  self.postMessage({ i, result })
 
 });

view/src/components/Editor/CustomOrbit.ts

@@ -1,5 +1,5 @@
+//@ts-nocheck
 import { Vec2, Vec3 } from 'ogl';
-
 const STATE = { NONE: -1, ROTATE: 0, DOLLY: 1, PAN: 2, DOLLY_PAN: 3 };
 const tempVec3 = new Vec3();
 const tempVec2a = new Vec2();

view/src/components/Editor/OrbView.ts

@@ -113,6 +113,7 @@ export default (image: Image, canvas2D: CanvasRenderingContext2D, canvas3D: HTML
   }
 
   function updateOverlay() {
+    //@ts-ignore
     overlayProgram.uniforms.tMap.value.image = canvas2D.getImageData(0, 0, image.width, image.height).data;
   }
 
@@ -126,6 +127,7 @@ export default (image: Image, canvas2D: CanvasRenderingContext2D, canvas3D: HTML
   }
 
   function setOpacity(o) {
+    //@ts-ignore
     overlayProgram.uniforms.opacity.value = o / 100;
   }
 

view/src/components/Editor/Painter.svelte

@@ -1,8 +1,9 @@
 <script lang="ts">
-  import { bufToImageUrl } from "helpers";
+  import { AI, bufToImageUrl } from "helpers";
   import { images as imageStore } from "stores";
   import OrbView from "./OrbView";
   import { onMount } from "svelte";
+  import Toast from "components/Toast";
 
   export let image: Image;
   export let activeTool = "pan";
@@ -20,7 +21,17 @@
   let cx2: CanvasRenderingContext2D;
 
   let orb: ReturnType<typeof OrbView>;
-  $: if (activeTool && orb) orb.setTool(activeTool);
+  $: if (activeTool) {
+    if (orb) {
+      orb.setTool(activeTool);
+    }
+    if (activeTool === "clear") {
+      cx1.clearRect(0, 0, image.width, image.height);
+      cx2.clearRect(0, 0, image.width, image.height);
+      saveToImage();
+      activeTool = "brush";
+    }
+  }
   $: if (layerOpacity && orb) orb.setOpacity(layerOpacity);
 
   let mode = "2d";
@@ -52,6 +63,33 @@
   let isSpacePressed = false;
   let lastActiveTool;
 
+  function scaleImageData(imageData, scale) {
+    var newCanvas = document.createElement("canvas");
+    newCanvas.width = imageData.width;
+    newCanvas.height = imageData.height;
+
+    newCanvas.getContext("2d").putImageData(imageData, 0, 0);
+
+    // Second canvas, for scaling
+    var scaleCanvas = document.createElement("canvas");
+    scaleCanvas.width = imageData.width * scale;
+    scaleCanvas.height = imageData.height * scale;
+
+    var scaleCtx = scaleCanvas.getContext("2d");
+
+    scaleCtx.scale(scale, scale);
+    scaleCtx.drawImage(newCanvas, 0, 0);
+
+    var scaledImageData = scaleCtx.getImageData(
+      0,
+      0,
+      scaleCanvas.width * scale,
+      scaleCanvas.height * scale
+    );
+
+    return scaledImageData;
+  }
+
   const saveToImage = () => {
     const imageData = cx1.getImageData(0, 0, image.width, image.height);
     image.overlayData = imageData.data;
@@ -225,6 +263,35 @@
     wrapperWidth = box.width;
   }
 
+  let aiState = "AI";
+  async function handleAi() {
+    if (aiState !== "AI") return;
+
+    if (
+      !(await Toast.confirm(
+        "For now this action will overwrite existing mask, continue?"
+      ))
+    ) {
+      return;
+    }
+
+    aiState = "loading...";
+    const res: any = await AI.analyze(image);
+    aiState = "Finished";
+    setTimeout(() => {
+      aiState = "AI";
+    }, 1000);
+
+    console.log(res);
+
+    var iData = new ImageData(res.segmentationMap, res.width, res.height);
+
+    const scale = image.width / res.width;
+    const s = scaleImageData(iData, scale);
+
+    cx1.putImageData(s, 0, 0);
+  }
+
   onMount(() => {
     canvas.width = image.width;
     canvas.height = image.height;
@@ -274,6 +341,10 @@
     {mode}
   </button>
 
+  <button id="ai" on:click={handleAi}>
+    {aiState}
+  </button>
+
   {#if activeTool === "brush" || activeTool === "erasor"}
     <div
       id="cursor"
@@ -313,6 +384,16 @@
     left: 0px;
     z-index: 1001;
   }
+
+  #ai {
+    position: absolute;
+    top: 60px;
+    right: 10px;
+    z-index: 1001;
+    cursor: pointer;
+    pointer-events: all;
+  }
+
   #mode {
     position: absolute;
     top: 10px;

view/src/components/Editor/ToolBox.svelte

@@ -1,4 +1,6 @@
 <script>
+  import Toast from "components/Toast";
+
   export let activeTool = "pan";
 
   const tools = ["pan", "brush", "erasor", "polygon"];
@@ -10,6 +12,17 @@
   </button>
 {/each}
 
+<button
+  on:click={async () => {
+    const clear = await Toast.confirm("Clear Mask?");
+    if (clear) {
+      activeTool = "clear";
+    }
+  }}
+>
+  clear
+</button>
+
 <style>
   button.active {
     background-color: red;

view/src/helpers/AI.ts

@@ -1,3 +1,5 @@
+import { bufToImageUrl } from ".";
+
 const worker = new Worker("build/ai-worker.js");
 
 let i = 0;
@@ -6,14 +8,33 @@ let cb = {};
 
 worker.addEventListener("message", ev => {
   if (ev.data.i in cb) {
-    cb[ev.data.i](ev.data.res);
+    cb[ev.data.i](ev.data.result);
+    delete cb[ev.data.i]
+    delete ev.data.i;
   }
 })
 
 const analyze = (img: Image) => new Promise((res, rej) => {
   i++;
   const _i = i;
-  worker.postMessage({ i: _i, pixels: img.data, width: img.width, height: img.height });
+
+  const canvas = document.createElement("canvas");
+  canvas.width = img.width;
+  canvas.height = img.height;
+
+  const cx = canvas.getContext("2d");
+  const image = document.createElement("img");
+  image.onload = () => {
+
+    cx.drawImage(image, 0, 0);
+    const pixels = cx.getImageData(0, 0, image.width, image.height).data
+
+    worker.postMessage({ i: _i, pixels, width: img.width, height: img.height });
+  }
+
+  image.src = bufToImageUrl(img.data, img.type);
+
+
   cb[_i] = res;
 });
 

view/src/helpers/BuffToImg.ts

@@ -3,6 +3,5 @@ const urlCreator = window.URL || window.webkitURL;
 export default (buff: ArrayBuffer, mimeType: string) => {
   const blob = new Blob([buff], { type: mimeType });
   const imageUrl = urlCreator.createObjectURL(blob);
-
   return imageUrl;
 }

view/src/routes/list.svelte

@@ -8,6 +8,7 @@
   import Toast from "../components/Toast";
   import Analyzer from "../components/Analyzer.svelte";
   import { onMount } from "svelte";
+  import ToastWrapper from "components/Toast/ToastWrapper.svelte";
 
   const imageStore: Writable<Image[]> = imageData.store;
 
@@ -59,7 +60,6 @@
                 showAnalyzerIndex = undefined;
               } else {
                 showAnalyzerIndex = i;
-                //AI.analyze(img);
               }
             } else {
               showAnalyzerIndex = undefined;
@@ -70,7 +70,11 @@
         <button on:click={() => route.set("editor/" + img.id)}>edit</button>
         <button
           on:click={() => {
+            if (img.overlayData && img.overlayData.byteLength) {
               downloadImage(img);
+            } else {
+              Toast.warn("Image has no mask, cant download");
+            }
           }}>download mask</button
         >
         <button