| 500 | |
| 501 | |
| 502 | class BboxDataset(Dataset): |
| 503 | def __init__(self, path_vid): |
| 504 | self.data_complete = pd.read_json('./data/data_5500_revised.json',orient='records',dtype=False,lines=True) |
| 505 | |
| 506 | self.vid = [] |
| 507 | with open('./data/vids/'+path_vid, "r") as fr: |
| 508 | for line in fr.readlines(): |
| 509 | self.vid.append(line.strip()) |
| 510 | self.data = self.data_complete[self.data_complete.video_id.isin(self.vid)] |
| 511 | self.data['video_id'] = self.data['video_id'].astype('category') |
| 512 | self.data['video_id'].cat.set_categories(self.vid, inplace=True) |
| 513 | self.data.sort_values('video_id', ascending=True, inplace=True) |
| 514 | self.data.reset_index(inplace=True) |
| 515 | |
| 516 | self.bboxfeapath = './data/bbox_vgg19/' |
| 517 | |
| 518 | def __len__(self): |
| 519 | return self.data.shape[0] |
| 520 | |
| 521 | def __getitem__(self, idx): |
| 522 | item = self.data.iloc[idx] |
| 523 | vid = item['video_id'] |
| 524 | |
| 525 | label = 1 if item['annotation']=='假' else 0 |
| 526 | label = torch.tensor(label) |
| 527 | |
| 528 | bbox_vgg = pickle.load(open(os.path.join(self.bboxfeapath,vid+'.pkl'),'rb')) |
| 529 | bbox_vgg = torch.FloatTensor(bbox_vgg) |
| 530 | |
| 531 | return { |
| 532 | 'label': label, |
| 533 | 'bbox_vgg': bbox_vgg |
| 534 | } |
| 535 | |
| 536 | |
| 537 | class Title_W2V_Dataset(Dataset): |