def create_model(num_classes=21, device=torch.device('cpu')):
# https://download.pytorch.org/models/resnet50-19c8e357.pth
# pre_train_path = "./src/resnet50.pth"
backbone = Backbone()
model = SSD300(backbone=backbone, num_classes=num_classes)
# https://ngc.nvidia.com/catalog/models -> search ssd -> download FP32
pre_ssd_path = "./src/nvidia_ssdpyt_fp32.pt"
pre_model_dict = torch.load(pre_ssd_path, map_location=device)
pre_weights_dict = pre_model_dict["model"]
# 删除类别预测器权重，注意，回归预测器的权重可以重用，因为不涉及num_classes
del_conf_loc_dict = {}
for k, v in pre_weights_dict.items():
split_key = k.split(".")
if "conf" in split_key:
continue
del_conf_loc_dict.update({k: v})
missing_keys, unexpected_keys = model.load_state_dict(del_conf_loc_dict,
strict=False)
if len(missing_keys) != 0 or len(unexpected_keys) != 0:
print("missing_keys: ", missing_keys)
print("unexpected_keys: ", unexpected_keys)
return model
def create_model(num_classes=21, device=torch.device('cpu')):
# https://download.pytorch.org/models/resnet50-19c8e357.pth
# pre_train_path = "./src/resnet50.pth"
backbone = Backbone()
model = SSD300(backbone=backbone, num_classes=num_classes)
# # https://ngc.nvidia.com/catalog/models -> search ssd -> download FP32
# pre_ssd_path = "./src/nvidia_ssdpyt_fp32_190826.pt"
# if os.path.exists(pre_ssd_path) is False:
# raise FileNotFoundError("nvidia_ssdpyt_fp32.pt not find in {}".format(pre_ssd_path))
# pre_model_dict = torch.load(pre_ssd_path, map_location=device)
# pre_weights_dict = pre_model_dict["model"]
#
# # 删除类别预测器权重，注意，回归预测器的权重可以重用，因为不涉及num_classes
# del_conf_loc_dict = {}
# for k, v in pre_weights_dict.items():
# split_key = k.split(".")
# if "conf" in split_key:
# continue
# del_conf_loc_dict.update({k: v})
#
# missing_keys, unexpected_keys = model.load_state_dict(del_conf_loc_dict, strict=False)
# if len(missing_keys) != 0 or len(unexpected_keys) != 0:
# print("missing_keys: ", missing_keys)
# print("unexpected_keys: ", unexpected_keys)
return model
def load_network(weights_path, device, num_classes):
backbone = Backbone()
model = SSD300(backbone=backbone, num_classes=num_classes)
#net = create_model(num_classes=6)
model.load_state_dict(
torch.load(weights_path, map_location=device)["model"])
return model
def create_model(num_classes, device=torch.device('cpu')):
# https://download.pytorch.org/models/resnet50-19c8e357.pth
pre_train_path = "./src/resnet50.pth"
backbone = Backbone(pretrain_path=pre_train_path)
model = SSD300(backbone=backbone, num_classes=num_classes)
pre_ssd_path = "./src/nvidia_ssdpyt_fp32.pt"
pre_model_dict = torch.load(pre_ssd_path, map_location=device)
pre_weights_dict = pre_model_dict["model"]
# delete weights of prediction categories
del_conf_loc_dict = {}
for k, v in pre_weights_dict.items():
split_key = k.split(".")
if "conf" in split_key:
continue
del_conf_loc_dict.update({k: v})
missing_keys, unexpected_keys = model.load_state_dict(del_conf_loc_dict,
strict=False)
if len(missing_keys) != 0 or len(unexpected_keys) != 0:
print("missing_keys: ", missing_keys)
print("unexpected_keys: ", unexpected_keys)
return model
def __init__(self, device="cuda:1"):
picfile = str(time.strftime("%Y%m%d"))
self.todaypath = os.path.join('/workspace/nologopics', picfile)
if not os.path.exists(self.todaypath):
os.mkdir(self.todaypath)
self.device = torch.device(device)
# logo检测模型
backbone = Backbone()
self.ssdmodel = SSD300(backbone=backbone, num_classes=2)
modelpath = './weights/ssd300-best.pth'
weights_dict = torch.load(modelpath, map_location=device)
self.ssdmodel.load_state_dict(weights_dict, strict=False)
json_file = open('./pascal_voc_classes.json', 'r')
class_dict = json.load(json_file)
self.category_index = {v: k for k, v in class_dict.items()}
self.data_transforms = transform.Compose([
transform.Resize(),
transform.ToTensor(),
transform.Normalization()
])
# 水印字体
self.font = ImageFont.truetype("./src/msyh.TTF", 24, encoding="utf-8")
# 爬虫网址
self.spiderurl = {
#clear_log
2: {
'url': 'http://adsoc.qknode.com/adagent/material/material?',
'topic': ["清理", "日历", "天气"]
},
0: {
'url': 'http://adsoc.qknode.com/adagent/material/center/rank?',
'topic': ["清理", "日历", "天气", "教育"]
},
# 排行榜
1: {
'url': 'http://adsoc.qknode.com/adagent/material/material?',
'topic': ["清理", "日历", "天气"]
}
# 素材洞察
}
# 推送地址
self.finalurl = 'http://adsoc.qknode.com/adagent/material/center/push'
# self.cnniqamodel = CNNIQAnet(ker_size=7, n_kers=50, n1_nodes=800, n2_nodes=800)
# self.cnniqamodel.load_state_dict(torch.load('./weights/CNNIQA-LIVE.pth',map_location=device))
if device != 'cpu':
# self.cnniqamodel = self.cnniqamodel.to(self.device)
# self.cnniqamodel.eval()
self.ssdmodel = self.ssdmodel.to(self.device)
self.ssdmodel.eval()
def create_model(num_classes):
backbone = Backbone()
model = SSD300(backbone=backbone, num_classes=num_classes)
return model