Python cv2.IMWRITE_PNG_COMPRESSION Examples

def process_image(fid):
    fid = fid + '.png'

    used_msks = []

    for pr_f in pred_folders:
        msk1 = cv2.imread(path.join('/wdata/', pr_f, '{0}.png'.format(fid.split('.')[0])), cv2.IMREAD_UNCHANGED)
        used_msks.append(msk1)

    msk = np.zeros_like(used_msks[0], dtype='float')

    for i in range(len(pred_folders)):
        p = used_msks[i]
        msk += (coefs[i] * p.astype('float'))
    msk /= np.sum(coefs)

    cv2.imwrite(path.join('/wdata/merged_oof', fid), msk.astype('uint8'), [cv2.IMWRITE_PNG_COMPRESSION, 9]) 

def process_image(img_id):
    if 'Pan-Sharpen_' in img_id:
        img_id = img_id.split('Pan-Sharpen_')[1]
    img = io.imread(path.join(test_dir, '_'.join(img_id.split('_')[:4]), 'Pan-Sharpen', 'Pan-Sharpen_' + img_id+'.tif'))
    nir = img[:, :, 3:]
    img = img[:, :, :3]
    np.clip(img, None, threshold, out=img)
    img = np.floor_divide(img, threshold / 255).astype('uint8')
    cv2.imwrite(path.join(test_png, img_id + '.png'), img, [cv2.IMWRITE_PNG_COMPRESSION, 9])

    img2 = io.imread(path.join(test_dir, '_'.join(img_id.split('_')[:4]), 'MS', 'MS_' + img_id+'.tif'))
    img2 = np.rollaxis(img2, 0, 3)
    img2 = cv2.resize(img2, (900, 900), interpolation=cv2.INTER_LANCZOS4)
    
    img_0_3_5 = (np.clip(img2[..., [0, 3, 5]], None, (2000, 3000, 3000)) / (np.array([2000, 3000, 3000]) / 255)).astype('uint8')
    cv2.imwrite(path.join(test_png2, img_id + '.png'), img_0_3_5, [cv2.IMWRITE_PNG_COMPRESSION, 9])
    
    pan = io.imread(path.join(test_dir, '_'.join(img_id.split('_')[:4]), 'PAN', 'PAN_' + img_id+'.tif'))
    pan = pan[..., np.newaxis]
    img_pan_6_7 = np.concatenate([pan, img2[..., 7:], nir], axis=2)
    img_pan_6_7 = (np.clip(img_pan_6_7, None, (3000, 5000, 5000)) / (np.array([3000, 5000, 5000]) / 255)).astype('uint8')
    cv2.imwrite(path.join(test_png3, img_id + '.png'), img_pan_6_7, [cv2.IMWRITE_PNG_COMPRESSION, 9]) 

def shape_chosen(self, event):
        tmp = self.get_item_text(event.GetItem())
        if not tmp:
            return
        else:
            chosen_shape = ('the shape is ' + tmp)
            set_goal(self.engine, chosen_shape)
            results, matched_facts, hit_rules = main_run(self.engine)
            source_image = cv2.imread(self.pic_path)
            detection_image = np.zeros(source_image.shape, np.uint8)
            draw_lines(detection_image, matched_facts, self.contour_num)
            cv2.imwrite('detection.png', detection_image, [int(cv2.IMWRITE_PNG_COMPRESSION), 0])
            self.show_picture('detection.png', (420, 30))
            self.resultText.Clear()
            self.resultText.WriteText(get_result(results))
            self.matchedFactText.Clear()
            self.matchedFactText.WriteText(get_matched_facts(matched_facts, self.contour_num))
            self.hitRuleText.Clear()
            self.hitRuleText.WriteText(get_hit_rules(hit_rules, self.contour_num)) 

def writer():
    # Writer is run in a separate thread, so that writing is concurrent
    # to taking the pictures.
    global q, job_finished, still_writing
    write_time = Stopwatch()
    #writeParams = (int(cv2.IMWRITE_PNG_COMPRESSION),7)
    while not job_finished:
	still_writing = True
	while not q.empty():
	    write_time.start()
            fn,img = q.get()
	    try:
                cv2.imwrite(fn,img, fileSaveParams)
                t=write_time.stop()
                print('Written {} in {:.02f} secs'.format(fn,t))
	    except:
		t=write_time.stop()
		print('Failed to write {} in {:.02f} secs'.format(fn,t))
    # Finished all jobs and queue is empty
    still_writing = False 

def save_img(img, filename):

    print("Save %s" %filename)

    if img.ndim == 3:
        img = img[:, :, ::-1] ### RGB to BGR
    
    ## clip to [0, 1]
    img = np.clip(img, 0, 1)

    ## quantize to [0, 255]
    img = np.uint8(img * 255.0)

    cv2.imwrite(filename, img, [cv2.IMWRITE_PNG_COMPRESSION, 0])


######################################################################################
##  Flow utility
###################################################################################### 

def save_img(img, img_name):
    save_img = img.squeeze().clamp(0, 1).numpy().transpose(1,2,0)
    # save img
    save_dir=os.path.join(opt.output,opt.test_dataset)
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
        
    save_fn = save_dir +'/'+ img_name
    cv2.imwrite(save_fn, cv2.cvtColor(save_img*255, cv2.COLOR_BGR2RGB),  [cv2.IMWRITE_PNG_COMPRESSION, 0]) 

def create_image_summary(name, val):
    """
    Args:
        name(str):
        val(np.ndarray): 4D tensor of NHWC. assume RGB if C==3.
            Can be either float or uint8. Range has to be [0,255].

    Returns:
        tf.Summary:
    """
    assert isinstance(name, six.string_types), type(name)
    n, h, w, c = val.shape
    val = val.astype('uint8')
    s = tf.Summary()
    imparams = [cv2.IMWRITE_PNG_COMPRESSION, 9]
    for k in range(n):
        arr = val[k]
        # CV2 will only write correctly in BGR chanel order
        if c == 3:
            arr = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
        elif c == 4:
            arr = cv2.cvtColor(arr, cv2.COLOR_RGBA2BGRA)
        tag = name if n == 1 else '{}/{}'.format(name, k)
        retval, img_str = cv2.imencode('.png', arr, imparams)
        if not retval:
            # Encoding has failed.
            continue
        img_str = img_str.tostring()

        img = tf.Summary.Image()
        img.height = h
        img.width = w
        # 1 - grayscale 3 - RGB 4 - RGBA
        img.colorspace = c
        img.encoded_image_string = img_str
        s.value.add(tag=tag, image=img)
    return s 

def create_mask(img_id, data_dir):
    labels_dir = os.path.join(data_dir, "labels")
    masks_dir = os.path.join(data_dir, "masks_all")
    os.makedirs(labels_dir, exist_ok=True)
    os.makedirs(masks_dir, exist_ok=True)

    labels = cv2.imread(path.join(labels_dir, '{0}.tif'.format(img_id)), cv2.IMREAD_UNCHANGED)
    final_mask = np.zeros((labels.shape[0], labels.shape[1], 3))
    if np.sum(labels) == 0:
        cv2.imwrite(path.join(masks_dir, '{0}.png'.format(img_id)), final_mask, [cv2.IMWRITE_PNG_COMPRESSION, 9])
        return final_mask

    ships_num = np.max(labels)
    if ships_num > 0:
        for i in range(1, ships_num + 1):
            ship_mask = np.zeros_like(labels, dtype='bool')
            ship_mask[labels == i] = 1
            area = np.sum(ship_mask)
            if area < 200:
                contour_size = 1
            elif area < 500:
                contour_size = 2
            else:
                contour_size = 3
            eroded = binary_erosion(ship_mask, iterations=contour_size)
            countour_mask = ship_mask ^ eroded
            final_mask[..., 0] += ship_mask
            final_mask[..., 1] += countour_mask
    final_mask[..., 2] = create_separation(labels)
    msk = np.clip(final_mask * 255, 0, 255)
    cv2.imwrite(path.join(masks_dir, '{0}.png'.format(img_id)), msk, [cv2.IMWRITE_PNG_COMPRESSION, 9]) 

def main():
    
    i_t_dir = os.path.join(cfg.data_dir, cfg.i_t_dir)
    i_s_dir = os.path.join(cfg.data_dir, cfg.i_s_dir)
    t_sk_dir = os.path.join(cfg.data_dir, cfg.t_sk_dir)
    t_t_dir = os.path.join(cfg.data_dir, cfg.t_t_dir)
    t_b_dir = os.path.join(cfg.data_dir, cfg.t_b_dir)
    t_f_dir = os.path.join(cfg.data_dir, cfg.t_f_dir)
    mask_t_dir = os.path.join(cfg.data_dir, cfg.mask_t_dir)
    
    makedirs(i_t_dir)
    makedirs(i_s_dir)
    makedirs(t_sk_dir)
    makedirs(t_t_dir)
    makedirs(t_b_dir)
    makedirs(t_f_dir)
    makedirs(mask_t_dir)

    mp_gen = multiprocess_datagen(cfg.process_num, cfg.data_capacity)
    mp_gen.multiprocess_runningqueue()
    digit_num = len(str(cfg.sample_num)) - 1
    for idx in range(cfg.sample_num):
        print ("Generating step {:>6d} / {:>6d}".format(idx + 1, cfg.sample_num))
        i_t, i_s, t_sk, t_t, t_b, t_f, mask_t = mp_gen.dequeue_data()
        i_t_path = os.path.join(i_t_dir, str(idx).zfill(digit_num) + '.png')
        i_s_path = os.path.join(i_s_dir, str(idx).zfill(digit_num) + '.png')
        t_sk_path = os.path.join(t_sk_dir, str(idx).zfill(digit_num) + '.png')
        t_t_path = os.path.join(t_t_dir, str(idx).zfill(digit_num) + '.png')
        t_b_path = os.path.join(t_b_dir, str(idx).zfill(digit_num) + '.png')
        t_f_path = os.path.join(t_f_dir, str(idx).zfill(digit_num) + '.png')
        mask_t_path = os.path.join(cfg.data_dir, cfg.mask_t_dir, str(idx).zfill(digit_num) + '.png')
        cv2.imwrite(i_t_path, i_t, [int(cv2.IMWRITE_PNG_COMPRESSION), 0])
        cv2.imwrite(i_s_path, i_s, [int(cv2.IMWRITE_PNG_COMPRESSION), 0])
        cv2.imwrite(t_sk_path, t_sk, [int(cv2.IMWRITE_PNG_COMPRESSION), 0])
        cv2.imwrite(t_t_path, t_t, [int(cv2.IMWRITE_PNG_COMPRESSION), 0])
        cv2.imwrite(t_b_path, t_b, [int(cv2.IMWRITE_PNG_COMPRESSION), 0])
        cv2.imwrite(t_f_path, t_f, [int(cv2.IMWRITE_PNG_COMPRESSION), 0])
        cv2.imwrite(mask_t_path, mask_t, [int(cv2.IMWRITE_PNG_COMPRESSION), 0])

    mp_gen.terminate_pool() 

def get_save_args(self):
        """ Return the save parameters for the file format """
        filetype = self.config["format"]
        args = list()
        if filetype == "jpg" and self.config["jpg_quality"] > 0:
            args = (cv2.IMWRITE_JPEG_QUALITY,  # pylint: disable=no-member
                    self.config["jpg_quality"])
        if filetype == "png" and self.config["png_compress_level"] > -1:
            args = (cv2.IMWRITE_PNG_COMPRESSION,  # pylint: disable=no-member
                    self.config["png_compress_level"])
        logger.debug(args)
        return args 

def save_img(img, filename):

    print("Save %s" %filename)
    ## clip to [0, 1]
    img = np.clip(img, 0, 1)
    ## quantize to [0, 255]
    img = np.uint8(img * 255.0)
    cv2.imwrite(filename, img, [cv2.IMWRITE_PNG_COMPRESSION, 0]) 

def save_result(save_dir, result, name, mode):

    # save output images
    o_sk, o_t, o_b, o_f = result
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    cv2.imwrite(os.path.join(save_dir, name + 'o_f.png'), o_f, [int(cv2.IMWRITE_PNG_COMPRESSION), 0])
    if mode == 1:
        cv2.imwrite(os.path.join(save_dir, name + 'o_sk.png'), o_sk, [int(cv2.IMWRITE_PNG_COMPRESSION), 0])
        cv2.imwrite(os.path.join(save_dir, name + 'o_t.png'), o_t, [int(cv2.IMWRITE_PNG_COMPRESSION), 0])
        cv2.imwrite(os.path.join(save_dir, name + 'o_b.png'), o_b, [int(cv2.IMWRITE_PNG_COMPRESSION), 0]) 

def worker(path, opt):
    crop_sz = opt['crop_sz']
    step = opt['step']
    thres_sz = opt['thres_sz']
    img_name = osp.basename(path)
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)

    n_channels = len(img.shape)
    if n_channels == 2:
        h, w = img.shape
    elif n_channels == 3:
        h, w, c = img.shape
    else:
        raise ValueError('Wrong image shape - {}'.format(n_channels))

    h_space = np.arange(0, h - crop_sz + 1, step)
    if h - (h_space[-1] + crop_sz) > thres_sz:
        h_space = np.append(h_space, h - crop_sz)
    w_space = np.arange(0, w - crop_sz + 1, step)
    if w - (w_space[-1] + crop_sz) > thres_sz:
        w_space = np.append(w_space, w - crop_sz)

    index = 0
    for x in h_space:
        for y in w_space:
            index += 1
            if n_channels == 2:
                crop_img = img[x:x + crop_sz, y:y + crop_sz]
            else:
                crop_img = img[x:x + crop_sz, y:y + crop_sz, :]
            crop_img = np.ascontiguousarray(crop_img)
            cv2.imwrite(
                osp.join(opt['save_folder'],
                         img_name.replace('.png', '_s{:03d}.png'.format(index))), crop_img,
                [cv2.IMWRITE_PNG_COMPRESSION, opt['compression_level']])
    return 'Processing {:s} ...'.format(img_name) 

def save_img(img, img_name):
    save_img = img.squeeze().clamp(0, 1).numpy().transpose(1,2,0)

    # save img
    save_dir=os.path.join(opt.output,opt.test_dataset)
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
        
    save_fn = save_dir +'/'+ img_name
    cv2.imwrite(save_fn, cv2.cvtColor(save_img*255, cv2.COLOR_BGR2RGB),  [cv2.IMWRITE_PNG_COMPRESSION, 0]) 

def save_depth_map(save_path, depth_map,
                   version='cv2', png_compression=3):
    """Saves depth map to disk as uint16 png

    Args:
        save_path: path to save depth map
        depth_map: depth map numpy array [h w]
        version: 'cv2' or 'pypng'
        png_compression: Only when version is 'cv2', sets png compression level.
            A lower value is faster with larger output,
            a higher value is slower with smaller output.
    """

    # Convert depth map to a uint16 png
    depth_image = (depth_map * 256.0).astype(np.uint16)

    if version == 'cv2':
        cv2.imwrite(save_path, depth_image, [cv2.IMWRITE_PNG_COMPRESSION, png_compression])

    elif version == 'pypng':
        with open(save_path, 'wb') as f:
            depth_image = (depth_map * 256.0).astype(np.uint16)
            writer = png.Writer(width=depth_image.shape[1],
                                height=depth_image.shape[0],
                                bitdepth=16,
                                greyscale=True)
            writer.write(f, depth_image)

    else:
        raise ValueError('Invalid version', version) 

def worker(path, save_folder, crop_sz, step, thres_sz, compression_level):
    img_name = os.path.basename(path)
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)

    n_channels = len(img.shape)
    if n_channels == 2:
        h, w = img.shape
    elif n_channels == 3:
        h, w, c = img.shape
    else:
        raise ValueError('Wrong image shape - {}'.format(n_channels))

    h_space = np.arange(0, h - crop_sz + 1, step)
    if h - (h_space[-1] + crop_sz) > thres_sz:
        h_space = np.append(h_space, h - crop_sz)
    w_space = np.arange(0, w - crop_sz + 1, step)
    if w - (w_space[-1] + crop_sz) > thres_sz:
        w_space = np.append(w_space, w - crop_sz)

    index = 0
    for x in h_space:
        for y in w_space:
            index += 1
            if n_channels == 2:
                crop_img = img[x:x + crop_sz, y:y + crop_sz]
            else:
                crop_img = img[x:x + crop_sz, y:y + crop_sz, :]
            crop_img = np.ascontiguousarray(crop_img)
            # var = np.var(crop_img / 255)
            # if var > 0.008:
            #     print(img_name, index_str, var)
            cv2.imwrite(
                os.path.join(save_folder, img_name.replace(
                    '.png', '_s{:03d}.png'.format(index))),
                crop_img, [cv2.IMWRITE_PNG_COMPRESSION, compression_level])
    return 'Processing {:s} ...'.format(img_name) 

def worker(path, opt):
    crop_sz = opt['crop_sz']
    step = opt['step']
    thres_sz = opt['thres_sz']
    img_name = osp.basename(path)
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)

    n_channels = len(img.shape)
    if n_channels == 2:
        h, w = img.shape
    elif n_channels == 3:
        h, w, c = img.shape
    else:
        raise ValueError('Wrong image shape - {}'.format(n_channels))

    h_space = np.arange(0, h - crop_sz + 1, step)
    if h - (h_space[-1] + crop_sz) > thres_sz:
        h_space = np.append(h_space, h - crop_sz)
    w_space = np.arange(0, w - crop_sz + 1, step)
    if w - (w_space[-1] + crop_sz) > thres_sz:
        w_space = np.append(w_space, w - crop_sz)

    index = 0
    for x in h_space:
        for y in w_space:
            index += 1
            if n_channels == 2:
                crop_img = img[x:x + crop_sz, y:y + crop_sz]
            else:
                crop_img = img[x:x + crop_sz, y:y + crop_sz, :]
            crop_img = np.ascontiguousarray(crop_img)
            cv2.imwrite(
                osp.join(opt['save_folder'],
                         img_name.replace('.png', '_s{:03d}.png'.format(index))), crop_img,
                [cv2.IMWRITE_PNG_COMPRESSION, opt['compression_level']])
    return 'Processing {:s} ...'.format(img_name) 

def map(self, data):
        image = data[self.src]
        image_id = data["image_id"]

        # Prepare output for image based on image_id
        output = image_id.split(os.path.sep)
        dirname = output[:-1]
        if len(dirname) > 0:
            dirname = os.path.join(*dirname)
            dirname = os.path.join(self.path, dirname)
            os.makedirs(dirname, exist_ok=True)
        else:
            dirname = self.path
        filename = f"{output[-1].rsplit('.', 1)[0]}.{self.image_ext}"
        path = os.path.join(dirname, filename)

        if self.image_ext == "jpg":
            cv2.imwrite(path, image,
                        (cv2.IMWRITE_JPEG_QUALITY, self.jpg_quality) if self.jpg_quality else None)
        elif self.image_ext == "png":
            cv2.imwrite(path, image,
                        (cv2.IMWRITE_PNG_COMPRESSION, self.png_compression) if self.png_compression else None)
        else:
            raise Exception("Unsupported image format")

        return data 

def pack_img(header, img, quality=95, img_fmt='.jpg'):
    """Pack an image into ``MXImageRecord``.

    Parameters
    ----------
    header : IRHeader
        Header of the image record.
        ``header.label`` can be a number or an array. See more detail in ``IRHeader``.
    img : numpy.ndarray
        Image to be packed.
    quality : int
        Quality for JPEG encoding in range 1-100, or compression for PNG encoding in range 1-9.
    img_fmt : str
        Encoding of the image (.jpg for JPEG, .png for PNG).

    Returns
    -------
    s : str
        The packed string.

    Examples
    --------
    >>> label = 4 # label can also be a 1-D array, for example: label = [1,2,3]
    >>> id = 2574
    >>> header = mx.recordio.IRHeader(0, label, id, 0)
    >>> img = cv2.imread('test.jpg')
    >>> packed_s = mx.recordio.pack_img(header, img)
    """
    assert cv2 is not None
    jpg_formats = ['.JPG', '.JPEG']
    png_formats = ['.PNG']
    encode_params = None
    if img_fmt.upper() in jpg_formats:
        encode_params = [cv2.IMWRITE_JPEG_QUALITY, quality]
    elif img_fmt.upper() in png_formats:
        encode_params = [cv2.IMWRITE_PNG_COMPRESSION, quality]

    ret, buf = cv2.imencode(img_fmt, img, encode_params)
    assert ret, 'failed to encode image'
    return pack(header, buf.tostring()) 

def create_image_summary(name, val):
    """
    Args:
        name(str):
        val(np.ndarray): 4D tensor of NHWC. assume RGB if C==3.
            Can be either float or uint8. Range has to be [0,255].

    Returns:
        tf.Summary:
    """
    assert isinstance(name, six.string_types), type(name)
    n, h, w, c = val.shape
    val = val.astype('uint8')
    s = tf.Summary()
    imparams = [cv2.IMWRITE_PNG_COMPRESSION, 9]
    for k in range(n):
        arr = val[k]
        # CV2 will only write correctly in BGR chanel order
        if c == 3:
            arr = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
        elif c == 4:
            arr = cv2.cvtColor(arr, cv2.COLOR_RGBA2BGRA)
        tag = name if n == 1 else '{}/{}'.format(name, k)
        retval, img_str = cv2.imencode('.png', arr, imparams)
        if not retval:
            # Encoding has failed.
            continue
        img_str = img_str.tostring()

        img = tf.Summary.Image()
        img.height = h
        img.width = w
        # 1 - grayscale 3 - RGB 4 - RGBA
        img.colorspace = c
        img.encoded_image_string = img_str
        s.value.add(tag=tag, image=img)
    return s 

def create_image_summary(name, val):
    """
    Args:
        name(str):
        val(np.ndarray): 4D tensor of NHWC. assume RGB if C==3.
            Can be either float or uint8. Range has to be [0,255].

    Returns:
        tf.Summary:
    """
    assert isinstance(name, six.string_types), type(name)
    n, h, w, c = val.shape
    val = val.astype('uint8')
    s = tf.Summary()
    imparams = [cv2.IMWRITE_PNG_COMPRESSION, 9]
    for k in range(n):
        arr = val[k]
        # CV2 will only write correctly in BGR chanel order
        if c == 3:
            arr = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
        elif c == 4:
            arr = cv2.cvtColor(arr, cv2.COLOR_RGBA2BGRA)
        tag = name if n == 1 else '{}/{}'.format(name, k)
        retval, img_str = cv2.imencode('.png', arr, imparams)
        if not retval:
            # Encoding has failed.
            continue
        img_str = img_str.tostring()

        img = tf.Summary.Image()
        img.height = h
        img.width = w
        # 1 - grayscale 3 - RGB 4 - RGBA
        img.colorspace = c
        img.encoded_image_string = img_str
        s.value.add(tag=tag, image=img)
    return s 

def saveTensorToImages(t, idxs, save_to_path):
    
    if os.path.exists(save_to_path)==False:
        os.mkdir(save_to_path)
    
    for i in range(t.size(0)):
        im = t[i,:,:,:].detach().data.cpu().numpy() 
        im = np.transpose(im, (1,2,0)).astype(np.uint16)
        cv2.imwrite(os.path.join(save_to_path, str(idxs[i].data.cpu().numpy()).zfill(10)+'.png'), im, [cv2.IMWRITE_PNG_COMPRESSION, 4] ) 

def save_img(img, img_name):
    save_img = img.squeeze().clamp(0, 1).numpy().transpose(1,2,0)

    # save img    
    save_fn = 'im'+img_name+'_'+opt.model_type+opt.prefix+'.png'
    cv2.imwrite(save_fn, cv2.cvtColor(save_img*255, cv2.COLOR_BGR2RGB),  [cv2.IMWRITE_PNG_COMPRESSION, 0]) 

def save_img(img, d_dir,img_name, pred_flag):
    save_img = img.squeeze().clamp(0, 1).numpy().transpose(1,2,0)
    filename = os.path.splitext(img_name)

    # save img
    save_dir=os.path.join(opt.output, d_dir)
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
        
    if pred_flag:
        save_fn = save_dir +'/'+ filename[0]+'_'+opt.model_type+filename[1]
    else:
        save_fn = save_dir +'/'+ img_name
    cv2.imwrite(save_fn, cv2.cvtColor(save_img*255, cv2.COLOR_BGR2RGB),  [cv2.IMWRITE_PNG_COMPRESSION, 0]) 

def save_img(img, d_dir,img_name, pred_flag):
    save_img = img.squeeze().clamp(0, 1).numpy().transpose(1,2,0)
    filename = os.path.splitext(img_name)

    # save img
    save_dir=os.path.join(opt.output, d_dir)
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
        
    if pred_flag:
        save_fn = save_dir +'/'+ filename[0]+'_'+opt.model_type+filename[1]
    else:
        save_fn = save_dir +'/'+ img_name
    cv2.imwrite(save_fn, cv2.cvtColor(save_img*255, cv2.COLOR_BGR2RGB),  [cv2.IMWRITE_PNG_COMPRESSION, 0]) 

def save_img(img, img_name):
    save_img = img.squeeze().clamp(0, 1).numpy().transpose(1,2,0)

    # save img    
    save_fn = 'im'+img_name+'_'+opt.model_type+opt.prefix+'.png'
    cv2.imwrite(save_fn, cv2.cvtColor(save_img*255, cv2.COLOR_BGR2RGB),  [cv2.IMWRITE_PNG_COMPRESSION, 0]) 

def save_img(img, img_name):
    save_img = img.squeeze().clamp(0, 1).numpy().transpose(1,2,0)

    # save img    
    save_fn = 'im'+img_name+'_'+opt.model_type+opt.prefix+'_'+opt.train_obj+'.png'
    cv2.imwrite(save_fn, cv2.cvtColor(save_img*255, cv2.COLOR_BGR2RGB),  [cv2.IMWRITE_PNG_COMPRESSION, 0]) 

def save_img(img, img_name, pred_flag):
    save_img = img.squeeze().clamp(0, 1).numpy().transpose(1,2,0)

    # save img
    save_dir=os.path.join(args.output, args.data_dir, os.path.splitext(args.file_list)[0]+'_'+str(args.upscale_factor)+'x')
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
        
    if pred_flag:
        save_fn = save_dir +'/'+ img_name+'_'+args.model_type+'F'+str(args.nFrames)+'.png'
    else:
        save_fn = save_dir +'/'+ img_name+'.png'
    cv2.imwrite(save_fn, cv2.cvtColor(save_img*255, cv2.COLOR_BGR2RGB),  [cv2.IMWRITE_PNG_COMPRESSION, 0]) 

def worker(path, save_folder, crop_sz, step, thres_sz, compression_level):
    img_name = os.path.basename(path)
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)

    n_channels = len(img.shape)
    if n_channels == 2:
        h, w = img.shape
    elif n_channels == 3:
        h, w, c = img.shape
    else:
        raise ValueError('Wrong image shape - {}'.format(n_channels))

    h_space = np.arange(0, h - crop_sz + 1, step)
    if h - (h_space[-1] + crop_sz) > thres_sz:
        h_space = np.append(h_space, h - crop_sz)
    w_space = np.arange(0, w - crop_sz + 1, step)
    if w - (w_space[-1] + crop_sz) > thres_sz:
        w_space = np.append(w_space, w - crop_sz)

    index = 0
    for x in h_space:
        for y in w_space:
            index += 1
            if n_channels == 2:
                crop_img = img[x:x + crop_sz, y:y + crop_sz]
            else:
                crop_img = img[x:x + crop_sz, y:y + crop_sz, :]
            crop_img = np.ascontiguousarray(crop_img)
            # var = np.var(crop_img / 255)
            # if var > 0.008:
            #     print(img_name, index_str, var)
            cv2.imwrite(
                os.path.join(save_folder, img_name.replace('.png', '_s{:03d}.png'.format(index))),
                crop_img, [cv2.IMWRITE_PNG_COMPRESSION, compression_level])
    return 'Processing {:s} ...'.format(img_name) 

def worker(path, save_folder, mode, compression_level):
    img_name = os.path.basename(path)
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)  # BGR
    if mode == 'gray':
        img_y = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    else:
        img_y = bgr2ycbcr(img, only_y=True)
    cv2.imwrite(os.path.join(save_folder, img_name), img_y,
                [cv2.IMWRITE_PNG_COMPRESSION, compression_level])
    return 'Processing {:s} ...'.format(img_name)