问OpenCV dnn模块生成与原始torch模型不同的预测

EN

### upgrade opencv to last version ###
!pip install --upgrade opencv-python0

### clone git with U2Net ###
%cd /content
!git clone https://github.com/shreyas-bk/U-2-Net

### download weights for U2Net ###
!gdown --id 1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ -O /content/U-2-Net/u2net.pth

###
%cd /content/U-2-Net

### imports ###
from google.colab import files
from model import U2NET
import torch
import os
import numpy as np
from torchvision import transforms
import cv2 as cv
from skimage import io, transform
from PIL import Image

### instantiate U2Net ###
model_dir = '/content/U-2-Net/u2net.pth'
net = U2NET(3, 1)
net.load_state_dict(torch.load(model_dir, map_location='cpu'))
net.eval() 

### export torch model to onnx ###
img = torch.randn(1, 3, 320, 320, requires_grad=False)
img = img.to(torch.device('cpu'))
output_dir = os.path.join('/content/u2net.onnx')
torch.onnx.export(net, img, output_dir, opset_version=11, verbose=True)

### load model from OpenCV ###
cv_net = cv.dnn.readNetFromONNX('/content/u2net.onnx')

### your test image here ###
IMG_PATH = '/content/<IMAGE_NAME>.png'

### load image ###
image = Image.open(IMG_PATH)

### preprocessing ###

def preprocess_image(image, output_size=320, for_torch=False):
  ''''''
  # resize image
  img = transform.resize(image, (output_size, output_size), mode="constant")

  # mean subtraction and normalization
  tmp_img = np.zeros((img.shape[0], img.shape[1], 3))
  img = img / np.max(img)

  tmp_img[:, :, 0] = (img[:, :, 0] - 0.485) / 0.229
  tmp_img[:, :, 1] = (img[:, :, 1] - 0.456) / 0.224
  tmp_img[:, :, 2] = (img[:, :, 2] - 0.406) / 0.225
  tmp_img = tmp_img.transpose((2, 0, 1))
  
  if for_torch:
    return torch.from_numpy(tmp_img)

  return tmp_img

### predictions norm ###
def norm_pred(d):
  '''
  normalize predictions
  '''
  ma = d.max()
  mi = d.min()
  dn = (d - mi) / (ma - mi)
  return dn

### the magic ###
def remove_bg(image, processed, for_torch=False):
  pred = None
  if for_torch:
    with torch.no_grad():
      inputs_test = torch.FloatTensor(processed.unsqueeze(0).float())
      preds, _, _, _, _, _, _ = net(inputs_test)
  else:
    cv_net.setInput(np.expand_dims(processed, axis=0))
    preds = cv_net.forward()  
    
  pred = preds[:, 0, :, :]
  # normalization
  pred_normalized = norm_pred(pred.cpu().detach().numpy() if for_torch else pred)
  # squeeze
  predict = pred_normalized.squeeze()
  # to RGB
  img_out = Image.fromarray(predict * 255).convert("RGB")
  image = image.resize((img_out.size), resample=Image.BILINEAR)
  empty_img = Image.new("RGBA", (image.size), 0)
  img_out = Image.composite(image, empty_img, img_out.convert("L"))

  # draw 
  img_out = img_out.resize((image.size), resample=Image.BILINEAR)
  empty_img = Image.new("RGBA", (image.size), 0)
  img_out = Image.composite(image, empty_img, img_out)

  return img_out

### preprocess image ###
sample = preprocess_image(np.array(image), for_torch=True)

### torch results ###
remove_bg(image, sample, for_torch=True)

### opencv results ###
remove_bg(image, sample, for_torch=False)

/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py:780: UserWarning: Note that order of the arguments: ceil_mode and return_indices will changeto match the args list in nn.MaxPool2d in a future release.
  warnings.warn("Note that order of the arguments: ceil_mode and return_indices will change"
/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py:3704: UserWarning: nn.functional.upsample is deprecated. Use nn.functional.interpolate instead.
  warnings.warn("nn.functional.upsample is deprecated. Use nn.functional.interpolate instead.")
/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py:1944: UserWarning: nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.
  warnings.warn("nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.")

### install onnxruntime (version 1.11.1) ###
!pip install onnxruntime

### loading model with onnxruntime ###
import onnxruntime
ort_session = onnxruntime.InferenceSession("/content/u2net.onnx")

### updating remove_bg ###
def remove_bg(image, processed, backend='torch'):
  ''''''
  if backend not in ['torch', 'OpenCV', 'onnx']: raise AttributeError('Wrong backend.')

  pred = None
  if backend == 'torch':
    print('NO')
    with torch.no_grad():
      inputs_test = torch.FloatTensor(processed.unsqueeze(0).float())
      preds, _, _, _, _, _, _ = net(inputs_test)
  if backend == 'OpenCV':
    print('NO')
    cv_net.setInput(np.expand_dims(processed, axis=0))
    preds = cv_net.forward()  
  if backend== 'onnx':
    print('HI')
    ort_inputs = {ort_session.get_inputs()[0].name: np.expand_dims(processed.astype(np.float32), axis=0)}
    ort_outs = ort_session.run(None, ort_inputs)
    preds = ort_outs[0]
    
  pred = preds[:, 0, :, :]
  # normalization
  pred_normalized = norm_pred(pred.numpy() if backend=='torch' else pred)
  # squeeze
  predict = pred_normalized.squeeze()
  # to RGB
  img_out = Image.fromarray(predict * 255).convert("RGB")
  image = image.resize((img_out.size), resample=Image.BILINEAR)
  empty_img = Image.new("RGBA", (image.size), 0)
  img_out = Image.composite(image, empty_img, img_out.convert("L"))

  # draw 
  img_out = img_out.resize((image.size), resample=Image.BILINEAR)
  empty_img = Image.new("RGBA", (image.size), 0)
  img_out = Image.composite(image, empty_img, img_out)

  return img_out


### inference with onnxruntime ###
remove_bg(image, sample, backend='onnx')