Main class for classfication model creation
| 35 | return model |
| 36 | |
| 37 | class PCBClassNet: |
| 38 | """ |
| 39 | Main class for classfication model creation |
| 40 | """ |
| 41 | def __init__(self, opt): |
| 42 | """ |
| 43 | Args: |
| 44 | opt: training options |
| 45 | """ |
| 46 | self.model_type = opt["model_type"] |
| 47 | self.image_height = opt["datasets"]["train"]["img_size_h"] |
| 48 | self.image_width = opt["datasets"]["train"]["img_size_w"] |
| 49 | self.num_classes = opt["train"]["num_classes"] |
| 50 | |
| 51 | def build(self): |
| 52 | """ |
| 53 | build encoder and final model |
| 54 | """ |
| 55 | encoder = get_encoder(self.image_height, self.image_width) |
| 56 | model = get_classification(encoder, self.num_classes) |
| 57 | return model |
| 58 | |
| 59 | |
| 60 | def get_model(opt): |