feat(gui): implement mask painting, flood fill

gui/src/api/client.ts

@@ -38,8 +38,9 @@ export interface Txt2ImgParams extends BaseImgParams {
 
 export type Txt2ImgResponse = Required<Txt2ImgParams>;
 
-export interface InpaintParams extends Img2ImgParams {
+export interface InpaintParams extends BaseImgParams {
   mask: Blob;
+  source: File;
 }
 
 export interface OutpaintParams extends Img2ImgParams {
@@ -68,6 +69,10 @@ export interface ApiClient {
 
 export const STATUS_SUCCESS = 200;
 
+export function equalResponse(a: ApiResponse, b: ApiResponse): boolean {
+  return a.output === b.output;
+}
+
 export function joinPath(...parts: Array<string>): string {
   return parts.join('/');
 }

gui/src/components/Inpaint.tsx

@@ -1,9 +1,11 @@
+/* eslint-disable @typescript-eslint/no-magic-numbers */
 import { doesExist, mustExist } from '@apextoaster/js-utils';
+import { FormatColorFill, Gradient } from '@mui/icons-material';
 import { Box, Button, Stack } from '@mui/material';
 import * as React from 'react';
 import { useMutation, useQuery } from 'react-query';
 
-import { ApiClient, ApiResponse, BaseImgParams } from '../api/client.js';
+import { ApiClient, ApiResponse, BaseImgParams, equalResponse } from '../api/client.js';
 import { Config, CONFIG_DEFAULTS, STALE_TIME } from '../config.js';
 import { SCHEDULER_LABELS } from '../strings.js';
 import { ImageCard } from './ImageCard.js';
@@ -14,6 +16,38 @@ import { QueryList } from './QueryList.js';
 
 const { useEffect, useRef, useState } = React;
 
+export const FULL_CIRCLE = 2 * Math.PI;
+
+export const COLORS = {
+  black: 0,
+  white: 255,
+};
+
+export function floodBelow(n: number): number {
+  if (n < 224) {
+    return COLORS.black;
+  } else {
+    return COLORS.white;
+  }
+}
+
+export function floodAbove(n: number): number {
+  if (n > 32) {
+    return COLORS.white;
+  } else {
+    return COLORS.black;
+  }
+}
+
+export function grayToRGB(n: number): string {
+  return `rgb(${n.toFixed(0)},${n.toFixed(0)},${n.toFixed(0)})`;
+}
+
+export interface Point {
+  x: number;
+  y: number;
+}
+
 export interface InpaintProps {
   client: ApiClient;
   config: Config;
@@ -35,7 +69,6 @@ export function Inpaint(props: InpaintProps) {
           model,
           platform,
           scheduler,
-          strength,
           mask,
           source: mustExist(source),
         }));
@@ -63,39 +96,60 @@ export function Inpaint(props: InpaintProps) {
     }
   }
 
+  function grayscaleMask() {
+    const canvas = mustExist(canvasRef.current);
+    const ctx = mustExist(canvas.getContext('2d'));
+    const image = ctx.getImageData(0, 0, canvas.width, canvas.height);
+    const pixels = image.data;
+
+    for (let x = 0; x < canvas.width; ++x) {
+      for (let y = 0; y < canvas.height; ++y) {
+        const i = (y * canvas.width * 4) + (x * 4);
+        const hue = (pixels[i] + pixels[i + 1] + pixels[i + 2]) / 3;
+        pixels[i] = hue;
+        pixels[i + 1] = hue;
+        pixels[i + 2] = hue;
+      }
+    }
+
+    ctx.putImageData(image, 0, 0);
+  }
+
+  function floodMask(flooder: (n: number) => number) {
+    const canvas = mustExist(canvasRef.current);
+    const ctx = mustExist(canvas.getContext('2d'));
+    const image = ctx.getImageData(0, 0, canvas.width, canvas.height);
+    const pixels = image.data;
+
+    for (let x = 0; x < canvas.width; ++x) {
+      for (let y = 0; y < canvas.height; ++y) {
+        const i = (y * canvas.width * 4) + (x * 4);
+        const hue = (pixels[i] + pixels[i + 1] + pixels[i + 2]) / 3;
+        const final = flooder(hue);
+
+        pixels[i] = final;
+        pixels[i + 1] = final;
+        pixels[i + 2] = final;
+      }
+    }
+
+    ctx.putImageData(image, 0, 0);
+  }
+
   const upload = useMutation(uploadSource);
   const schedulers = useQuery('schedulers', async () => client.schedulers(), {
     staleTime: STALE_TIME,
   });
 
-  interface Point {
-    x: number;
-    y: number;
-  };
-
   // eslint-disable-next-line @typescript-eslint/ban-types, no-null/no-null
   const canvasRef = useRef<HTMLCanvasElement | null>(null);
   const [clicks, setClicks] = useState<Array<Point>>([]);
-  useEffect(() => {
-    const canvas = mustExist(canvasRef.current);
-    const ctx = mustExist(canvas.getContext('2d'));
-    ctx.strokeStyle = 'black';
 
-    for (const click of clicks) {
+  const [painting, setPainting] = useState(false);
-      // eslint-disable-next-line no-console
+  const [brushColor, setBrushColor] = useState(0);
-      console.log(click.x, click.y, canvas.width, canvas.height);
+  const [brushSize, setBrushSize] = useState(4);
-
-      ctx.beginPath();
-      // eslint-disable-next-line @typescript-eslint/no-magic-numbers
-      ctx.arc(click.x, click.y, 8, 0, 2 * Math.PI);
-      ctx.fill();
-    }
-
-    clicks.length = 0;
-  }, [clicks.length]);
 
   const [source, setSource] = useState<File>();
-  const [strength, setStrength] = useState(CONFIG_DEFAULTS.strength.default);
   const [params, setParams] = useState<BaseImgParams>({
     cfg: CONFIG_DEFAULTS.cfg.default,
     seed: CONFIG_DEFAULTS.seed.default,
@@ -104,6 +158,20 @@ export function Inpaint(props: InpaintProps) {
   });
   const [scheduler, setScheduler] = useState(config.default.scheduler);
 
+  useEffect(() => {
+    const canvas = mustExist(canvasRef.current);
+    const ctx = mustExist(canvas.getContext('2d'));
+    ctx.fillStyle = grayToRGB(brushColor);
+
+    for (const click of clicks) {
+      ctx.beginPath();
+      ctx.arc(click.x, click.y, brushSize, 0, FULL_CIRCLE);
+      ctx.fill();
+    }
+
+    clicks.length = 0;
+  }, [clicks.length]);
+
   return <Box>
     <Stack spacing={2}>
       <Stack direction='row' spacing={2}>
@@ -136,24 +204,75 @@ export function Inpaint(props: InpaintProps) {
             y: event.clientY - bounds.top,
           }]);
         }}
+        onMouseDown={() => {
+          setPainting(true);
+        }}
+        onMouseLeave={() => {
+          setPainting(false);
+        }}
+        onMouseOut={() => {
+          setPainting(false);
+        }}
+        onMouseUp={() => {
+          setPainting(false);
+        }}
+        onMouseMove={(event) => {
+          if (painting) {
+            const canvas = mustExist(canvasRef.current);
+            const bounds = canvas.getBoundingClientRect();
+
+            setClicks([...clicks, {
+              x: event.clientX - bounds.left,
+              y: event.clientY - bounds.top,
+            }]);
+          }
+        }}
       />
+      <Stack direction='row' spacing={4}>
+        <NumericField
+          decimal
+          label='Brush Shade'
+          min={0}
+          max={255}
+          step={1}
+          value={brushColor}
+          onChange={(value) => {
+            setBrushColor(value);
+          }}
+        />
+        <NumericField
+          decimal
+          label='Brush Size'
+          min={4}
+          max={64}
+          step={1}
+          value={brushSize}
+          onChange={(value) => {
+            setBrushSize(value);
+          }}
+        />
+        <Button
+          startIcon={<FormatColorFill htmlColor='black' />}
+          onClick={() => floodMask(floodBelow)}>
+          Gray to black
+        </Button>
+        <Button
+          startIcon={<Gradient />}
+          onClick={() => grayscaleMask()}>
+          Grayscale
+        </Button>
+        <Button
+          startIcon={<FormatColorFill htmlColor='white' sx={{ bgcolor: 'text.primary' }} />}
+          onClick={() => floodMask(floodAbove)}>
+            Gray to white
+        </Button>
+      </Stack>
       <ImageControl params={params} onChange={(newParams) => {
         setParams(newParams);
       }} />
-      <NumericField
-        decimal
-        label='Strength'
-        min={CONFIG_DEFAULTS.strength.min}
-        max={CONFIG_DEFAULTS.strength.max}
-        step={CONFIG_DEFAULTS.strength.step}
-        value={strength}
-        onChange={(value) => {
-          setStrength(value);
-        }}
-      />
       <Button onClick={() => upload.mutate()}>Generate</Button>
       <MutationHistory result={upload} limit={4} element={ImageCard}
-        isEqual={(a, b) => a.output === b.output}
+        isEqual={equalResponse}
       />
     </Stack>
   </Box>;