| 62 | return len(self.flows) |
| 63 | |
| 64 | def __iter__(self): |
| 65 | for flow_path in self.flows: |
| 66 | input_path = flow_path.replace( |
| 67 | self.path_prefix, os.path.join(self.data_path, 'clean')) |
| 68 | frame_id = int(input_path[-8:-4]) |
| 69 | input_a_path = '%s%04i.png' % (input_path[:-8], frame_id) |
| 70 | input_b_path = '%s%04i.png' % (input_path[:-8], frame_id + 1) |
| 71 | |
| 72 | input_a = cv2.imread(input_a_path) |
| 73 | input_b = cv2.imread(input_b_path) |
| 74 | flow = Flow.read(flow_path) |
| 75 | |
| 76 | # most implementation just crop the center |
| 77 | # which seems to be accepted practise |
| 78 | h, w = input_a.shape[:2] |
| 79 | newh = (h // 64) * 64 |
| 80 | neww = (w // 64) * 64 |
| 81 | aug = imgaug.CenterCrop((newh, neww)) |
| 82 | input_a = aug.augment(input_a) |
| 83 | input_b = aug.augment(input_b) |
| 84 | flow = aug.augment(flow) |
| 85 | yield [input_a, input_b, flow] |
| 86 | |
| 87 | |
| 88 | def inference(model, model_path, sintel_path): |