update lightrise.py.

lightrise.py

@@ -8,6 +8,10 @@ import argparse
 from models.modeling import VisionTransformer, CONFIGS
 import time
 
+
+
+#复杂场景小模型测试单张图片
+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()
@@ -21,10 +25,6 @@ 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):
     # 准备模型
     model = torch.load(args.pretrained_model,map_location=torch.device('cpu')) #自己预训练模型
     model.to(args.device)
@@ -41,8 +41,8 @@ def riseempty(imgdata):
     
     return riseresult
 
-if __name__ == "__main__":
-    riseresult = riseempty("light.jpg")
-    print(riseresult)
+# if __name__ == "__main__":
+#     riseresult = riseempty("light.jpg")
+#     print(riseresult)