Magic wand selection tool with OpenCV Python

magic_wand_opencv.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Magic wand selection with OpenCV Python.
The code highly credits to https://github.com/alkasm/magicwand, but this script only relies on tkinter for user interface.
Please run the code by simply calling: $ python magic_wand_opencv.py
Requirements: Python 3.9+, Pillow, opencv-python>=4.6
"""

import os
import sys

import cv2 as cv
import numpy as np
from PIL import Image, ImageTk
import tkinter as tk
from tkinter import filedialog


def _find_exterior_contours(img):
    ret = cv.findContours(img, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
    if len(ret) == 2:
        return ret[0]
    elif len(ret) == 3:
        return ret[1]
    raise Exception("Check the signature for `cv.findContours()`.")

class WandAreaApp(tk.Tk):
    def __init__(self):
        # Initialize window
        super().__init__()
        self.geometry('600x600')
        self.resizable(False, False)
        self.title('Magic Wand Area App')

        # Load/save image button / status bar
        self.top_frame = tk.Frame(self)
        self.top_frame.pack(side=tk.TOP)
        self.img_load_btn = tk.Button(self.top_frame, text="Load image",
                                      command=self.img_load_callback)
        self.img_load_btn.pack(side=tk.LEFT)
        self.img_save_btn = tk.Button(self.top_frame, text="Save image",
                                      command=self.img_save_callback, state=tk.DISABLED)
        self.img_save_btn.pack(side=tk.LEFT)
        self.status_txt = tk.StringVar()
        self.status_lbl = tk.Label(self.top_frame, textvariable=self.status_txt, width=480, anchor="w")
        self.status_lbl.pack(side=tk.RIGHT, expand=True)

        # Sliders
        self.slider_frame = tk.Frame(self)
        self.slider_frame.pack(side=tk.TOP, anchor="nw")
        self.slider_frame.columnconfigure(0, pad=3)
        self.slider_frame.columnconfigure(1, pad=3)
        self.slider_frame.columnconfigure(2, pad=3)
        self.slider_frame.columnconfigure(3, pad=3)
        self.slider_frame.rowconfigure(0, pad=3)
        self.slider_frame.rowconfigure(1, pad=3)
        tk.Label(self.slider_frame, text="Blur").grid(row=0, column=0)
        self.blur_slider = tk.Scale(self.slider_frame, from_=1, to=21, resolution=2, orient=tk.HORIZONTAL, length=250,
                                    command=self.callback_blur_slider)
        self.blur_slider.set(5)
        self.blur_slider.grid(row=0, column=1)
        self.gaussian_var = tk.IntVar()
        self.gaussian_chk = tk.Checkbutton(self.slider_frame, text='Gaussian Blur',variable=self.gaussian_var, onvalue=1, offvalue=0, command=self.callback_chk)
        self.gaussian_var.set(1)
        self.gaussian_chk.grid(row=0, column=2)
        self.denoise_var = tk.IntVar()
        self.denoise_chk = tk.Checkbutton(self.slider_frame, text='Denoise',variable=self.denoise_var, onvalue=1, offvalue=0, command=self.callback_chk)
        self.denoise_chk.grid(row=0, column=3)
        tk.Label(self.slider_frame, text="Tolerance").grid(row=1, column=0)
        self.tole_slider = tk.Scale(self.slider_frame, from_=0, to=30, resolution=1, orient=tk.HORIZONTAL, length=250)
        self.tole_slider.set(5)
        self.tole_slider.grid(row=1, column=1)

        # Image canvas
        self.canvas_height = 500
        self.canvas_width = 600
        self.img_canvas = tk.Canvas(self, height=self.canvas_height, width=self.canvas_width)
        self.img_canvas.pack(side=tk.TOP)
        self.img_container = self.img_canvas.create_image(0, 0, anchor="nw")
        self.img_canvas.bind("<Button-1>", self.callback_img_click)
        self.img_canvas["state"] = tk.DISABLED

        # Initial setting
        self.img_fpath = ""
        self.img_orig = np.zeros((self.canvas_height, self.canvas_width, 3), dtype=np.uint8)
        self.mask = np.zeros((self.canvas_height, self.canvas_width), dtype=np.uint8)
        self._flood_mask = np.zeros((self.canvas_height + 2, self.canvas_width + 2), dtype=np.uint8)
        self._flood_fill_flags = (
            4 | cv.FLOODFILL_FIXED_RANGE | cv.FLOODFILL_MASK_ONLY | 255 << 8
        )
        self.blur_image()
        self.set_status_text()

    def img_load_callback(self):
        self.img_fpath = filedialog.askopenfilename()
        if self.img_fpath:
            self.img_orig = cv.imread(self.img_fpath)
            h, w = self.img_orig.shape[:2]
            self.mask = np.zeros((h, w), dtype=np.uint8)
            self._flood_mask = np.zeros((h + 2, w + 2), dtype=np.uint8)
            self.blur_image()
            self.display_image()
            self.set_status_text("Loaded image: " + os.path.basename(self.img_fpath))

            self.img_save_btn["state"] = tk.NORMAL
            self.img_canvas["state"] = tk.NORMAL

    def img_save_callback(self):
        save_fpath = filedialog.asksaveasfilename(
            filetypes=(
                ("JPG files", "*.jpg"),
                ("PNG files", "*.png"),
                ("All files", "*.*"),
            )
        )
        img_output = self.img_disp.copy()
        mask_cnt = np.sum(self.mask > 0)
        img_output = cv.putText(img_output, f'(Pixels: {mask_cnt})',
                                (0, img_output.shape[0]-20),
                                cv.FONT_HERSHEY_PLAIN, 1, (0, 0, 255), 1, cv.LINE_4)
        cv.imwrite(save_fpath, img_output)

    def set_status_text(self, text=None):
        if text:
            self.status_txt.set(text)
            return
        if not self.img_fpath:
            self.status_txt.set("Please select an image")
            return
        text = os.path.basename(self.img_fpath)
        mask_cnt = np.sum(self.mask > 0)
        text += f' ({mask_cnt} pixels selected)'
        self.status_txt.set(text)

    # Gaussian blur and/or denosing
    def blur_image(self):
        k = self.blur_slider.get()
        self.img_blur = self.img_orig.copy()
        if self.denoise_var.get():
            self.img_blur = cv.fastNlMeansDenoising(self.img_blur, h=k)            
        if self.gaussian_var.get():
            self.img_blur = cv.GaussianBlur(self.img_blur, (k, k), 0)

    # Display image
    def display_image(self):
        # Mask image
        self.img_disp = self.img_blur.copy()
        contours = _find_exterior_contours(self.mask)
        self.img_disp = cv.drawContours(self.img_disp, contours, -1, color=(255,) * 3, thickness=-1)
        self.img_disp = cv.addWeighted(self.img_blur, 0.75, self.img_disp, 0.25, 0)
        self.img_disp = cv.drawContours(self.img_disp, contours, -1, color=(255,) * 3, thickness=1)

        # Display image
        self.disp_image = ImageTk.PhotoImage(Image.fromarray(self.img_disp))
        self.img_canvas.itemconfig(self.img_container, image=self.disp_image)

    # Slider callback
    def callback_blur_slider(self, _):
        self.blur_image()
        self.display_image()
        pass

    # Checkbox callback -> same as slider callback
    def callback_chk(self):
        self.callback_blur_slider(None)

    # Click callback
    def callback_img_click(self, e):
        # print(f"x: {e.x}, y: {e.y}, state: {e.state}")
        x, y = e.x, e.y
        ctrl  = (e.state & 0x4) != 0
        shift = (e.state & 0x1) != 0
        tole = self.tole_slider.get()

        self._flood_mask[:] = 0
        cv.floodFill(
            self.img_blur,
            self._flood_mask,
            (x, y),
            0,
            (tole,) * 3,
            (tole,) * 3,
            self._flood_fill_flags,
        )
        flood_mask = self._flood_mask[1:-1, 1:-1].copy()

        if ctrl & shift:
            self.mask = cv.bitwise_and(self.mask, flood_mask)
        elif shift:
            self.mask = cv.bitwise_or(self.mask, flood_mask)
        elif ctrl:
            self.mask = cv.bitwise_and(self.mask, cv.bitwise_not(flood_mask))
        else:
            self.mask = flood_mask

        self.display_image()
        self.set_status_text()


if __name__ == "__main__":
    app = WandAreaApp()
    app.mainloop()