import os
from pathlib import Path
import argparse
from tqdm import tqdm
import cv2
import torch

from mycv.utils.image import pad_divisible


device = torch.device('cuda:0')

@torch.no_grad()
def main(args):
    print()

    if args.model.startswith('cheng'):
        from mycv.external.compressai.zoo import cheng2020_anchor
        q = int(args.model[-1])
        model = cheng2020_anchor(q, metric="mse", pretrained=True)
    else:
        raise ValueError()
    model.eval()
    model = model.to(device=device)

    src_dir = Path(args.src)
    tgt_dir = src_dir.parent / f'{src_dir.stem}_{args.model}'
    print(f'Saving to {tgt_dir}...')

    # from mycv.datasets.imcoding import nic_evaluate
    # results = nic_evaluate(model, input_norm=False, bar=True)
    # print(results)

    print('Loading image paths...')
    img_names = os.listdir(src_dir)

    for imname in tqdm(img_names):
        # print(impath)
        impath = src_dir / imname

        im = cv2.imread(str(impath))
        # cv2.imwrite('ori.png', im)
        im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
        imh, imw = im.shape[:2]
        im, _ = pad_divisible(im, div=16)
        im = torch.from_numpy(im).permute(2,0,1).float().div_(255.0)
        im = im.unsqueeze_(0)

        x = im.to(device=device)
        try:
            z = model.g_a(x)
            zhat = z.round_()
            xhat = model.g_s(zhat)
        except Exception as e:
            print(e)
            print(impath, imh, imw, im.shape)
            exit()

        rec = xhat.squeeze(0).clamp_(0,1).cpu()
        rec = rec.mul_(255).round_().to(dtype=torch.uint8).permute(1,2,0).numpy()
        rec = rec[:imh, :imw, :]

        rec = cv2.cvtColor(rec, cv2.COLOR_RGB2BGR)
        # cv2.imwrite('rec.png', rec)

        savepath = tgt_dir / f'{impath.stem}.png'
        if not savepath.parent.is_dir():
            savepath.parent.mkdir(parents=True)

        cv2.imwrite(str(savepath), rec)
        # break


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--src',   type=str, default='../datasets/coco/images/val2017')
    parser.add_argument('--model', type=str, default='cheng3')
    args = parser.parse_args()
    main(args)

import os
from pathlib import Path
import argparse
from tqdm import tqdm
import cv2
import torch

from mycv.utils.image import pad_divisible


device = torch.device('cuda:0')

@torch.no_grad()
def main(args):
    print()

    if args.model.startswith('cheng'):
        from mycv.external.compressai.zoo import cheng2020_anchor
        q = int(args.model[-1])
        model = cheng2020_anchor(q, metric="mse", pretrained=True)
    else:
        raise ValueError()
    model.eval()
    model = model.to(device=device)

    src_dir = Path(args.src)
    tgt_dir = src_dir.parent / f'{src_dir.stem}_{args.model}'
    print(f'Saving to {tgt_dir}...')

    # from mycv.datasets.imcoding import nic_evaluate
    # results = nic_evaluate(model, input_norm=False, bar=True)
    # print(results)

    print('Loading image paths...')
    img_names = os.listdir(src_dir)

    for imname in tqdm(img_names):
        # print(impath)
        impath = src_dir / imname

        im = cv2.imread(str(impath))
        # cv2.imwrite('ori.png', im)
        im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
        imh, imw = im.shape[:2]
        im, _ = pad_divisible(im, div=16)
        im = torch.from_numpy(im).permute(2,0,1).float().div_(255.0)
        im = im.unsqueeze_(0)

        x = im.to(device=device)
        try:
            z = model.g_a(x)
            zhat = z.round_()
            xhat = model.g_s(zhat)
        except Exception as e:
            print(e)
            print(impath, imh, imw, im.shape)
            exit()

        rec = xhat.squeeze(0).clamp_(0,1).cpu()
        rec = rec.mul_(255).round_().to(dtype=torch.uint8).permute(1,2,0).numpy()
        rec = rec[:imh, :imw, :]

        rec = cv2.cvtColor(rec, cv2.COLOR_RGB2BGR)
        # cv2.imwrite('rec.png', rec)

        savepath = tgt_dir / f'{impath.stem}.png'
        if not savepath.parent.is_dir():
            savepath.parent.mkdir(parents=True)

        cv2.imwrite(str(savepath), rec)
        # break


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--src',   type=str, default='../datasets/coco/images/val2017')
    parser.add_argument('--model', type=str, default='cheng3')
    args = parser.parse_args()
    main(args)

# from mycv.datasets.imcoding import nic_evaluate # results = nic_evaluate(model, input_norm=False, bar=True) # print(results) # print(impath) # cv2.imwrite('ori.png', im) # cv2.imwrite('rec.png', rec) # break