94 lines
3.2 KiB
Python
94 lines
3.2 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
lmy 2024-12-10 for esanet infer sunrgbd dataset
|
|
cd 1_script
|
|
run: python .\1_save.py --dataset sunrgbd
|
|
"""
|
|
|
|
import sys
|
|
sys.path.append(R"../0_esanet")
|
|
import argparse
|
|
from glob import glob
|
|
import os
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from src.args import ArgumentParserRGBDSegmentation
|
|
from src.build_model import build_model
|
|
from src.prepare_data import prepare_data
|
|
|
|
|
|
def _load_img(fp):
|
|
img = cv2.imread(fp, cv2.IMREAD_UNCHANGED)
|
|
if img.ndim == 3:
|
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
return img
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# arguments
|
|
parser = ArgumentParserRGBDSegmentation(
|
|
description='Efficient RGBD Indoor Sematic Segmentation (Inference)',
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
|
parser.set_common_args()
|
|
parser.add_argument('--ckpt_path', type=str,default='../weights/sunrgbd/r34_NBt1D.pth',
|
|
required=False,
|
|
help='Path to the checkpoint of the trained model.')
|
|
parser.add_argument('--depth_scale', type=float,
|
|
default=1.0,
|
|
help='Additional depth scaling factor to apply.')
|
|
args = parser.parse_args()
|
|
|
|
# dataset
|
|
args.pretrained_on_imagenet = False # we are loading other weights anyway
|
|
dataset, preprocessor = prepare_data(args, with_input_orig=True)
|
|
n_classes = dataset.n_classes_without_void
|
|
|
|
# model and checkpoint loading
|
|
model, device = build_model(args, n_classes=n_classes)
|
|
checkpoint = torch.load(args.ckpt_path,
|
|
map_location=lambda storage, loc: storage)
|
|
model.load_state_dict(checkpoint['state_dict'])
|
|
print('Loaded checkpoint from {}'.format(args.ckpt_path))
|
|
model.eval()
|
|
model.to(device)
|
|
|
|
# load sample
|
|
img_rgb = _load_img('../2_compile/qtset/sample/sample_rgb.png')
|
|
img_depth = _load_img('../2_compile/qtset/sample/sample_depth.png').astype('float32') * args.depth_scale
|
|
h, w, _ = img_rgb.shape
|
|
|
|
# preprocess sample
|
|
sample = preprocessor({'image': img_rgb, 'depth': img_depth})
|
|
|
|
# add batch axis and copy to device
|
|
image = sample['image'][None].to(device)
|
|
depth = sample['depth'][None].to(device)
|
|
|
|
# apply network
|
|
pred = model(image, depth)
|
|
print('-----trace model------')
|
|
torch.onnx.export(model,(image, depth),'../2_compile/fmodel/esanet_480x640_sunrgbd.onnx',opset_version=17)
|
|
trace_model = torch.jit.trace(model,(image,depth))
|
|
torch.jit.save(trace_model,'../2_compile/fmodel/esanet_480x640_sunrgbd.pt')
|
|
|
|
|
|
pred = F.interpolate(pred, (h, w),
|
|
mode='bilinear', align_corners=False)
|
|
pred = torch.argmax(pred, dim=1)
|
|
pred = pred.cpu().numpy().squeeze().astype(np.uint8)
|
|
|
|
# show result
|
|
pred_colored = dataset.color_label(pred, with_void=False)
|
|
fig, axs = plt.subplots(1, 3, figsize=(16, 3))
|
|
[ax.set_axis_off() for ax in axs.ravel()]
|
|
axs[0].imshow(img_rgb)
|
|
axs[1].imshow(img_depth, cmap='gray')
|
|
axs[2].imshow(pred_colored)
|
|
plt.suptitle(f"Image depth" f"Model: {args.ckpt_path}")
|
|
plt.show()
|