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

update ieemoo-ai-isempty.py.

Browse Source
This commit is contained in:
Brainway
2022-11-22 08:19:05 +00:00
committed by Gitee
parent 2a89533188
commit 5d5788fe6f
1 changed files with 26 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

82
ieemoo-ai-isempty.py
View File

@ -12,7 +12,6 @@ from PIL import Image
from torchvision import transforms
from models.modeling import VisionTransformer, CONFIGS
from vit_pytorch import ViT
#import lightrise
# import logging.config as log_config
sys.path.insert(0, ".")
@ -46,11 +45,11 @@ print(torch.__version__)
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--img_size", default=600, type=int, help="Resolution size")
parser.add_argument("--img_size", default=320, type=int, help="Resolution size")
parser.add_argument('--split', type=str, default='overlap', help="Split method")
parser.add_argument('--slide_step', type=int, default=2, help="Slide step for overlap split")
parser.add_argument('--smoothing_value', type=float, default=0.0, help="Label smoothing value")
parser.add_argument("--pretrained_model", type=str, default="../module/ieemoo-ai-isempty/model/new/ieemooempty_vit_checkpoint.pth", help="load pretrained model")
parser.add_argument("--pretrained_model", type=str, default="../module/ieemoo-ai-isempty/model/now/emptyjudge5_checkpoint.bin", help="load pretrained model")
#parser.add_argument("--pretrained_model", type=str, default="output/ieemooempty_vit_checkpoint.pth", help="load pretrained model") #使用自定义VIT
opt, unknown = parser.parse_known_args()
return opt
@ -66,7 +65,7 @@ class Predictor(object):
self.num_classes = 0
self.model = None
self.prepare_model()
self.test_transform = transforms.Compose([transforms.Resize((600, 600), Image.BILINEAR),
self.test_transform = transforms.Compose([transforms.Resize((320, 320), Image.BILINEAR),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
@ -83,7 +82,7 @@ class Predictor(object):
# self.model = torch.load(self.args.pretrained_model)
# else:
# self.model = torch.load(self.args.pretrained_model,map_location='cpu')
self.model = torch.load(self.args.pretrained_model)
self.model = torch.load(self.args.pretrained_model,map_location=torch.device('cpu'))
self.model.eval()
self.model.to("cuda")
@ -96,18 +95,16 @@ class Predictor(object):
else:
with torch.no_grad():
x = self.test_transform(img_data)
if torch.cuda.is_available():
x = x.cuda()
# if torch.cuda.is_available():
# x = x.cuda()
part_logits = self.model(x.unsqueeze(0))
probs = torch.nn.Softmax(dim=-1)(part_logits)
topN = torch.argsort(probs, dim=-1, descending=True).tolist()
clas_ids = topN[0][0]
clas_ids = 0 if 0==int(clas_ids) or 2 == int(clas_ids) or 3 == int(clas_ids) else 1
#print("cur_img result: class id: %d, score: %0.3f" % (clas_ids, probs[0, clas_ids].item()))
result={}
print("cur_img result: class id: %d, score: %0.3f" % (clas_ids, probs[0, clas_ids].item()))
result["success"] = "true"
result["rst_cls"] = str(clas_ids)
return result
@ -115,65 +112,38 @@ class Predictor(object):
args = parse_args()
predictor = Predictor(args)
def riseempty(imgdata):
risemodel = torch.load("../module/ieemoo-ai-isempty/model/new/ieemooempty_vitlight_checkpoint.pth",map_location=torch.device('cpu')) #自己预训练模型
risemodel.to("cpu")
risemodel.eval()
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])])
x = test_transform(imgdata)
part_logits = risemodel(x.unsqueeze(0).to('cpu'))
probs = torch.nn.Softmax(dim=-1)(part_logits)
top2 = torch.argsort(probs, dim=-1, descending=True)
riseclas_ids = top2[0][0]
#print("cur_img result: class id: %d, score: %0.3f" % (riseclas_ids, probs[0, riseclas_ids].item()))
riseresult={}
riseresult["success"] = "true"
riseresult["rst_cls"] = int(riseclas_ids)
return riseresult
@app.route("/isempty", methods=['POST'])
def get_isempty():
start = time.time()
#print('--------------------EmptyPredict-----------------')
print("begin")
data = request.get_data()
ip = request.remote_addr
#print('------ ip = %s ------' % ip)
print(ip)
json_data = json.loads(data.decode("utf-8"))
getdateend = time.time()
#print('get date use time: {0:.2f}s'.format(getdateend - start))
pic = json_data.get("pic")
result = {}
imgdata = base64.b64decode(pic)
result ={}
imgdata_np = np.frombuffer(imgdata, dtype='uint8')
img_src = cv2.imdecode(imgdata_np, cv2.IMREAD_COLOR)
cv2.imwrite('huanyuan.jpg',img_src)
#img_data = Image.fromarray(np.uint8(img_src)) #这个转换不能要，会导致判空错误增加
img_data = Image.open('huanyuan.jpg')
img_data = Image.fromarray(np.uint8(img_src))
result = predictor.normal_predict(img_data, result) # 1==empty, 0==nonEmpty
riseresult = riseempty(img_data)
#print(riseresult["rst_cls"])
if(int(result["rst_cls"])==1):
if(int(riseresult["rst_cls"])==1):
result = {}
result["success"] = "true"
result["rst_cls"] = 1
else:
result = {}
result["success"] = "true"
result["rst_cls"] = 0
return repr(result)
def getByte(path):
with open(path, 'rb') as f:
img_byte = base64.b64encode(f.read())
img_str = img_byte.decode('utf-8')
return img_str
if __name__ == "__main__":
app.run(host='0.0.0.0', port=8888)
# result ={}
# imgdata = base64.b64decode(getByte("img.jpg"))
# imgdata_np = np.frombuffer(imgdata, dtype='uint8')
# img_src = cv2.imdecode(imgdata_np, cv2.IMREAD_COLOR)
# img_data = Image.fromarray(np.uint8(img_src))
# result = predictor.normal_predict(img_data, result)
# print(result)