1122334455

1122334455

import torch, os, cv2
from model.model import parsingNet
from utils.common import merge_config
from utils.dist_utils import dist_print
import torch
import scipy, tqdm
import numpy as np
import torchvision.transforms as transforms
from data.dataset import LaneTestDataset

def softmax(x, axis=0):
    """Compute softmax values for each sets of scores in x."""
    return np.exp(x) / np.sum(np.exp(x), axis=axis)
if __name__ == "__main__":
    torch.backends.cudnn.benchmark = True
    
    cls_num_per_lane = 18
    

    net = parsingNet(pretrained = False, backbone='18',cls_dim = (200 + 1,cls_num_per_lane,4),
                    use_aux=False).cuda() # we dont need auxiliary segmentation in testing

    state_dict = torch.load('culane_18.pth', map_location = 'cpu')

    net.load_state_dict(state_dict['model'])
    net.eval()

    img_transforms = [
        transforms.ToPILImage(),
        transforms.Resize((288, 800)),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ]
    grid_num = 200
    cap = cv2.VideoCapture("E:/datasets/road_line_dataset_mine/videos/not_rural/road_line/Road trip soundtrip dashcam.mp4")
    while(True):
        ret, frame = cap.read()
        img = img_transforms[0](frame)
        img = img_transforms[1](img)
        
        img = img_transforms[2](img)        
        # frame = np.transpose(img.numpy(), (1,2,0))
        img = img_transforms[3](img)

        img.unsqueeze_(0)
        print(np.shape(frame))
        imgs = img.cuda()
        with torch.no_grad():
            out = net(imgs)

        col_sample = np.linspace(0, 800 - 1, grid_num)
        col_sample_w = col_sample[1] - col_sample[0]


        out_j = out[0].data.cpu().numpy()
        out_j = out_j[:, ::-1, :]
        prob = softmax(out_j[:-1, :, :], axis=0)
        idx = np.arange(grid_num) + 1
        idx = idx.reshape(-1, 1, 1)
        loc = np.sum(prob * idx, axis=0)
        out_j = np.argmax(out_j, axis=0)
        loc[out_j == grid_num] = 0
        out_j = loc
        # frame = np.ascontiguousarray(frame, dtype=np.float32)
        # import pdb; pdb.set_trace()
        # vis = cv2.imread(os.path.join(grid_num,names[0]))
        for i in range(out_j.shape[1]):
            if np.sum(out_j[:, i] != 0) > 2:
                for k in range(out_j.shape[0]):
                    if out_j[k, i] > 0:
                        # ppp = (int(out_j[k, i] * col_sample_w * 1640 / 800) - 1, int(590 - k * 20) - 1)
                        ppp = (int(out_j[k, i] * col_sample_w * frame.shape[1] / 800) - 1, int(frame.shape[0] - k * 10 * frame.shape[0]/288) - 1)
                        cv2.circle(frame,ppp,5,(0,255,0),-1)
        cv2.imshow("image", frame)
        key = cv2.waitKey(3)
        if(key == 27):
            cv2.destroyAllWindows()
            break