feat: add noise feature (#11)

fal-ai · Mar 9, 2024 · a385ec3 · a385ec3
1 parent 2eef279
commit a385ec3
Show file tree

Hide file tree

Showing 3 changed files with 124 additions and 13 deletions.
diff --git a/app/page.tsx b/app/page.tsx
@@ -17,7 +17,8 @@ import {
 } from "@/components/model-dropdown";
 import { Button } from "@/components/ui/button";
 import { LoaderIcon } from "lucide-react";
-import { resizeImage } from "@/lib/image";
+import { resizeAndNoiseImage } from "@/lib/image";
+import Switch from "@/components/ui/switch";
 
 fal.config({ proxyUrl: "/api/proxy" });
 
@@ -29,11 +30,16 @@ interface ModelResult {
 
 type CompareMode = "original" | "model";
 
+const NOISE_LEVEL = 0.2;
+const MAX_IMAGE_SIZE = 512;
+
 export default function UpscalerBattleground() {
   const [mode, setMode] = useState<CompareMode>("original");
   const [position, setPosition] = useState<number>(50);
   const [originalImage, setOriginalImage] = useState<string | null>(null);
-  const [imageFile, setImageFile] = useState<File | null>(null);
+  const [rawImageFile, setRawImageFile] = useState<File | Blob | null>(null);
+  const [imageFile, setImageFile] = useState<File | Blob | null>(null);
+  const [noise, setNoise] = useState<boolean>(true);
 
   const [firstModelLoading, setFirstModelLoading] = useState<boolean>(false);
   const [secondModelLoading, setSecondModelLoading] = useState<boolean>(false);
@@ -54,10 +60,9 @@ export default function UpscalerBattleground() {
 
     let inferenceTime;
 
-    const resizedImage = await resizeImage(file);
     const result: Record<string, any> = await fal.subscribe(firstModel.model, {
       input: {
-        image_url: resizedImage,
+        image_url: file,
         ...(firstModel.meta || {}),
       },
       logs: true,
@@ -86,10 +91,9 @@ export default function UpscalerBattleground() {
 
     let inferenceTime;
 
-    const resizedImage = await resizeImage(file);
     const result: Record<string, any> = await fal.subscribe(secondModel.model, {
       input: {
-        image_url: resizedImage,
+        image_url: file,
         ...(secondModel.meta || {}),
       },
       logs: true,
@@ -111,8 +115,34 @@ export default function UpscalerBattleground() {
     setSecondModelLoading(false);
   };
 
-  const handleImageSelect = (e: React.ChangeEvent<HTMLInputElement>) => {
-    const image = e.target?.files?.[0];
+  const handleNoise = async (checked) => {
+    const noise = checked;
+    setNoise(noise);
+    if (!rawImageFile) return;
+
+    const newImage = await resizeAndNoiseImage(
+      rawImageFile as File,
+      MAX_IMAGE_SIZE,
+      noise ? NOISE_LEVEL : 0
+    );
+
+    const blobUrl = URL.createObjectURL(newImage);
+    setImageFile(newImage);
+    setOriginalImage(blobUrl);
+  };
+
+  const handleImageSelect = async (e: React.ChangeEvent<HTMLInputElement>) => {
+    let image: File | Blob | undefined = e.target?.files?.[0];
+
+    setRawImageFile(image || null);
+
+    if (image) {
+      image = await resizeAndNoiseImage(
+        image as File,
+        MAX_IMAGE_SIZE,
+        noise ? NOISE_LEVEL : 0
+      );
+    }
 
     setFirstModelOutput(null);
     setSecondModelOutput(null);
@@ -174,6 +204,24 @@ export default function UpscalerBattleground() {
                     </span>
                   </div>
                 </div>
+                <div className="h-10 flex items-center justify-center">
+                  <label
+                    htmlFor="noise"
+                    className="h-6 relative text-neutral-500 dark:text-neutral-400 uppercase text-xs flex items-center"
+                  >
+                    Noise:
+                    <Switch
+                      disabled={firstModelLoading || secondModelLoading}
+                      defaultChecked={noise}
+                      className={cn(
+                        "ml-1",
+                        firstModelLoading && "cursor-not-allowed"
+                      )}
+                      id="noise"
+                      onCheckedChange={handleNoise}
+                    />
+                  </label>
+                </div>
                 <div className="flex flex-col md:flex-row space-y-2 md:space-y-0 text-sm items-center justify-center">
                   <label className="text-neutral-500 md:mr-1 dark:text-neutral-400 uppercase text-xs">
                     Image:
@@ -190,6 +238,7 @@ export default function UpscalerBattleground() {
                   />
                 </div>
               </div>
+
               <Button
                 size="lg"
                 className="mx-auto md:mx-0"

diff --git a/components/ui/switch.tsx b/components/ui/switch.tsx
@@ -0,0 +1,19 @@
+import React from "react";
+import * as SwitchPrimitive from "@radix-ui/react-switch";
+import { cn } from "@/lib/utils";
+
+const Switch = ({ className, ...props }) => (
+  <SwitchPrimitive.Root
+    {...props}
+    className={cn(
+      "w-[42px] h-[25px] dark:bg-neutral-800 bg-neutral-200 rounded-full relative dark:data-[state=checked]:bg-white data-[state=checked]:bg-black outline-none cursor-default",
+      className
+    )}
+    // @ts-expect-error
+    style={{ "-webkit-tap-highlight-color": "rgba(0, 0, 0, 0)" }}
+  >
+    <SwitchPrimitive.Thumb className="block w-[21px] h-[21px] bg-white dark:bg-black rounded-full transition-transform duration-100 translate-x-0.5 will-change-transform data-[state=checked]:translate-x-[19px]" />
+  </SwitchPrimitive.Root>
+);
+
+export default Switch;
diff --git a/lib/image.ts b/lib/image.ts
@@ -1,4 +1,8 @@
-export async function resizeImage(file: File, max: number = 512): Promise<Blob> {
+export async function resizeAndNoiseImage(
+  file: File,
+  max: number = 512,
+  noiseLevel: number = 0.3
+): Promise<Blob> {
   const canvas = document.createElement("canvas");
   const ctx = canvas.getContext("2d");
 
@@ -11,18 +15,57 @@ export async function resizeImage(file: File, max: number = 512): Promise<Blob>
   const targetSize = Math.min(max, Math.max(image.width, image.height));
 
   const targetWidth =
-    image.width > image.height ? targetSize : Math.round((image.width / image.height) * targetSize);
+    image.width > image.height
+      ? targetSize
+      : Math.round((image.width / image.height) * targetSize);
   const targetHeight =
-    image.height > image.width ? targetSize : Math.round((image.height / image.width) * targetSize);
+    image.height > image.width
+      ? targetSize
+      : Math.round((image.height / image.width) * targetSize);
 
   canvas.width = targetWidth;
   canvas.height = targetHeight;
 
-  ctx.drawImage(image, 0, 0, image.width, image.height, 0, 0, targetWidth, targetHeight);
+  ctx.drawImage(
+    image,
+    0,
+    0,
+    image.width,
+    image.height,
+    0,
+    0,
+    targetWidth,
+    targetHeight
+  );
+
+  if (noiseLevel > 0) {
+    addNoise(ctx, targetWidth, targetHeight, noiseLevel);
+  }
 
   return await canvasToImage(canvas);
 }
 
+function addNoise(
+  ctx: CanvasRenderingContext2D,
+  width: number,
+  height: number,
+  noiseLevel: number
+) {
+  const imageData = ctx.getImageData(0, 0, width, height);
+  const data = imageData.data;
+
+  for (let i = 0; i < data.length; i += 4) {
+    const random = Math.random();
+    const noise = Math.round(random * noiseLevel * 255);
+
+    data[i] = Math.min(data[i] + noise, 255); // Red
+    data[i + 1] = Math.min(data[i + 1] + noise, 255); // Green
+    data[i + 2] = Math.min(data[i + 2] + noise, 255); // Blue
+  }
+
+  ctx.putImageData(imageData, 0, 0);
+}
+
 async function loadImage(file: File): Promise<HTMLImageElement> {
   return new Promise((resolve, reject) => {
     const reader = new FileReader();
@@ -48,7 +91,7 @@ async function canvasToImage(canvas: HTMLCanvasElement): Promise<Blob> {
         resolve(blob);
       },
       "image/jpeg",
-      0.7,
+      0.7
     );
   });
 }