import gradio as gr
import cv2
import os

ed     = cv2.ximgproc.createEdgeDrawing()
params = cv2.ximgproc.EdgeDrawing.Params()
params.PFmode = True
ed.setParams(params)

def center_crop_mod64(img):
    h, w = img.shape[:2]
    shortside, longside = min(w, h), max(w, h)
    longside_crop = (longside // 64) * 64
    left   = (w - longside_crop) // 2 if w > h else 0
    top    = (h - longside_crop) // 2 if h > w else 0
    right  = left + longside_crop if w > h else shortside
    bottom = top  + longside_crop if h > w else shortside
    return img[top:bottom, left:right]
    
def edpf(image_rgb):
    img_gray = cv2.cvtColor(image_rgb, cv2.COLOR_BGR2GRAY)
    img_crop = center_crop_mod64(img_gray)
    edges    = ed.detectEdges(img_crop)
    edge_map = ed.getEdgeImage(edges)
    return edge_map

examples=[
        os.path.join(os.path.dirname(__file__), "images/bag.png"),
        os.path.join(os.path.dirname(__file__), "images/beard.png"),
        os.path.join(os.path.dirname(__file__), "images/bird.png"),
        os.path.join(os.path.dirname(__file__), "images/cat.png"),
        os.path.join(os.path.dirname(__file__), "images/dog2.png"),
        os.path.join(os.path.dirname(__file__), "images/house.png"),
        os.path.join(os.path.dirname(__file__), "images/house2.png"),
        os.path.join(os.path.dirname(__file__), "images/human.png"),
        os.path.join(os.path.dirname(__file__), "images/kitten.png"),
        os.path.join(os.path.dirname(__file__), "images/lion.png"),
        os.path.join(os.path.dirname(__file__), "images/man.png"),
        os.path.join(os.path.dirname(__file__), "images/robot.png"),
        os.path.join(os.path.dirname(__file__), "images/robotics.png"),
        os.path.join(os.path.dirname(__file__), "images/room.png"),
        os.path.join(os.path.dirname(__file__), "images/room2.png"),
        os.path.join(os.path.dirname(__file__), "images/suit.png"),
        os.path.join(os.path.dirname(__file__), "images/tree.png"),
        os.path.join(os.path.dirname(__file__), "images/vermeer.png"),
    ]
    
app = gr.Interface(fn=edpf, inputs=gr.Image(value="images/dog2.png"), outputs="image", title="Edge Drawing Parameter Free", description="A modern edge detection algorithm which requires no parameter tuning. Generate edge maps for the edpf ControlNet model at https://huggingface.co/GeroldMeisinger/control-edgedrawing", examples=examples)

if __name__ == "__main__":
    app.launch()