Update youtubevis.py

3 years ago · 81a10937c7
parent f8a26e2719
commit 81a10937c7
1 changed files with 2 additions and 2 deletions
--- a/dataset/youtubevis.py
+++ b/dataset/youtubevis.py
@ -38,7 +38,7 @@ class YouTubeVISDataset(Dataset):
        
        self.index = []
        for video_id in self.videos.keys():
-            for frame in range(len(self.videos[video_id])):
+            for frame in range(len(self.videos[video_id]['file_names'])):
                self.index.append((video_id, frame))
                
    def __len__(self):
@ -120,4 +120,4 @@ class YouTubeVISAugmentation:
            imgs = F.hflip(imgs)
            segs = F.hflip(segs)
        
-        return imgs, segs
+        return imgs, segs