| 46 | |
| 47 | |
| 48 | class random_3d_loader: |
| 49 | def __init__(self, batch_size): |
| 50 | np.random.seed(12345) |
| 51 | self.subdirs = ["SER00004", "SER00006", "SER00008", "SER00009", "SER00011", "SER00015"] |
| 52 | self.dirs = [os.path.join(db_3d_folder, x) for x in self.subdirs] |
| 53 | self.batch_size = batch_size |
| 54 | np.random.seed(1234) |
| 55 | self.n = 0 |
| 56 | self.order = list(range(len(self.subdirs))) |
| 57 | np.random.shuffle(self.order) |
| 58 | |
| 59 | def __iter__(self): |
| 60 | return self |
| 61 | |
| 62 | def __next__(self): |
| 63 | return [self.get_one() for _ in range(self.batch_size)] |
| 64 | |
| 65 | def get_one(self): |
| 66 | idx = self.get_index() |
| 67 | dir = self.dirs[idx] |
| 68 | |
| 69 | imgs = [] |
| 70 | i = 0 |
| 71 | path = os.path.join(dir, "%i.jpg") |
| 72 | while True: |
| 73 | fname = path % i |
| 74 | img = cv2.imread(fname, cv2.IMREAD_GRAYSCALE) |
| 75 | if img is None: |
| 76 | break |
| 77 | i += 1 |
| 78 | imgs.append(img[::-1, :, np.newaxis]) |
| 79 | return np.stack(imgs, axis=0) |
| 80 | |
| 81 | def get_index(self): |
| 82 | if self.n >= len(self.order): |
| 83 | np.random.shuffle(self.order) |
| 84 | self.n = 0 |
| 85 | idx = self.order[self.n] |
| 86 | self.n += 1 |
| 87 | return idx |
| 88 | |
| 89 | |
| 90 | def layout_str(dim, channel_first): |
no outgoing calls