| 50 | |
| 51 | |
| 52 | class SintelData(DataFlow): |
| 53 | |
| 54 | def __init__(self, data_path): |
| 55 | super(SintelData, self).__init__() |
| 56 | self.data_path = data_path |
| 57 | self.path_prefix = os.path.join(data_path, 'flow') |
| 58 | assert os.path.isdir(self.path_prefix), self.path_prefix |
| 59 | self.flows = glob.glob(os.path.join(self.path_prefix, '*', '*.flo')) |
| 60 | |
| 61 | def size(self): |
| 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): |
no outgoing calls
no test coverage detected
searching dependent graphs…