STDC-Seg
STDC-Seg copied to clipboard
visulization is not so good
Thanks for your great work! I am using your provided model to evaluate single image and find that the visualization is not so good? How could I get better results? Images are shown here: Codes are shown below:
code
import os
import cv2
import numpy as np
import torch
import time
from tqdm import tqdm
from models.model_stages_trt import BiSeNet
import torchvision.transforms as transforms
from PIL import Image
import torch.nn.functional as F
def render_semantic(label, colors):
label_rgb = cv2.cvtColor(label.astype("uint8"), cv2.COLOR_GRAY2BGR)
rendered_label = np.array(cv2.LUT(label_rgb, colors))
rendered_label = cv2.cvtColor(rendered_label, cv2.COLOR_RGB2BGR)
return rendered_label
def cityscape_color_maps():
colors = np.zeros((256, 1, 3), dtype='uint8')
colors[0, :, :] = [128, 64, 128] # road
colors[1, :, :] = [244, 35, 232] # sidewalk
colors[2, :, :] = [70, 70, 70] # building
colors[3, :, :] = [102, 102, 156] # wall
colors[4, :, :] = [190, 153, 153] # fence
colors[5, :, :] = [153, 153, 153] # pole
colors[6, :, :] = [250, 170, 30] # traffic light
colors[7, :, :] = [220, 220, 0] # traffic sign
colors[8, :, :] = [107, 142, 35] # vegetation
colors[9, :, :] = [152, 251, 152] # terrain
colors[10, :, :] = [70, 130, 180] # sky
colors[11, :, :] = [220, 20, 60] # person
colors[12, :, :] = [255, 0, 0] # rider
colors[13, :, :] = [0, 0, 142] # car
colors[14, :, :] = [0, 0, 70] # truck
colors[15, :, :] = [0, 60, 100] # bus
colors[16, :, :] = [0, 80, 100] # train
colors[17, :, :] = [0, 0, 230] # motocycle
colors[18, :, :] = [119, 11, 32] # bicycle
return colors
def compute_latency_ms_pytorch(input_file_list=None, color_maps=None, device=None):
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
# Configuration ##############
use_boundary_2 = False
use_boundary_4 = False
use_boundary_8 = True
use_boundary_16 = False
use_conv_last = False
n_classes = 19
# STDC2Seg-75 97.0FPS on NVIDIA GTX 1080Ti
backbone = 'STDCNet1446'
methodName = 'STDC2-Seg'
inputSize = 768
inputScale = 75
inputDimension = (1, 3, 768, 1536)
model = BiSeNet(backbone=backbone, n_classes=n_classes,
use_boundary_2=use_boundary_2, use_boundary_4=use_boundary_4,
use_boundary_8=use_boundary_8, use_boundary_16=use_boundary_16,
input_size=inputSize, use_conv_last=use_conv_last)
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
print('loading parameters...')
respth = './checkpoints/{}/'.format(methodName)
save_pth = os.path.join(respth, 'model_maxmIOU{}.pth'.format(inputScale))
model.load_state_dict(torch.load(save_pth))
model = model.to(device)
model.eval()
# input = torch.randn(*inputDimension).to(device)
input = Image.open(input_file_list[0]).convert('RGB')
input = to_tensor(input).to(device)[None]
N, C, H, W = input.shape
input = F.interpolate(input, (768, 1536), mode='bilinear', align_corners=True)
result = model(input)
# result = F.interpolate(result, [H, W], mode='bilinear', align_corners=True)
result = torch.softmax(result, dim=1)
result = torch.argmax(result, dim=1).cpu().numpy()
result = result.transpose(1, 2, 0)
result = render_semantic(result, color_maps)
cv2.imwrite("./test.png", result)
if __name__ == "__main__":
device = torch.device('cuda:0')
file_list = ["./demo.png"]
color_maps = cityscape_color_maps()
compute_latency_ms_pytorch(input_file_list=file_list, color_maps=color_maps, device=device)
@DRosemei I have same issue during inference. Did you solve the problem?
@DRosemei I have same issue during inference. Did you solve the problem?
Not yet
@DRosemei Any progress?