Update tools/edit.py

tools/edit.py

@@ -1,38 +1,127 @@
-import torch
-from PIL import Image
-from diffusers import AutoPipelineForImage2Image
+import os
+from PIL import Image, ImageEnhance, ImageFilter
 
-pipe = None
 
+def edit_image(image_path: str, prompt: str):
+    """
+    Edit an existing image based on the user's prompt.
 
-def load_model():
-    global pipe
+    Parameters
+    ----------
+    image_path : str
+        Path to uploaded/current image
 
-    if pipe is None:
-        pipe = AutoPipelineForImage2Image.from_pretrained(
-            "stabilityai/sd-turbo",
-            torch_dtype=torch.float32
-        )
-        pipe.to("cpu")
+    prompt : str
+        User request
 
-    return pipe
+    Returns
+    -------
+    PIL.Image
+    """
 
+    if image_path is None:
+        raise ValueError("No image selected.")
 
-def edit_image(image_path: str, prompt: str):
+    if not os.path.exists(image_path):
+        raise FileNotFoundError(image_path)
 
-    if image_path is None:
-        return None
+    img = Image.open(image_path).convert("RGB")
+
+    prompt = prompt.lower()
+
+    # ---------------------------------
+    # Brightness
+    # ---------------------------------
+
+    if "bright" in prompt:
+
+        enhancer = ImageEnhance.Brightness(img)
+
+        img = enhancer.enhance(1.4)
+
+    # ---------------------------------
+    # Contrast
+    # ---------------------------------
+
+    elif "contrast" in prompt:
+
+        enhancer = ImageEnhance.Contrast(img)
+
+        img = enhancer.enhance(1.6)
+
+    # ---------------------------------
+    # Blur
+    # ---------------------------------
+
+    elif "blur" in prompt:
+
+        img = img.filter(ImageFilter.GaussianBlur(4))
+
+    # ---------------------------------
+    # Sharpen
+    # ---------------------------------
+
+    elif "sharpen" in prompt:
+
+        img = img.filter(ImageFilter.SHARPEN)
+
+    # ---------------------------------
+    # Black & White
+    # ---------------------------------
+
+    elif "black" in prompt or "white" in prompt:
+
+        img = img.convert("L").convert("RGB")
+
+    # ---------------------------------
+    # Grayscale
+    # ---------------------------------
+
+    elif "grayscale" in prompt:
+
+        img = img.convert("L").convert("RGB")
+
+    # ---------------------------------
+    # Flip
+    # ---------------------------------
+
+    elif "flip" in prompt:
+
+        img = img.transpose(Image.FLIP_LEFT_RIGHT)
+
+    # ---------------------------------
+    # Rotate
+    # ---------------------------------
+
+    elif "rotate" in prompt:
+
+        img = img.rotate(90, expand=True)
+
+    # ---------------------------------
+    # Enhance
+    # ---------------------------------
+
+    elif "enhance" in prompt:
+
+        img = img.filter(ImageFilter.DETAIL)
+
+    # ---------------------------------
+    # Anime / Pixar / Ghibli
+    # ---------------------------------
+
+    elif (
+        "anime" in prompt
+        or "pixar" in prompt
+        or "ghibli" in prompt
+        or "cartoon" in prompt
+    ):
+
+        # Placeholder
 
-    model = load_model()
+        print("Connect Image Editing API here")
 
-    image = Image.open(image_path).convert("RGB")
+    else:
 
-    result = model(
-        prompt=prompt,
-        image=image,
-        strength=0.75,
-        guidance_scale=3.0,
-        num_inference_steps=20
-    )
+        print("No matching edit command.")
 
-    return result.images[0]
+    return img