<reponame>PhiDCH/mmsegmentation
from mmseg.apis import inference_segmentor, init_segmentor, show_result_pyplot
from mmseg.core.evaluation import get_palette

import os 
from tqdm import tqdm 
import numpy as np
import cv2

import argparse

def get_parser():
    parser = argparse.ArgumentParser(description="predict image depthmap")
    parser.add_argument("--img-path", type=str, help="path to image dataset")
    parser.add_argument("--save-path", type=str, help="path to save image depthmap")
    
    return parser

# img_path = '../gen_hcm/image'
# save_path = '../gen_hcm/seg'

def main():
    
    args = get_parser().parse_args()
    img_path = args.img_path
    save_path = args.save_path
    
    config_file = 'configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py'
    checkpoint_file = 'checkpoints/pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-2966598c.pth'
    model = init_segmentor(config_file, checkpoint_file, device='cuda:0')
    
    for path in tqdm(os.listdir(img_path)):
        if path.endswith('.jpg'):
            fullpath = os.path.join(img_path, path)
            result = inference_segmentor(model, fullpath)[0].astype(np.uint8)
            
            seg_path = os.path.join(save_path, path)
            cv2.imwrite(seg_path, result, [cv2.IMWRITE_JPEG_QUALITY, 75])
            
            
if __name__ == '__main__':
    main()
from mmseg.apis import inference_segmentor, init_segmentor, show_result_pyplot
from mmseg.core.evaluation import get_palette

import os 
from tqdm import tqdm 
import numpy as np
import cv2

import argparse

def get_parser():
    parser = argparse.ArgumentParser(description="predict image depthmap")
    parser.add_argument("--img-path", type=str, help="path to image dataset")
    parser.add_argument("--save-path", type=str, help="path to save image depthmap")
    
    return parser

# img_path = '../gen_hcm/image'
# save_path = '../gen_hcm/seg'

def main():
    
    args = get_parser().parse_args()
    img_path = args.img_path
    save_path = args.save_path
    
    config_file = 'configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py'
    checkpoint_file = 'checkpoints/pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-2966598c.pth'
    model = init_segmentor(config_file, checkpoint_file, device='cuda:0')
    
    for path in tqdm(os.listdir(img_path)):
        if path.endswith('.jpg'):
            fullpath = os.path.join(img_path, path)
            result = inference_segmentor(model, fullpath)[0].astype(np.uint8)
            
            seg_path = os.path.join(save_path, path)
            cv2.imwrite(seg_path, result, [cv2.IMWRITE_JPEG_QUALITY, 75])
            
            
if __name__ == '__main__':
    main()
# img_path = '../gen_hcm/image' # save_path = '../gen_hcm/seg'