| 463 | |
| 464 | |
| 465 | class VGGDataset(Dataset): |
| 466 | |
| 467 | def __init__(self, path_vid): |
| 468 | self.data_complete = pd.read_json('./data/data.json',orient='records',dtype=False,lines=True) |
| 469 | |
| 470 | self.vid = [] |
| 471 | with open('./data/vids/'+path_vid, "r") as fr: |
| 472 | for line in fr.readlines(): |
| 473 | self.vid.append(line.strip()) |
| 474 | self.data = self.data_complete[self.data_complete.video_id.isin(self.vid)] |
| 475 | self.data['video_id'] = self.data['video_id'].astype('category') |
| 476 | self.data['video_id'].cat.set_categories(self.vid, inplace=True) |
| 477 | self.data.sort_values('video_id', ascending=True, inplace=True) |
| 478 | self.data.reset_index(inplace=True) |
| 479 | |
| 480 | self.framefeapath='./data/ptvgg19_frames/' |
| 481 | |
| 482 | |
| 483 | def __len__(self): |
| 484 | return self.data.shape[0] |
| 485 | |
| 486 | def __getitem__(self, idx): |
| 487 | item = self.data.iloc[idx] |
| 488 | vid = item['video_id'] |
| 489 | |
| 490 | label = 1 if item['annotation']=='假' else 0 |
| 491 | label = torch.tensor(label) |
| 492 | |
| 493 | frames=pickle.load(open(os.path.join(self.framefeapath,vid+'.pkl'),'rb')) |
| 494 | frames=torch.FloatTensor(frames) |
| 495 | |
| 496 | return { |
| 497 | 'label': label, |
| 498 | 'frames': frames, |
| 499 | } |
| 500 | |
| 501 | |
| 502 | class BboxDataset(Dataset): |