File size: 1,519 Bytes
c10aebf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
import cv2
import numpy as np
from PIL import Image
from PIL import ImageOps
from sd_bmab.base.context import Context
from sd_bmab.base.processorbase import ProcessorBase
class EdgeEnhancement(ProcessorBase):
def __init__(self) -> None:
super().__init__()
self.edge_low_threadhold = 50
self.edge_high_threadhold = 200
self.edge_strength = 0.5
def preprocess(self, context: Context, image: Image):
if context.args['edge_flavor_enabled']:
self.edge_low_threadhold = context.args['edge_low_threadhold']
self.edge_high_threadhold = context.args['edge_high_threadhold']
self.edge_strength = context.args['edge_strength']
return context.args['edge_flavor_enabled']
def process(self, context: Context, image: Image):
context.add_generation_param('BMAB edge flavor low threadhold', self.edge_low_threadhold)
context.add_generation_param('BMAB edge flavor high threadhold', self.edge_high_threadhold)
context.add_generation_param('BMAB edge flavor strength', self.edge_strength)
numpy_image = np.array(image)
base = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
arcanny = cv2.Canny(base, self.edge_low_threadhold, self.edge_high_threadhold)
canny = Image.fromarray(arcanny)
canny = ImageOps.invert(canny)
img = image.convert('RGB')
newdata = [(0, 0, 0) if mdata == 0 else ndata for mdata, ndata in zip(canny.getdata(), img.getdata())]
newbase = Image.new('RGB', img.size)
newbase.putdata(newdata)
return Image.blend(img, newbase, alpha=self.edge_strength).convert("RGB")
|