def setup_args():
parser = argparse.ArgumentParser(description='Run model on image dataset')
parser.add_argument('model',
type=str,
choices=models.keys(),
help='model architecture')
parser.add_argument('-m',
'--metric',
type=str,
choices=['mse'],
default='mse',
help='metric trained against (default: %(default)s)')
parser.add_argument('dataset', type=str, help='dataset path')
parser.add_argument('-c',
'--entropy-coder',
choices=compressai.available_entropy_coders(),
default=compressai.available_entropy_coders()[0],
help='Entropy coder (default: %(default)s)')
parser.add_argument('-q',
'--quality',
dest='qualities',
nargs='+',
type=int,
default=range(1, 9))
parser.add_argument(
'--entropy-estimation',
action='store_true',
help='Use evaluated entropy estimation (no entropy coding)')
return parser
def encode(argv):
parser = argparse.ArgumentParser(description='Encode image to bit-stream')
parser.add_argument('image', type=str)
parser.add_argument('--model',
choices=models.keys(),
default=list(models.keys())[0],
help='NN model to use (default: %(default)s)')
parser.add_argument('-m',
'--metric',
choices=['mse'],
default='mse',
help='metric trained against (default: %(default)s')
parser.add_argument('-q',
'--quality',
choices=list(range(1, 9)),
type=int,
default=3,
help='Quality setting (default: %(default)s)')
parser.add_argument('-c',
'--coder',
choices=compressai.available_entropy_coders(),
default=compressai.available_entropy_coders()[0],
help='Entropy coder (default: %(default)s)')
parser.add_argument('-o', '--output', help='Output path')
args = parser.parse_args(argv)
if not args.output:
args.output = Path(Path(args.image).resolve().name).with_suffix('.bin')
_encode(args.image, args.model, args.metric, args.quality, args.coder,
args.output)
def __init__(self, method):
if not isinstance(method, str):
raise ValueError(f'Invalid method type "{type(method)}"')
from compressai import available_entropy_coders
if method not in available_entropy_coders():
methods = ", ".join(available_entropy_coders())
raise ValueError(f'Unknown entropy coder "{method}"'
f" (available: {methods})")
if method == "ans":
from compressai import ans
encoder = ans.RansEncoder()
decoder = ans.RansDecoder()
elif method == "rangecoder":
import range_coder
encoder = range_coder.RangeEncoder()
decoder = range_coder.RangeDecoder()
self.name = method
self._encoder = encoder
self._decoder = decoder
def decode(argv):
parser = argparse.ArgumentParser(description='Decode bit-stream to imager')
parser.add_argument('input', type=str)
parser.add_argument('-c',
'--coder',
choices=compressai.available_entropy_coders(),
default=compressai.available_entropy_coders()[0],
help='Entropy coder (default: %(default)s)')
parser.add_argument('--show', action='store_true')
parser.add_argument('-o', '--output', help='Output path')
args = parser.parse_args(argv)
_decode(args.input, args.coder, args.show, args.output)
def decode(argv):
parser = argparse.ArgumentParser(description="Decode bit-stream to imager")
parser.add_argument("input", type=str)
parser.add_argument(
"-c",
"--coder",
choices=compressai.available_entropy_coders(),
default=compressai.available_entropy_coders()[0],
help="Entropy coder (default: %(default)s)",
)
parser.add_argument("--show", action="store_true")
parser.add_argument("-o", "--output", help="Output path")
args = parser.parse_args(argv)
_decode(args.input, args.coder, args.show, args.output)
def decode(argv):
parser = argparse.ArgumentParser(
description="Decode bit-stream to image/video")
parser.add_argument("input", type=str)
parser.add_argument(
"-c",
"--coder",
choices=compressai.available_entropy_coders(),
default=compressai.available_entropy_coders()[0],
help="Entropy coder (default: %(default)s)",
)
parser.add_argument("--show", action="store_true")
parser.add_argument("-o", "--output", help="Output path")
parser.add_argument("--cuda", action="store_true", help="Use cuda")
args = parser.parse_args(argv)
device = "cuda" if args.cuda and torch.cuda.is_available() else "cpu"
_decode(args.input, args.coder, args.show, device, args.output)
def __init__(self, method):
if not isinstance(method, str):
raise ValueError(f'Invalid method type "{type(method)}"')
from compressai import available_entropy_coders
if method not in available_entropy_coders():
methods = ', '.join(available_entropy_coders())
raise ValueError(f'Unknown entropy coder "{method}"'
f' (available: {methods})')
if method == 'ans':
from compressai import ans # pylint: disable=E0611
encoder = ans.RansEncoder()
decoder = ans.RansDecoder()
elif method == 'rangecoder':
import range_coder # pylint: disable=E0401
encoder = range_coder.RangeEncoder()
decoder = range_coder.RangeDecoder()
self._encoder = encoder
self._decoder = decoder
def encode(argv):
parser = argparse.ArgumentParser(description="Encode image to bit-stream")
parser.add_argument("image", type=str)
parser.add_argument(
"--model",
choices=models.keys(),
default=list(models.keys())[0],
help="NN model to use (default: %(default)s)",
)
parser.add_argument(
"-m",
"--metric",
choices=["mse"],
default="mse",
help="metric trained against (default: %(default)s",
)
parser.add_argument(
"-q",
"--quality",
choices=list(range(1, 9)),
type=int,
default=3,
help="Quality setting (default: %(default)s)",
)
parser.add_argument(
"-c",
"--coder",
choices=compressai.available_entropy_coders(),
default=compressai.available_entropy_coders()[0],
help="Entropy coder (default: %(default)s)",
)
parser.add_argument("-o", "--output", help="Output path")
args = parser.parse_args(argv)
if not args.output:
args.output = Path(Path(args.image).resolve().name).with_suffix(".bin")
_encode(args.image, args.model, args.metric, args.quality, args.coder,
args.output)
def encode(argv):
parser = argparse.ArgumentParser(
description="Encode image/video to bit-stream")
parser.add_argument(
"input",
type=str,
help=
"Input path, the first frame will be encoded with a NN image codec if the input is a raw yuv sequence",
)
parser.add_argument(
"-f",
"--num_of_frames",
default=-1,
type=int,
help=
"Number of frames to be coded. -1 will encode all frames of input (default: %(default)s)",
)
parser.add_argument(
"--model",
choices=models.keys(),
default=list(models.keys())[0],
help="NN model to use (default: %(default)s)",
)
parser.add_argument(
"-m",
"--metric",
choices=metric_ids.keys(),
default="mse",
help="metric trained against (default: %(default)s)",
)
parser.add_argument(
"-q",
"--quality",
choices=list(range(1, 9)),
type=int,
default=3,
help="Quality setting (default: %(default)s)",
)
parser.add_argument(
"-c",
"--coder",
choices=compressai.available_entropy_coders(),
default=compressai.available_entropy_coders()[0],
help="Entropy coder (default: %(default)s)",
)
parser.add_argument("-o", "--output", help="Output path")
parser.add_argument("--cuda", action="store_true", help="Use cuda")
args = parser.parse_args(argv)
if not args.output:
args.output = Path(Path(args.input).resolve().name).with_suffix(".bin")
device = "cuda" if args.cuda and torch.cuda.is_available() else "cpu"
_encode(
args.input,
args.num_of_frames,
args.model,
args.metric,
args.quality,
args.coder,
device,
args.output,
)
def test_available_entropy_coders():
rv = compressai.available_entropy_coders()
assert isinstance(rv, list)
assert "ans" in rv
def setup_args():
parent_parser = argparse.ArgumentParser(add_help=False, )
# Common options.
parent_parser.add_argument("dataset", type=str, help="dataset path")
parent_parser.add_argument(
"-a",
"--arch",
type=str,
choices=pretrained_models.keys(),
help="model architecture",
required=True,
)
parent_parser.add_argument(
"-c",
"--entropy-coder",
choices=compressai.available_entropy_coders(),
default=compressai.available_entropy_coders()[0],
help="entropy coder (default: %(default)s)",
)
parent_parser.add_argument(
"--entropy-estimation",
action="store_true",
help="use evaluated entropy estimation (no entropy coding)",
)
parent_parser.add_argument(
"-v",
"--verbose",
action="store_true",
help="verbose mode",
)
parser = argparse.ArgumentParser(
description="Evaluate a model on an image dataset.", add_help=True)
subparsers = parser.add_subparsers(help="model source",
dest="source",
required=True)
# Options for pretrained models
pretrained_parser = subparsers.add_parser("pretrained",
parents=[parent_parser])
pretrained_parser.add_argument(
"-m",
"--metric",
type=str,
choices=["mse", "ms-ssim"],
default="mse",
help="metric trained against (default: %(default)s)",
)
pretrained_parser.add_argument(
"-q",
"--quality",
dest="qualities",
nargs="+",
type=int,
default=(1, ),
)
checkpoint_parser = subparsers.add_parser("checkpoint",
parents=[parent_parser])
checkpoint_parser.add_argument(
"-p",
"--path",
dest="paths",
type=str,
nargs="*",
required=True,
help="checkpoint path",
)
return parser
def create_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description="Video compression network evaluation.",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parent_parser = argparse.ArgumentParser(add_help=False)
parent_parser.add_argument("dataset", type=str, help="sequences directory")
parent_parser.add_argument("output", type=str, help="output directory")
parent_parser.add_argument(
"-a",
"--architecture",
type=str,
choices=models.keys(),
help="model architecture",
required=True,
)
parent_parser.add_argument(
"-f", "--force", action="store_true", help="overwrite previous runs"
)
parent_parser.add_argument("--cuda", action="store_true", help="use cuda")
parent_parser.add_argument("--half", action="store_true", help="use AMP")
parent_parser.add_argument(
"--entropy-estimation",
action="store_true",
help="use evaluated entropy estimation (no entropy coding)",
)
parent_parser.add_argument(
"-c",
"--entropy-coder",
choices=compressai.available_entropy_coders(),
default=compressai.available_entropy_coders()[0],
help="entropy coder (default: %(default)s)",
)
parent_parser.add_argument(
"--keep_binaries",
action="store_true",
help="keep bitstream files in output directory",
)
parent_parser.add_argument(
"-v",
"--verbose",
action="store_true",
help="verbose mode",
)
parent_parser.add_argument(
"-m",
"--metric",
type=str,
choices=["mse", "ms-ssim"],
default="mse",
help="metric trained against (default: %(default)s)",
)
subparsers = parser.add_subparsers(help="model source", dest="source")
subparsers.required = True
# Options for pretrained models
pretrained_parser = subparsers.add_parser("pretrained", parents=[parent_parser])
pretrained_parser.add_argument(
"-q",
"--quality",
dest="qualities",
nargs="+",
type=int,
default=(1,),
)
checkpoint_parser = subparsers.add_parser("checkpoint", parents=[parent_parser])
checkpoint_parser.add_argument(
"-p",
"--path",
dest="paths",
type=str,
nargs="*",
required=True,
help="checkpoint path",
)
return parser
