Skip to content

Commit

Permalink
fix: loadable dataset
Browse files Browse the repository at this point in the history
We should fix the transform as it generates ugly code in the dataset side.
  • Loading branch information
eloy-encord committed Feb 23, 2024
1 parent 7e86787 commit e1091b3
Showing 1 changed file with 8 additions and 4 deletions.
12 changes: 8 additions & 4 deletions clip_eval/dataset/encord_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,9 +32,12 @@ def __getitem__(self, idx):
img = Image.open(frame_path)
label = self._labels[idx]

if self.transform:
img, label = self.transform(img, label)
return img, label
if self.transform is not None:
_d = self.transform(dict(image=[img], label=[label]))
res_item = dict(image=_d["image"][0], label=_d["label"][0])
else:
res_item = dict(image=img, label=label)
return res_item

def __len__(self):
return len(self._frame_paths)
Expand Down Expand Up @@ -95,6 +98,7 @@ def _setup(

self._frame_paths = []
self._labels = []
class_name_to_idx = {name: idx for idx, name in enumerate(self.class_names)} # Fast lookup of class indices
for label_row in self._project.list_label_rows_v2():
anns_path = self._get_label_row_annotations(label_row)
label_row.from_labels_dict(json.loads(anns_path.read_text(encoding="utf-8")))
Expand All @@ -111,4 +115,4 @@ def _setup(
continue

self._frame_paths.append(self._get_frame_path(label_row, frame_view.frame))
self._labels.append(clf_answer.title)
self._labels.append(class_name_to_idx[clf_answer.title])

0 comments on commit e1091b3

Please sign in to comment.