lichen/ieemoo-ai-isempty
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update lightrise.py.

Browse Source
This commit is contained in:
Brainway
2022-11-09 03:30:40 +00:00
committed by Gitee
parent e6270608d1
commit 7a1c101307
1 changed files with 16 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
lightrise.py
View File

@ -8,23 +8,23 @@ import argparse
from models.modeling import VisionTransformer, CONFIGS from models.modeling import VisionTransformer, CONFIGS
import time import time
parser = argparse.ArgumentParser()
parser.add_argument("--pretrained_model", type=str, default="../module/ieemoo-ai-isempty/model/new/ieemooempty_vitlight_checkpoint.pth", help="load pretrained model") #使用自定义VIT
args = parser.parse_args()
#args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
args.device = torch.device("cpu")
num_classes = 2
cls_dict = {0: "noemp", 1: "yesemp"}
test_transform = transforms.Compose([transforms.Resize((600, 600), Image.BILINEAR),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
#复杂场景小模型测试单张图片 #复杂场景小模型测试单张图片
def riseempty(imgdata): def riseempty(imgdata):
parser = argparse.ArgumentParser()
parser.add_argument("--pretrained_model", type=str, default="../module/ieemoo-ai-isempty/model/new/ieemooempty_vitlight_checkpoint.pth", help="load pretrained model") #使用自定义VIT
args = parser.parse_args()
#args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
args.device = torch.device("cpu")
num_classes = 2
cls_dict = {0: "noemp", 1: "yesemp"}
test_transform = transforms.Compose([transforms.Resize((600, 600), Image.BILINEAR),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
# 准备模型 # 准备模型
model = torch.load(args.pretrained_model,map_location=torch.device('cpu')) #自己预训练模型 model = torch.load(args.pretrained_model,map_location=torch.device('cpu')) #自己预训练模型
model.to(args.device) model.to(args.device)
@ -41,8 +41,8 @@ def riseempty(imgdata):
return riseresult return riseresult
if __name__ == "__main__": # if __name__ == "__main__":
riseresult = riseempty("light.jpg") # riseresult = riseempty("light.jpg")
print(riseresult) # print(riseresult)