| import torch |
| from torch.utils.data import Dataset, DataLoader |
| import random |
| import json |
| import os |
| from tqdm import tqdm |
|
|
| class PairwiseOriginalDataset(Dataset): |
| def __init__( |
| self, |
| json_list, |
| soft_label=False, |
| confidence_threshold=None, |
| ): |
| self.samples = [] |
| for json_file in json_list: |
| with open(json_file, "r") as f: |
| data = json.load(f) |
| self.samples.extend(data) |
|
|
| self.soft_label = soft_label |
| self.confidence_threshold = confidence_threshold |
|
|
| if confidence_threshold is not None: |
| new_samples = [] |
| for sample in tqdm( |
| self.samples, desc="Filtering samples according to confidence threshold" |
| ): |
| if sample.get("confidence", float("inf")) >= confidence_threshold: |
| new_samples.append(sample) |
| self.samples = new_samples |
|
|
| def __len__(self): |
| return len(self.samples) |
|
|
| def __getitem__(self, idx): |
| while True: |
| index = idx |
| try: |
| return self.get_single_item(index) |
| except Exception as e: |
| print(f"Error processing sample at index {idx}: {e}") |
| import traceback |
| traceback.print_exc() |
| index = random.randint(0, len(self.samples) - 1) |
| if index == idx: |
| continue |
| idx = index |
|
|
| def get_single_item(self, idx): |
| sample = self.samples[idx] |
| |
| image_1 = sample["path1"] |
| image_2 = sample["path2"] |
| assert os.path.exists(image_1) and os.path.exists(image_2), f'{image_1} or {image_2}' |
| text_1 = sample["prompt"] |
| text_2 = sample["prompt"] |
|
|
| |
| if self.soft_label: |
| choice_dist = sorted(sample["choice_dist"], reverse=True) |
| assert ( |
| torch.sum(torch.tensor(choice_dist)) > 0 |
| ), "Choice distribution cannot be zero." |
| label = torch.tensor(choice_dist[0]) / torch.sum(torch.tensor(choice_dist)) |
| else: |
| label = torch.tensor(1).float() |
| |
| return { |
| "image_1": image_1, |
| "image_2": image_2, |
| "text_1": text_1, |
| "text_2": text_2, |
| "label": label, |
| "confidence": sample.get("confidence", 1.0), |
| "choice_dist": torch.tensor(sample.get("choice_dist", [1.0, 0.0])), |
| } |
|
|
