YOLACT pytorch模型转tensorflow savedModel格式

YOLACT： pytorch模型 -> onnx -> tensorflow savedModel模型。

pytorch源码需要修改部分代码才能转onnx，其中decode和nms需要用tf自己在实现。

下面代码是decode和nms以及转savedModel格式：

import time
import cv2
import tensorflow as tf
import numpy as npMEANS = np.array([103.94, 116.78, 123.68])[None, :, None, None]
STD = np.array([57.38, 57.12, 58.40])[None, :, None, None]def crop(pred, boxes):pred_shape = tf.shape(pred)w = tf.cast(tf.range(pred_shape[1]), tf.float32)h = tf.expand_dims(tf.cast(tf.range(pred_shape[2]), tf.float32), axis=-1)rows = tf.broadcast_to(w, pred_shape) / tf.cast(pred_shape[1], tf.float32)cols = tf.broadcast_to(h, pred_shape) / tf.cast(pred_shape[2], tf.float32)ymin = tf.broadcast_to(tf.reshape(boxes[:, 0], [-1, 1, 1]), pred_shape)xmin = tf.broadcast_to(tf.reshape(boxes[:, 1], [-1, 1, 1]), pred_shape)ymax = tf.broadcast_to(tf.reshape(boxes[:, 2], [-1, 1, 1]), pred_shape)xmax = tf.broadcast_to(tf.reshape(boxes[:, 3], [-1, 1, 1]), pred_shape)mask_left = (rows >= xmin)mask_right = (rows <= xmax)mask_bottom = (cols >= ymin)mask_top = (cols <= ymax)crop_mask = tf.math.logical_and(tf.math.logical_and(mask_left, mask_right),tf.math.logical_and(mask_bottom, mask_top))crop_mask = tf.cast(crop_mask, tf.float32)return pred * crop_mask# conf_preds [1, 2+1, 19248]  mask_data [1, 19248, 32]  decoded_boxes [19248, 4]  proto_data [138, 138, 32]
def detect(batch_idx, conf_preds, mask_data, decoded_boxes, proto_data, conf_thresh=0.15, nms_thresh=0.5, top_k=100):cur_scores = conf_preds[batch_idx, 1:, :]conf_scores = tf.math.reduce_max(cur_scores, axis=0)conf_score_id = tf.argmax(cur_scores, axis=0)keep = tf.squeeze(tf.where(conf_scores > conf_thresh))if tf.size(keep) == 0:return Nonescores = tf.gather(conf_scores, keep)boxes = tf.gather(decoded_boxes, keep)  # 获取符合conf阈值的bboxmasks = tf.gather(mask_data[batch_idx], keep)  # 获取符合阈值的mask  coefficientclasses = tf.gather(conf_score_id, keep)selected_indices = tf.image.non_max_suppression(boxes, scores, top_k, nms_thresh)boxes = tf.gather(boxes, selected_indices)scores = tf.gather(scores, selected_indices)masks = tf.gather(masks, selected_indices)classes = tf.gather(classes, selected_indices)masks = tf.linalg.matmul(proto_data, masks, transpose_a=False, transpose_b=True)masks = tf.nn.sigmoid(masks)masks = tf.transpose(masks, perm=(2, 0, 1))masks = crop(masks, boxes)masks = tf.image.resize(tf.expand_dims(masks, axis=-1), [550, 550], method="bilinear")masks = tf.cast(masks + 0.5, tf.int32)# masks = tf.squeeze(tf.cast(masks, tf.float32))return boxes, masks, scores, classesdef decode(loc, priors):variances = [0.1, 0.2]cxy = priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:]  # 按照prior进行偏移， 获得中心点坐标wh = priors[:, 2:] * tf.exp(loc[:, 2:] * variances[1])  # 获得why1 = cxy[:, 1] - wh[:, 1] / 2x1 = cxy[:, 0] - wh[:, 0] / 2y2 = cxy[:, 1] + wh[:, 1] / 2x2 = cxy[:, 0] + wh[:, 0] / 2boxes = tf.stack((y1, x1, y2, x2), 1)return boxesdef create():with tf.compat.v1.Session() as sess:output_graph_def = tf.compat.v1.GraphDef()with open("./output.pb", "rb") as f:output_graph_def.ParseFromString(f.read())_ = tf.import_graph_def(output_graph_def, name="")input = sess.graph.get_tensor_by_name("input.1:0")proto_data = sess.graph.get_tensor_by_name("1090:0")priors = sess.graph.get_tensor_by_name("1289:0")boxes_ = sess.graph.get_tensor_by_name("1286:0")conf_preds = sess.graph.get_tensor_by_name("1290:0")mask_data = sess.graph.get_tensor_by_name("1288:0")decoded_boxes = decode(boxes_[0], priors)conf_preds = tf.transpose(conf_preds, [0, 2, 1])boxes, masks, scores, classes = detect(0, conf_preds, mask_data, decoded_boxes, proto_data[0])tf.compat.v1.saved_model.simple_save(sess, "../output/savedmodel/", inputs={"input": input},outputs={"output0": boxes, "output1": masks, "output2": scores,"output3": classes})print("create savedmodel files success!")def run_savedmode():with tf.compat.v1.Session() as sess:meta_graph_def = tf.compat.v1.saved_model.loader.load(sess,[tf.compat.v1.saved_model.tag_constants.SERVING],"../output/savedmodel")signature = meta_graph_def.signature_def# get tensor namein_tensor_name = signature['serving_default'].inputs['input'].nameboxes = signature['serving_default'].outputs['output0'].namemasks = signature['serving_default'].outputs['output1'].namescores = signature['serving_default'].outputs['output2'].nameclasses = signature['serving_default'].outputs['output3'].nameinput_ = sess.graph.get_tensor_by_name(in_tensor_name)boxes = sess.graph.get_tensor_by_name(boxes)masks = sess.graph.get_tensor_by_name(masks)scores = sess.graph.get_tensor_by_name(scores)classes = sess.graph.get_tensor_by_name(classes)img_roi = cv2.imread("../output/samples/bus3.jpg")img_roi = cv2.resize(img_roi, (550, 550))img = img_roi.astype(np.float32)img = np.transpose(img, [2, 0, 1])img = np.expand_dims(img, 0)img = (img - MEANS) / STDimg = img[:, (2, 1, 0), :, :]for _ in range(2):start_ = time.time()b, m, s, c = sess.run([boxes, masks, scores, classes], feed_dict={input_: img})print("run time:", time.time() - start_)for i in range(m.shape[0]):cv2.imshow(f"mask{i}", m[i].astype(np.uint8) * 255)for one in b:x1 = int(one[1] * 550)y1 = int(one[0] * 550)x2 = int(one[3] * 550)y2 = int(one[2] * 550)img_roi = cv2.rectangle(img_roi, (x1, y1), (x2, y2), [255, 0, 0], 2)cv2.imshow("box", img_roi)cv2.waitKey(0)# 1. onnx format to tf format
# onnx-tf convert -i ./YOLACT.onnx -o ./output.pb
# 2. create tf model code to savedmodel format
# create()
# 3. test savedmodel format model
run_savedmode()

YOLACT pytorch模型转tensorflow savedModel格式相关推荐

记将MobileNet的pytorch模型转换成TorchScript格式出现的问题
这里写自定义目录标题使用torch.jit.script()将mobilenetv3的pytorch模型转换成TorchScript时出现 RuntimeError: Tried to access ...
(五)将YOLOv5 PyTorch模型权重转换为TensorFlow Lite格式
目录介绍建议使用Ultralytics进行重量转换的方法在Google Colab上测试TFLite权重 TFLite解释器,可在边缘设备上实现良好性能在本地测试YOLOv5模型权重下一步 ...
Netron展示pytorch模型结构
第一步:需要将pytorch模型转化为.ONNX格式,转化程序如下: model_path = './best_model/best_model.pth' Model = YL_Effientnat( ...
模型转换：pytorch模型转onnx, onnx转tensorflow, tensorflow转tflite
文章目录软件版本: pytorch模型转onnx onnx模型转tensorflow tensorflow模型转tflite 软件版本: tensorflow 2.3.1 pytorch 1.6.0 ...
TensorFlow与PyTorch模型部署性能比较
TensorFlow与PyTorch模型部署性能比较前言 2022了,选 PyTorch 还是 TensorFlow?之前有一种说法:TensorFlow 适合业界,PyTorch 适合学界.这种说 ...
keras模型转TensorFlow模型-tensorrt UFF格式
最近在学习tensorrt,需要将keras训练好保存的.hdf5格式模型转为tensorflow的.pb模型,然后转为tensorrt支持的uff格式. 做个记录. 代码如下: 转为tensorfl ...
深度学习之格式转换笔记(三)：keras(.hdf5)模型转TensorFlow(.pb) 转TensorRT(.uff)格式
环境: tensorflow1.15,cuda10.0,cudnn7.6.4 将keras训练好保存的.hdf5格式模型转为tensorflow的.pb模型,然后转为tensorrt支持的uff格式. ...
pytorch .pth模型转tensorflow .pb模型
训练好的pytorch模型如何转化为tensorflow的pb模型? 本人初步使用的是onnx框架: pytorch ---> onnx ----> tensorflow 使用onnx转 ...
如何使用TensorRT对训练好的PyTorch模型进行加速？
点击上方"3D视觉工坊",选择"星标" 干货第一时间送达作者丨伯恩legacy@知乎来源丨https://zhuanlan.zhihu.com/p/8831 ...
PyTorch模型部署：pth转onnx跨框架部署详解+代码
文章目录引言基础概念 onnx:跨框架的模型表达标准 onnxruntime:部署模型的推理引擎示例代码 0)安装onnx和onnxruntime 1)pytorch模型转onnx模型 2)on ...

YOLACT pytorch模型转tensorflow savedModel格式

YOLACT pytorch模型转tensorflow savedModel格式相关推荐

最新文章

热门文章