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