main.py 文件源码

def test(m):
    class DrawingState:
        def __init__(self):
            self.x_prev = 0
            self.y_prev = 0
            self.drawing = False
            self.update = True

    def interactive_drawing(event, x, y, flags, param):
        image = param[0]
        state = param[1]
        if event == cv2.EVENT_LBUTTONDOWN:
            state.drawing = True
            state.x_prev, state.y_prev = x, y
        elif event == cv2.EVENT_MOUSEMOVE:
            if state.drawing:
                cv2.line(image, (state.x_prev, state.y_prev), (x, y), (1, 1, 1), 1)
                state.x_prev = x
                state.y_prev = y
                state.update = True
        elif event == cv2.EVENT_LBUTTONUP:
            state.drawing = False
        elif event == cv2.EVENT_RBUTTONDOWN:
            image.fill(0)
            state.update = True

    cv2.namedWindow('Canvas')
    image_input = np.zeros((FLAGS.input_height, FLAGS.input_width, 3), np.float32)
    state = DrawingState()
    cv2.setMouseCallback('Canvas', interactive_drawing, [image_input, state])
    while cv2.getWindowProperty('Canvas', 0) >= 0:
        if state.update:
            reshaped_image_input = np.array([image_input])
            image_output = m.test(reshaped_image_input)
            concatenated = np.concatenate((image_input, image_output[0]), axis=1)
            color_converted = cv2.cvtColor(concatenated, cv2.COLOR_RGB2BGR)
            cv2.imshow('Canvas', color_converted)
            state.update = False

        k = cv2.waitKey(1) & 0xFF
        if k == 27:  # esc
            break
    cv2.destroyAllWindows()