text: # -*- coding: utf-8 -*-
  # Copyright 2021 Huawei Technologies Co., Ltd
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  # http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ============================================================================
  
  import argparse
  import torch
  from lib.config import config, update_config, hrnet
  
  
  def parse_args():
      parser = argparse.ArgumentParser(description='Train classification network')
      parser.add_argument('--cfg',
                          help='experiment configure file name',
                          required=False,
                          type=str,
                          default='./experiments/cls_hrnet_w18_sgd_lr5e-2_wd1e-4_bs32_x100.yaml')
  
      parser.add_argument('--modelDir',
                          help='model directory',
                          type=str,
                          default='')
      parser.add_argument('--logDir',
                          help='log directory',
                          type=str,
                          default='')
      parser.add_argument('--dataDir',
                          help='data directory',
                          type=str,
                          default='')
      parser.add_argument('--testModel',
                          help='testModel',
                          type=str,
                          default='')
  
      args = parser.parse_args()
      update_config(config, args)
  
      return args
  
  
  def build_model():
      args = parse_args()
      checkpoint = torch.load("./output/imagenet/cls_hrnet_w18_sgd_lr5e-2_wd1e-4_bs32_x100/model_best.pth.tar", map_location='cpu')
      model = hrnet.get_cls_net(config)
      model.load_state_dict(checkpoint)
      model.eval()
      return model
  
  
  def get_raw_data():
      from PIL import Image
      from urllib.request import urlretrieve
      IMAGE_URL = 'https://bbs-img.huaweicloud.com/blogs/img/thumb/1591951315139_8989_1363.png'
      urlretrieve(IMAGE_URL, 'tmp.jpg')
      img = Image.open("tmp.jpg")
      img = img.convert('RGB')
      return img
  
  
  def pre_process(raw_data):
      from torchvision import transforms
      normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                       std=[0.229, 0.224, 0.225])
      transforms_list = transforms.Compose([
          transforms.Resize(256),
          transforms.CenterCrop(224),
          transforms.ToTensor(),
          normalize
      ])
      input_data = transforms_list(raw_data)
      return input_data.unsqueeze(0)
  
  
  def post_process(output_tensor):
      return  torch.argmax(output_tensor, 1)
  
  
  if __name__ == '__main__':
      # 1. get raw data
      raw_data = get_raw_data()
  
      # 2. buid model
      model = build_model()
  
      # 3. pre process data
      input_tensor = pre_process(raw_data)
  
      # 4. run forward
      output_tensor = model(input_tensor)
  
      # 5. post process
      result = post_process(output_tensor)
  
      # 6. print result
      print(result)
  