arte.py

from dataclasses import dataclass, field
from enum import Enum
import time
import typing
import cv2
import numpy as np
import pyautogui
import pyperclip

import mouse

class Mode(Enum):
    SCANLINE = "scanline",
    CONTOUR = "contour"

@dataclass
class Canvas:
    width: int = field(default=0)
    height: int = field(default=0)
    color_x: int = field(default=0)
    color_picker_x: int = field(default=0)
    color_picker_y: int = field(default=0)

class Albedo:
    def __init__(self, canvas: Canvas) -> None:
        self.canvas = canvas
        self.mouse = mouse.Mouse()

    def load_image(self, 
                   path: str) -> cv2.typing.MatLike:
        return cv2.imread(path)

    def apply_filters(self, 
                      img: cv2.typing.MatLike) -> cv2.typing.MatLike:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        return img
    
    def resize_and_crop(self, 
                        img: cv2.typing.MatLike) -> cv2.typing.MatLike:
        width = img.shape[1]
        height = img.shape[0]

        w_factor = width / self.canvas.width
        h_factor = height / self.canvas.height

        w_factor = 1 / w_factor if w_factor > 1 else 1 + (1 - w_factor)
        h_factor = 1 / h_factor if h_factor > 1 else 1 + (1 - h_factor)

        return cv2.resize(img, 
                          None, 
                          fx=h_factor, 
                          fy=h_factor, 
                          interpolation=cv2.INTER_LANCZOS4)
    
    def find_contours(self, 
                   img: cv2.typing.MatLike) -> tuple[typing.Sequence[cv2.typing.MatLike], cv2.typing.MatLike]:
        edges = cv2.Canny(img, 200, 290, apertureSize=3)
        contours, _ = cv2.findContours(edges, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
        return contours # type: ignore
    
    def draw_contours(self,
                      img: cv2.typing.MatLike,
                      contours: typing.Any) -> cv2.typing.MatLike:
        imagem_branca = np.zeros_like(img)
        return cv2.drawContours(image=imagem_branca, 
                                contours=contours, 
                                contourIdx=-1, 
                                color=(255, 255, 255), 
                                thickness=cv2.FILLED)
    
    def show(self, img: cv2.typing.MatLike) -> None:
        cv2.imshow("display", img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    def prepare_image(self, 
                      img: cv2.typing.MatLike) -> tuple[cv2.typing.MatLike, 
                                                        tuple[typing.Sequence[cv2.typing.MatLike], cv2.typing.MatLike]]:
        img = self.apply_filters(img=img)
        img = self.resize_and_crop(img=img)

        contours = self.find_contours(img=img)
        img = self.draw_contours(img=img, contours=contours)

        return (img, contours)
    
    def follow_contours(self,
                        contours: tuple[typing.Sequence[cv2.typing.MatLike], cv2.typing.MatLike],
                        offset: tuple[int, int] = (0, 0)):
        
        for contour in contours:
            shapes = contour.reshape(-1, 2).tolist() # type: ignore
            
            if len(shapes) <= 8:
                continue

            for shape in shapes:
                self.mouse.drag(initial=(int(offset[0] + int(shape[0])), 
                                        int(offset[1] + int(shape[1]))))
            self.mouse.release_drag()

    def follow_scanline(self, 
                        img: cv2.typing.MatLike, 
                        offset: tuple[int, int] = (0, 0)):
        height = img.shape[0]
        width = img.shape[1]

        for y in range(height):
            for x in range(width):
                pixel = img[y, x]
                if pixel == 255:
                    self.mouse.move_to(coord=(offset[0] + x, offset[1] + y))
                    self.mouse.click()

    def select_color(self, 
                     color: tuple[int, int, int]) -> None:
        self.mouse.move_to((self.canvas.color_picker_x, self.canvas.color_picker_y))
        self.mouse.click()
        time.sleep(0.01)
        self.mouse.move_to((560, 1186))
        self.mouse.click()
        pyautogui.write(str(color[0]))
        time.sleep(0.01)
        self.mouse.move_to((650, 1186))
        self.mouse.click()
        pyautogui.write(str(color[1]))
        time.sleep(0.01)
        self.mouse.move_to((720, 1186))
        self.mouse.click()
        pyautogui.write(str(color[2]))

        
    def draw(self, 
             image_path: str,
             mode: Mode = Mode.SCANLINE) -> None:
        image = self.load_image(path=image_path)
        image, contours = self.prepare_image(img=image)

        self.show(image)

        time.sleep(3)

        if mode.value == 'contour':
            self.follow_contours(contours=contours,
                                 offset=self.mouse.position())
        else:
            self.follow_scanline(img=image,
                                 offset=self.mouse.position())
                                 


if __name__ == '__main__':
    albedo = Albedo(canvas=Canvas(width=1303,
                                  height=729,
                                  color_picker_x=619,
                                  color_picker_y=918))
    
    albedo.draw(image_path="./images/howdy.jpg",
                mode=Mode.CONTOUR)