text: import argparse
  from pathlib import Path
  import os
  from tqdm import tqdm
  from PIL import Image
  import cv2
  import torch
  import torchvision.transforms.functional as tvf
  from torch.utils.data.dataset import Dataset
  from torch.utils.data.dataloader import DataLoader
  
  from mycv.paths import IMAGENET_DIR
  from mycv.utils.image import psnr_dB
  from mycv.utils.coding import MS_SSIM
  
  
  def get_img_names(root):
      classes = os.listdir(root)
      image_names = []
      for cname in classes:
          cdir = root / cname
          imnames = os.listdir(cdir)
          new = [f'{cname}/{imn}' for imn in imnames]
          image_names.extend(new)
      return image_names
  
  
  def get_img_paths(root):
      classes = os.listdir(root)
      image_paths = []
      for cname in classes:
          cdir = root / cname
          imnames = os.listdir(cdir)
          new = [cdir / imn for imn in imnames]
          image_paths.extend(new)
      return image_paths
  
  
  class JPEGEncode(Dataset):
      def __init__(self, src_dir, tgt_dir, quality) -> None:
          self.img_names = get_img_names(src_dir)
          self.src_dir = Path(src_dir)
          self.tgt_dir = Path(tgt_dir)
          self.cv2_param = [cv2.IMWRITE_JPEG_QUALITY, quality]
  
          assert not tgt_dir.is_dir(), f'{tgt_dir} already exists.'
          print(f'Creating {tgt_dir}...')
          tgt_dir.mkdir(parents=False, exist_ok=False)
          classes = os.listdir(self.src_dir)
          for cname in classes:
              (tgt_dir / cname).mkdir(parents=False, exist_ok=False)
  
      def __len__(self):
          return len(self.img_names)
  
      def __getitem__(self, index):
          imname = self.img_names[index]
          src_path = str(self.src_dir / imname)
          im = cv2.imread(src_path)
          # save image to a new folder
          assert len(imname.split('.')) == 2, f"{imname} should contain only one '.'"
          svname = imname.split('.')[0] + '.jpg'
          svpath = self.tgt_dir / svname
          assert svpath.parent.is_dir()
          cv2.imwrite(str(svpath), im, self.cv2_param)
          return 0
  
  
  def jpeg_encode(split='val', quality=50, workers=8):
      src_dir = IMAGENET_DIR / split
      tgt_dir = IMAGENET_DIR / f'{split}_jpeg{quality}'
      jpeg = JPEGEncode(src_dir, tgt_dir, quality)
      dataloader = DataLoader(jpeg, batch_size=max(workers*8, 16), num_workers=workers)
      for imgs in tqdm(dataloader):
          do_nothing = 1
  
  
  def compute_bpp(split='val_jpeg50'):
      img_dir = IMAGENET_DIR / split
      img_names = get_img_paths(img_dir)
  
      total_bits = 0
      total_pixels = 0
      pbar = tqdm(img_names)
      for imname in pbar:
          impath = img_dir / imname
          bytes = os.path.getsize(impath)
  
          im = Image.open(impath)
          pixels = im.height * im.width
          im.close()
  
          bits = bytes * 8
          total_bits += bits
          total_pixels += pixels
          bpp = total_bits / total_pixels
          pbar.set_description(f'bpp: {bpp:.4g}')
  
      print(total_bits / total_pixels)
  
  
  def jpeg_stats(split='val', quality=50):
      img_dir = IMAGENET_DIR / split
      # save_dir = IMAGENET_DIR / f'{split}_jpeg{quality}'
      # if not save_dir.is_dir():
      #     print(f'{save_dir} does not exist. Creating it...')
      #     save_dir.mkdir(parents=False, exist_ok=False)
  
      encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]
      msssim_func = MS_SSIM(max_val=1.0)
      img_paths = get_img_paths(img_dir)
      pbar = tqdm(enumerate(img_paths), total=len(img_paths))
      psnr_sum, msssim_sum, bpp_sum = 0, 0, 0
      for bi, impath in pbar:
          im = cv2.imread(str(impath))
          assert im is not None
          nH, nW, _ = im.shape
  
          # encode and decode
          flag, bits = cv2.imencode('.jpg', im, params=encode_param)
          assert flag
          im_dec = cv2.imdecode(bits, 1)
          # plt.imshow(im_dec[:,:,::-1]); plt.show()
  
          # compute metrics
          ps = psnr_dB(im, im_dec)
          with torch.no_grad():
              ms = msssim_func(
                  img1=tvf.to_tensor(im_dec).unsqueeze(0).cuda(),
                  img2=tvf.to_tensor(im).unsqueeze(0).cuda()
              )
          bpp = len(bits)*8 / (nH*nW)
  
          # loggging
          wnid = impath.parent.stem
          _num = bi + 1
          msssim_sum += ms
          psnr_sum += ps
          bpp_sum += bpp
          msssim = msssim_sum / _num
          psnr   = psnr_sum   / _num
          bpp    = bpp_sum    / _num
          msg = f'{wnid}... MS: {msssim:.4g}, PSNR: {psnr:.4g}, bpp: {bpp:.4g}'
          pbar.set_description(msg)
      print(bpp, psnr, msssim)
  
  
  if __name__ == '__main__':
      parser = argparse.ArgumentParser()
      parser.add_argument('--split',   type=str, default='val224l')
      parser.add_argument('--quality', type=int, default=28)
      parser.add_argument('--mode',    type=str, default='encode', choices=['stats', 'encode', 'bpp'])
      parser.add_argument('--workers', type=int, default=8)
      args = parser.parse_args()
  
      if args.mode == 'stats':
          jpeg_stats(split=args.split, quality=args.quality)
      elif args.mode == 'encode':
          jpeg_encode(split=args.split, quality=args.quality, workers=args.workers)
      elif args.mode == 'bpp':
          compute_bpp(split=f'{args.split}_jpeg{args.quality}')
      else:
          raise ValueError()
      # get_bpp(split=f'{args.split}_jpeg{args.quality}')
  