Spaces:
Sleeping
Sleeping
| import logging | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torchvision | |
| from torchvision.models.feature_extraction import create_feature_extractor | |
| import feature_extractor_models as smp | |
| import torch | |
| from .base import BaseModel | |
| logger = logging.getLogger(__name__) | |
| class FeatureExtractor(BaseModel): | |
| default_conf = { | |
| "pretrained": True, | |
| "input_dim": 3, | |
| "output_dim": 128, # # of channels in output feature maps | |
| "encoder": "resnet50", # torchvision net as string | |
| "remove_stride_from_first_conv": False, | |
| "num_downsample": None, # how many downsample block | |
| "decoder_norm": "nn.BatchNorm2d", # normalization ind decoder blocks | |
| "do_average_pooling": False, | |
| "checkpointed": False, # whether to use gradient checkpointing | |
| "architecture":"FPN" | |
| } | |
| mean = [0.485, 0.456, 0.406] | |
| std = [0.229, 0.224, 0.225] | |
| # self.fmodel=None | |
| def build_encoder(self, conf): | |
| assert isinstance(conf.encoder, str) | |
| if conf.pretrained: | |
| assert conf.input_dim == 3 | |
| # return encoder, layers | |
| def _init(self, conf): | |
| # Preprocessing | |
| self.register_buffer("mean_", torch.tensor(self.mean), persistent=False) | |
| self.register_buffer("std_", torch.tensor(self.std), persistent=False) | |
| if conf.architecture=="FPN": | |
| # Encoder | |
| self.fmodel = smp.FPN( | |
| encoder_name=conf.encoder, # choose encoder, e.g. mobilenet_v2 or efficientnet-b7 | |
| encoder_weights="imagenet", # use `imagenet` pre-trained weights for encoder initialization | |
| in_channels=conf.input_dim, # model input channels (1 for gray-scale images, 3 for RGB, etc.) | |
| classes=conf.output_dim, # model output channels (number of classes in your dataset) | |
| upsampling=2, # optional, final output upsampling, default is 8 | |
| activation=None | |
| ) | |
| elif conf.architecture == "LightFPN": | |
| self.fmodel = smp.L( | |
| encoder_name=conf.encoder, # choose encoder, e.g. mobilenet_v2 or efficientnet-b7 | |
| encoder_weights="imagenet", # use `imagenet` pre-trained weights for encoder initialization | |
| in_channels=conf.input_dim, # model input channels (1 for gray-scale images, 3 for RGB, etc.) | |
| classes=conf.output_dim, # model output channels (number of classes in your dataset) | |
| upsampling=2, # optional, final output upsampling, default is 8 | |
| activation=None | |
| ) | |
| elif conf.architecture=="PSP": | |
| self.fmodel =smp.PSPNet( | |
| encoder_name=conf.encoder, # choose encoder, e.g. mobilenet_v2 or efficientnet-b7 | |
| encoder_weights="imagenet", # use `imagenet` pre-trained weights for encoder initialization | |
| in_channels=conf.input_dim, # model input channels (1 for gray-scale images, 3 for RGB, etc.) | |
| classes=conf.output_dim, # model output channels (number of classes in your dataset) | |
| upsampling=4, # optional, final output upsampling, default is 8 | |
| activation=None | |
| ) | |
| else: | |
| raise ValueError("Only FPN") | |
| # elif conf.architecture=="Unet": | |
| # self.fmodel = smp.FPN( | |
| # encoder_name=conf.encoder, # choose encoder, e.g. mobilenet_v2 or efficientnet-b7 | |
| # encoder_weights="imagenet", # use `imagenet` pre-trained weights for encoder initialization | |
| # in_channels=conf.input_dim, # model input channels (1 for gray-scale images, 3 for RGB, etc.) | |
| # classes=conf.output_dim, # model output channels (number of classes in your dataset) | |
| # # upsampling=int(conf.upsampling), # optional, final output upsampling, default is 8 | |
| # activation="relu" | |
| # ) | |
| def _forward(self, data): | |
| image = data["image"] | |
| image = (image - self.mean_[:, None, None]) / self.std_[:, None, None] | |
| output = self.fmodel(image) | |
| # output = self.decoder(skip_features) | |
| pred = {"feature_maps": [output]} | |
| return pred | |
| if __name__ == '__main__': | |
| model=FeatureExtractor() |