Update utils.py

utils.py

@@ -136,12 +136,21 @@ def generate_similiarity_map(images, attn_map, all_bpe_strings, vis_list, target
     #     print("target_height",images.size[1])
     #     print("target_width",images.size[0])
 
-    images = images[0]
-    images_vis = T.ToTensor()(images) # images [<PIL.Image.Image image mode=RGB size=1024x608 at 0x7F9B6FC24B80>]
-    print("images",images)
-    print("images_vis",images_vis)
-    target_height = images.size[1]
-    target_width = images.size[0]
+    if len(images) == 1:
+        images_vis = torch.stack([T.ToTensor()(image) for image in images])
+    else:
+        images_vis = torch.stack([T.ToTensor()(image) for image in images[:-1]])
+    images_vis = make_grid(images_vis, nrow=target_aspect_ratio[0], padding=0, normalize=False)
+    target_width = image_size * target_aspect_ratio[0]
+    target_height = image_size * target_aspect_ratio[1]
+
+    
+    # images = images[0]
+    # images_vis = T.ToTensor()(images) # images [<PIL.Image.Image image mode=RGB size=1024x608 at 0x7F9B6FC24B80>]
+    # print("images",images)
+    # print("images_vis",images_vis)
+    # target_height = images.size[1]
+    # target_width = images.size[0]
 
     
     print("attn_map",attn_map.shape)# torch.Size([4, 76, 128])