利用flask建立pytorch 服务

客官°小女子只卖身不卖艺 2023-06-26 08:11 20阅读 0赞

pytorch服务不像tf那样可以进行直接部署，需要自己搭建服务才行

我这个代码是拿自己之前手动部署tf基础上改的，所以上面有tf的痕迹

服务端：

from flask import Flask, render_template, request, redirect, url_for, flash, jsonify
    #from tensorflow.keras.models import load_model
    import cv2
    import numpy as np
    #import tensorflow as tf
    import uuid
    import numpy as np
    # load model
    def get_model():
        global graph
        graph = tf.compat.v1.get_default_graph()
        with graph.as_default():
            global model
            model = load_model('model.h5')
    
    app = Flask(__name__)
    #get_model()
    """
    @app.route('/predict', methods=['POST'])
    def upload_image():
        with graph.as_default():
            if request.method == 'POST':
                #print(request.get_json())
                response = {'success': True}
                print(response)
                # Prepare image for model
                get_image=request.get_json()['input_image']
                #print(get_image)
                #get_image_tensor=tf.convert_to_tensor(get_image)
                #resized = tf.image.resize(get_image_tensor, [256, 256])
                #print(get_image_tensor)
                #normalized = (get_image_tensor / 127.5) - 1
                #print(normalized)
                #batched = tf.expand_dims(get_image_tensor, 0)
                #print(batched)
                # Apply model!
                model.summary()
                predicted = model.predict(get_image,steps=1)
                print(predicted)
                predicted = predicted[0,:,:,:]
                print(predicted)
                unnormalized = (predicted + 1) / 2
    
                # Save result on disk to eventually show to user
                result_name = 'static/tmp/{}.jpg'.format(uuid.uuid4())
                cv2.imwrite(result_name, unnormalized * 255)
    
                response['processed_image'] = result_name
                return jsonify(response)
    """
    
    import torch
    
    new_model = torch.load('./checkpoints/coconew/latest_net_G.pth')
    new_model.eval()
    print(new_model)
    import numpy
    @app.route('/predict', methods=['POST'])
    def upload_image():
        if request.method == 'POST':
            #print(request.get_json()['input_image'])
            response = {'success': True}
            print(response)
            # Prepare image for model
            get_image=request.get_json()['input_image']
            #print(get_image)
            #get_image_tensor=tf.convert_to_tensor(get_image)
            #resized = tf.image.resize(get_image_tensor, [256, 256])
            #print(get_image_tensor)
            #normalized = (get_image_tensor / 127.5) - 1
            #print(normalized)
            #batched = tf.expand_dims(get_image_tensor, 0)
            #print(batched)
            # Apply model!
            #new_model.summary()
            get_image=torch.from_numpy(numpy.array(get_image))
            predicted =  new_model(get_image)
            print(predicted)
            predicted = predicted[0,:,:,:]
            unnormalized = (predicted + 1) / 2
            # Save result on disk to eventually show to user
            result_name = 'static/tmp/{}.jpg'.format(uuid.uuid4())
            cv2.imwrite(result_name, unnormalized * 255)
            response['processed_image'] = result_name
            return jsonify(response)
    
    
    
    # flask will check if raised exception is of type 'SomeException' (or lower)
    # if so, will just execute this method
    @app.errorhandler(ValueError)
    def handle_error(error):
        print("完蛋")
        response = jsonify({'message': str(error)})
        return response
    
    if __name__ == "__main__":
        app.run(debug=True)

客户端

from __future__ import print_function
    #import tensorflow as tf
    import base64
    import requests
    import cv2
    # model with the name "pix2pix" and using the predict interface.
    SERVER_URL = 'http://localhost:5000/predict'
    import numpy as np
    
    #from __future__ import absolute_import, division, print_function, unicode_literals
    import json
    import tensorflow as tf
    import  numpy as np
    import os
    import time
    
    from matplotlib import pyplot as plt
    from IPython import display
    #import tensorflow as tf
    #import tensorflow.contrib.eager as tfe
    #tf.enable_eager_execution()
    
    #import tensorflow.contrib.eager as tfe
    #tfe.enable_eager_execution()
    
    #import tensorflow as tf
    #tf.enable_eager_execution()
    
    BUFFER_SIZE = 400
    BATCH_SIZE = 1
    IMG_WIDTH = 256
    IMG_HEIGHT = 256
    
    ##定义加载图片的函数
    def load(image_file):
        image = tf.io.read_file(image_file)
        image = tf.image.decode_jpeg(image)
        w = tf.shape(image)[1]
        w = w // 2
        real_image = image[:,:w, :]
        input_image = image[:, w:, :]
    
        input_image = tf.cast(input_image, tf.float32)
        real_image = tf.cast(real_image, tf.float32)
        return input_image, real_image
    ##定义resize函数
    def resize(input_image, real_image, height, width):
        input_image = tf.image.resize(input_image, [height, width],
                                    method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
        real_image = tf.image.resize(real_image, [height, width],
                                   method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
    
        return input_image, real_image
    ##定义随机剪裁函数
    def random_crop(input_image, real_image):
        stacked_image = tf.stack([input_image, real_image], axis=0)
        cropped_image = tf.image.random_crop(
          stacked_image, size=[2, IMG_HEIGHT, IMG_WIDTH, 3])
    
        return cropped_image[0], cropped_image[1]
    
    # normalizing the images to [-1, 1]
    
    def normalize(input_image, real_image):
        input_image = (input_image / 127.5) - 1
        real_image = (real_image / 127.5) - 1
    
        return input_image, real_image
    
    
    @tf.function()
    def random_jitter(input_image, real_image):
        """先把size设置为286*286，然后进行随机剪裁"""
      # resizing to 286 x 286 x 3
        input_image, real_image = resize(input_image, real_image, 286, 286)
    
      # randomly cropping to 256 x 256 x 3
        input_image, real_image = random_crop(input_image, real_image)
    
        if tf.random.uniform(()) > 0.5:
        # random mirroring
            input_image = tf.image.flip_left_right(input_image)
            real_image = tf.image.flip_left_right(real_image)
    
        return input_image, real_image
    
    ##加载训练数据
    def load_image_train(image_file):
        input_image, real_image = load(image_file)
        input_image, real_image = random_jitter(input_image, real_image)
        input_image, real_image = normalize(input_image, real_image)
    
        return input_image, real_image
    ##加载测试数据
    def load_image_test(image_file):
        input_image, real_image = load(image_file)
        input_image, real_image = resize(input_image, real_image,
                                       IMG_HEIGHT, IMG_WIDTH)
        input_image, real_image = normalize(input_image, real_image)
    
        return input_image, real_image
    
    
    
    
    PATH = os.path.join(os.path.dirname('/opt/AI/facades/'), '/opt/AI/facades/')
    train_dataset = tf.data.Dataset.list_files(PATH+'train/*.jpg')
    train_dataset = train_dataset.map(load_image_train,
                                      num_parallel_calls=tf.data.experimental.AUTOTUNE)
    train_dataset = train_dataset.shuffle(BUFFER_SIZE)
    train_dataset = train_dataset.batch(BATCH_SIZE)
    
    
    test_dataset = tf.data.Dataset.list_files(PATH+'test/*.jpg')
    test_dataset = test_dataset.map(load_image_test)
    test_dataset = test_dataset.batch(BATCH_SIZE)
    import numpy
    a=numpy.ones([184,256,256])
    print(a)
    num=1
    i=0
    for inp, tar in test_dataset.take(num):
        payload = {
      "instances": [{'input_1': inp.numpy()[0].tolist()}]
        }
        #predict_request='{"input_image":%s}' %inp.numpy().tolist()
        predict_request={"input_image":inp.numpy().tolist()}
        #predict_request={"input_image":a.tolist()}
        print(inp.numpy()[0])
        #print(predict_request)
        #json_request = '{
       { "instances" : {} }}'.format(np.array2string(inp.numpy(), separator=',')
        #print(predict_request)
        response=requests.post(SERVER_URL, json=predict_request)
        response.raise_for_status()
        print(response)
        prediction = response.json()['predictions']
        print(prediction)
        tmp=np.ones((256, 256,3))
        prediction=127.5*(prediction+tmp)
        print(prediction)
       # prediction=prediction.astype(int)
        #print(prediction)
        cv2.imwrite('./md'+str(i)+'.jpg',prediction[0])
        i=i+1

这里可以自己写客户端，这个里面是直接使用facades里面的数据集，就没改了