第12章 PyTorch圖像分割代碼框架-3：推理與部署

# 導(dǎo)入相關(guān)庫import numpy as npimport torchfrom PIL import Image# 定義推理函數(shù)def inference(model, test_img):    img = Image.open(test_img)    img = val_transform(img)    img = img.unsqueeze(0).to('cuda')
    with torch.no_grad():        outputs = model(img)        preds = outputs.detach().max(dim=1)[1].cpu().numpy()        print(preds.shape)        pred = VOCSegmentation.decode_target(preds[0]).astype(np.uint8)        Image.fromarray(pred).save(os.path.join('s%_pred.png' % test_img.split('.')[0]))

python inference.py --data_root 2007_000676.jpg --model deeplabv3plus_resnet101

# 導(dǎo)入flask相關(guān)模塊from flask import Flask, jsonify# 創(chuàng)建應(yīng)用app = Flask(__name__)# 定義預(yù)測路由@app.route('/predict', methods=['POST'])def predict():    return jsonify({'class_id': 'IMAGE_NET_XXX', 'class_name': 'Cat'})

# 導(dǎo)入相關(guān)庫import torchfrom torchvision import transformsfrom PIL import Imageimport ioimport numpy as npfrom utils import ext_transforms_new as etfrom datasets import VOCSegmentationfrom flask import Flask, request, jsonifyimport models
# 創(chuàng)建應(yīng)用app = Flask(__name__)
# 模型字典model_map = {    'deeplabv3plus_resnet50': models.deeplabv3plus_resnet50,    'deeplabv3plus_resnet101': models.deeplabv3plus_resnet101,}
# 創(chuàng)建模型model = model_map['deeplabv3plus_resnet101'](num_classes=21,output_stride=16)  # 導(dǎo)入模型model.load_state_dict(torch.load('../checkpoints/deeplabv3plus_resnet101_voc.pth')['model_state'])model.to('cuda')model.eval()print('model loaded.')
# 定義數(shù)據(jù)轉(zhuǎn)換方法transform = et.Compose([    transforms.ToTensor(),    transforms.Normalize(mean=[0.485, 0.456, 0.406],                         std=[0.229, 0.224, 0.225]),])
# 定義模型預(yù)測路由@app.route('/predict', methods=['POST'])def predict():    if request.method == 'POST':        # 從請求中讀取輸入圖像        image = request.files['image'].read()        image = Image.open(io.BytesIO(image))
        # 圖像變換        input_tensor = transform(image).unsqueeze(0).to('cuda')                # 模型預(yù)測        with torch.no_grad():            output = model(input_tensor)            preds = output.detach().max(dim=1)[1].cpu().numpy()            print(preds.shape)            # 對輸出進(jìn)行解碼，轉(zhuǎn)換為mask            preds = VOCSegmentation.decode_target(preds[0]).astype(np.uint8)
        # 轉(zhuǎn)換成list        result = preds.tolist()        return jsonify(result)
if __name__ == '__main__':    app.run(debug=True)

python app.py

resp = requests.post("http://localhost:5000/predict", files={"image": open('./deployment/2007_000676.jpg', 'rb')})

# 創(chuàng)建應(yīng)用app = Flask(__name__)# 定義上傳和預(yù)測路由@app.route('/', methods=['GET', 'POST'])def upload_predict():  # POST請求后讀取圖像    if request.method == 'POST':        image_file = request.files['image']        if image_file:            image_location = os.path.join(                app.config['UPLOAD_FOLDER'],                image_file.filename            )            image_file.save(image_location)            # 圖像變換            image = Image.open(image_location).convert('RGB')            input_tensor = transform(image).unsqueeze(0).to('cuda')
            # 模型預(yù)測            with torch.no_grad():                output = model(input_tensor)                preds = output.detach().max(dim=1)[1].cpu().numpy()                print(preds.shape)                # mask解碼                preds = VOCSegmentation.decode_target(preds[0]).astype(np.uint8)            # 保存圖像到指定路徑            segmented_image = Image.fromarray(preds)            segmented_image_path = image_location.replace('.jpg', '_segmented.jpg')            segmented_image.save(segmented_image_path)            display_input_path = '../' + image_location            display_segmented_path = '../' + segmented_image_path            # 渲染結(jié)果到網(wǎng)頁            return render_template('index.html', input_image=display_input_path, segmented_image=display_segmented_path)            return render_template('index.html')