Main class for segmentation model creation
| 12 | from . import get_encoder, get_decoder, get_classification |
| 13 | |
| 14 | class PCBSegNet: |
| 15 | """ |
| 16 | Main class for segmentation model creation |
| 17 | """ |
| 18 | def __init__(self, opt): |
| 19 | """ |
| 20 | Args: |
| 21 | opt: training options |
| 22 | """ |
| 23 | self.model_type = opt["model_type"] |
| 24 | self.image_height = opt["datasets"]["train"]["img_size_h"] |
| 25 | self.image_width = opt["datasets"]["train"]["img_size_w"] |
| 26 | self.num_classes = opt["train"]["num_classes"] + 1 |
| 27 | |
| 28 | def build(self): |
| 29 | """ |
| 30 | build encoder, decoder and final model |
| 31 | """ |
| 32 | encoder, learning_layer1, learning_layer2 = get_encoder(self.image_height, self.image_width) |
| 33 | model = get_decoder(encoder, learning_layer1, learning_layer2, self.num_classes) |
| 34 | print(model.summary()) |
| 35 | return model |
| 36 | |
| 37 | class PCBClassNet: |
| 38 | """ |