Dennis Trujillo commited on
Commit
6459c2a
·
1 Parent(s): 55ab5b2

added frame num to __get_item__

Browse files
Files changed (1) hide show
  1. ds1.py +3 -1
ds1.py CHANGED
@@ -39,6 +39,7 @@ class PatchDataset(Dataset):
39
  self.length = 0
40
  self.xy_positions = []
41
  self.patches=[]
 
42
  with open(ffile, 'rb') as f:
43
  for fNr in range(1,nFrames+1):
44
  BytesToSkip = fHead + fNr*NrPixels*NrPixels*2
@@ -65,6 +66,7 @@ class PatchDataset(Dataset):
65
  sub_img = thisFrame[start_y:end_y,start_x:end_x]
66
  self.patches.append(normalize_patch(sub_img))
67
  self.xy_positions.append([start_y,start_x])
 
68
  self.length += len(self.patches)
69
  f.close()
70
 
@@ -72,5 +74,5 @@ class PatchDataset(Dataset):
72
  return self.length
73
 
74
  def __getitem__(self, index):
75
- return self.patches[index], self.xy_positions[index]
76
 
 
39
  self.length = 0
40
  self.xy_positions = []
41
  self.patches=[]
42
+ self.f_nums=[]
43
  with open(ffile, 'rb') as f:
44
  for fNr in range(1,nFrames+1):
45
  BytesToSkip = fHead + fNr*NrPixels*NrPixels*2
 
66
  sub_img = thisFrame[start_y:end_y,start_x:end_x]
67
  self.patches.append(normalize_patch(sub_img))
68
  self.xy_positions.append([start_y,start_x])
69
+ self.f_nums.append(fNr)
70
  self.length += len(self.patches)
71
  f.close()
72
 
 
74
  return self.length
75
 
76
  def __getitem__(self, index):
77
+ return self.patches[index], self.xy_positions[index], self.f_nums[index]
78