1515"""Common utils."""
1616import contextlib
1717import os
18- from typing import Text , Tuple , Union
18+ from typing import Text , Tuple , Union , Dict
1919from absl import logging
2020import numpy as np
2121import tensorflow .compat .v1 as tf
@@ -482,11 +482,12 @@ def archive_ckpt(ckpt_eval, ckpt_objective, ckpt_path):
482482 return True
483483
484484
485- def parse_image_size (image_size : Union [Text , int , Tuple [int , int ]]):
485+ def parse_image_size (image_size : Union [Text , int , Tuple [int , int ], Dict [ int , int ] ]):
486486 """Parse the image size and return (height, width).
487487
488488 Args:
489- image_size: A integer, a tuple (H, W), or a string with HxW format.
489+ image_size: A integer, a tuple (H, W), a string with HxW format,
490+ or a dict with keys 'height' and 'width' with corresponding int values.
490491
491492 Returns:
492493 A tuple of integer (height, width).
@@ -496,10 +497,13 @@ def parse_image_size(image_size: Union[Text, int, Tuple[int, int]]):
496497 return (image_size , image_size )
497498
498499 if isinstance (image_size , str ):
499- # image_size is a string with format WxH
500- width , height = image_size .lower ().split ('x' )
500+ # image_size is a string with format HxW
501+ height , width = image_size .lower ().split ('x' )
501502 return (int (height ), int (width ))
502503
504+ if isinstance (image_size , dict ):
505+ return (image_size ['height' ], image_size ['width' ])
506+
503507 if isinstance (image_size , tuple ):
504508 return image_size
505509
0 commit comments