import collections
import torch


def imdict_to_img(img_dict, n_per_row=8, n_rows=1):
  """Returns an image where each row corresponds to a key in the dictionary.
  Values are expected to be of format BxCxHxW.
  """
  od = collections.OrderedDict(sorted(img_dict.items()))
  imgs = torch.stack(list(od.values()), dim=0)      # AxBxCxHxW
  # make the rows
  imgs = torch.stack([imgs[:, i*n_per_row:(i+1)*n_per_row] for i in range(n_rows)], dim=0)    # n_rows x A x n_per_row x CxHxW
  imgs = imgs.flatten(0,2)
  return imgs