tensorflow-自定義op

在處理數據時需要自定義處理方法，這時就可以使用tf.py_func(func, inp, Tout, stateful=True, name=None)來進行處理

具體例子

import tensorflow as tf

import cv2
import numpy as np


def forward_convert(coordinate, with_label=True):
    """
    :param coordinate: format [y_c, x_c, h, w, theta]
    :return: format [y1, x1, y2, x2, y3, x3, y4, x4]
    """
    boxes = []
    if with_label:
        for rect in coordinate:
            box = cv2.boxPoints(((rect[1], rect[0]), (rect[3], rect[2]), rect[4]))
            box = np.reshape(box, [-1, ])
            boxes.append([box[1], box[0], box[3], box[2], box[5], box[4], box[7], box[6], rect[5]])
    else:
        for rect in coordinate:
            box = cv2.boxPoints(((rect[1], rect[0]), (rect[3], rect[2]), rect[4]))
            box = np.reshape(box, [-1, ])
            boxes.append([box[1], box[0], box[3], box[2], box[5], box[4], box[7], box[6]])

    return np.array(boxes, dtype=np.float32)


coord = np.array([[150, 150, 50, 100, -90, 1],
                  [150, 150, 100, 50, -90, 1],
                  [150, 150, 50, 100, -45, 1],
                  [150, 150, 100, 50, -45, 1]])
coord = tf.constant(coord)

convert = tf.py_func(forward_convert, inp=[coord], Tout=tf.float32)

with tf.Session() as sess:
    print(sess.run(convert))