class ModelLoader:
def __init__(self, path: str):
self.path = path
self.net = Net()
def load_model(self):
self.net.load_state_dict(
torch.load(self.path, map_location=torch.device('cpu')))
self.net.eval()
return self.net
from flask import Flask, jsonify, request
from flask_restful import Api, Resource
import numpy as np
import traceback
import json
from model.model import Net
import torch
import cv2
app = Flask(__name__)
api = Api(app)
model = Net()
model.load_state_dict(torch.load('mnist_cnn.pt'))
model.eval()
# convert request_input dict to input accepted by model.
def parse_input(request_input):
"""parse input to make it ready for model input.
Arguments:
request_input::file- input received from API call.
Returns:
model_ready_input::torch.Tensor- data ready to be fed into model.
"""
img = request_input.read()
img = np.frombuffer(img, np.uint8)
img = cv2.imdecode(img, cv2.IMREAD_COLOR)
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
default="",
help="path to load model checkpoint")
parser.add_argument("--test",
type=str,
default="../ITS_data/test/data",
help="path to load test images")
opt = parser.parse_args()
# for i in range(669,670):
net = Net()
print(opt)
print(opt.checkpoint)
net.load_state_dict(torch.load(opt.checkpoint)['state_dict'])
net.eval()
net = nn.DataParallel(net, device_ids=[0]).cuda()
images = utils.load_all_image(opt.test)
ssims = []
psnrs = []
psnrs2 = []
str_label = "../ITS_data/test/label/"
def compute_psnr(ref_im, res_im):
ref_img = scipy.misc.fromimage(ref_im).astype(float) / 255.0
res_img = scipy.misc.fromimage(res_im).astype(float) / 255.0
squared_error = np.square(ref_img - res_img)