- def img_descr(x):
- u = x.split('<img>', 1)
- return u[1] if len(u) > 1 else ''
-
- img = torch.full((len(descr), 3, height, width), 255)
- d = [ img_descr(x) for x in descr ]
- d = [ u.strip().split(' ')[:height * width] for u in d ]
- d = [ u + [ '<unk>' ] * (height * width - len(u)) for u in d ]
- d = [ [ token2color(t) for t in u ] for u in d ]
- img = torch.tensor(d).permute(0, 2, 1)
- img = img.reshape(img.size(0), 3, height, width)
+ d = descr.split('<img>', 1)
+ d = d[-1] if len(d) > 1 else ''
+ d = d.strip().split(' ')[:height * width]
+ d = d + [ '<unk>' ] * (height * width - len(d))
+ d = [ token2color(t) for t in d ]
+ img = torch.tensor(d).permute(1, 0)
+ img = img.reshape(1, 3, height, width)