| 60 | class TempVideoDataset: |
| 61 | |
| 62 | def __init__(self, video_list, model_3dmm, if_align=False, cross_id=False, image_list=None, resize=256): |
| 63 | self.video_list = video_list |
| 64 | self.model_3dmm = model_3dmm |
| 65 | self.cross_id = cross_id |
| 66 | self.image_list = image_list |
| 67 | self.if_align = if_align |
| 68 | if self.cross_id and len(self.video_list) != len(self.image_list) and len(self.video_list) > 0: |
| 69 | self.video_list = self.video_list * (len(self.image_list) // len(self.video_list) + 1) |
| 70 | self.video_list = self.video_list[:len(self.image_list)] |
| 71 | |
| 72 | self.video_index = -1 |
| 73 | self.transform = transforms.Compose([ |
| 74 | transforms.Resize((resize, resize)), |
| 75 | transforms.ToTensor(), |
| 76 | transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True), |
| 77 | ]) |
| 78 | self.image_transform = transforms.Compose([ |
| 79 | transforms.ToTensor(), |
| 80 | transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True), |
| 81 | ]) |
| 82 | self.semantic_radius = 13 |
| 83 | self.croper = Croper() |
| 84 | |
| 85 | def __len__(self): |
| 86 | if len(self.video_list) > 0: |