# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved. # # 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. # Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved. # # 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 os import sys __dir__ = os.path.dirname(__file__) sys.path.append(os.path.join(__dir__, '')) import argparse import shutil import cv2 import numpy as np import tarfile import requests from tqdm import tqdm from tools.infer.utils import get_image_list, preprocess, save_prelabel_results from tools.infer.predict import Predictor __all__ = ['PaddleClas'] BASE_DIR = os.path.expanduser("~/.paddleclas/") BASE_INFERENCE_MODEL_DIR = os.path.join(BASE_DIR, 'inference_model') BASE_IMAGES_DIR = os.path.join(BASE_DIR, 'images') model_names = { 'Xception71', 'SE_ResNeXt101_32x4d', 'ShuffleNetV2_x0_5', 'ResNet34', 'ShuffleNetV2_x2_0', 'ResNeXt101_32x4d', 'HRNet_W48_C_ssld', 'ResNeSt50_fast_1s1x64d', 'MobileNetV2_x2_0', 'MobileNetV3_large_x1_0', 'Fix_ResNeXt101_32x48d_wsl', 'MobileNetV2_ssld', 'ResNeXt101_vd_64x4d', 'ResNet34_vd_ssld', 'MobileNetV3_small_x1_0', 'VGG11', 'ResNeXt50_vd_32x4d', 'MobileNetV3_large_x1_25', 'MobileNetV3_large_x1_0_ssld', 'MobileNetV2_x0_75', 'MobileNetV3_small_x0_35', 'MobileNetV1_x0_75', 'MobileNetV1_ssld', 'ResNeXt50_32x4d', 'GhostNet_x1_3_ssld', 'Res2Net101_vd_26w_4s', 'ResNet152', 'Xception65', 'EfficientNetB0', 'ResNet152_vd', 'HRNet_W18_C', 'Res2Net50_14w_8s', 'ShuffleNetV2_x0_25', 'HRNet_W64_C', 'Res2Net50_vd_26w_4s_ssld', 'HRNet_W18_C_ssld', 'ResNet18_vd', 'ResNeXt101_32x16d_wsl', 'SE_ResNeXt50_32x4d', 'SqueezeNet1_1', 'SENet154_vd', 'SqueezeNet1_0', 'GhostNet_x1_0', 'ResNet50_vc', 'DPN98', 'HRNet_W48_C', 'DenseNet264', 'SE_ResNet34_vd', 'HRNet_W44_C', 'MobileNetV3_small_x1_25', 'MobileNetV1_x0_5', 'ResNet200_vd', 'VGG13', 'EfficientNetB3', 'EfficientNetB2', 'ShuffleNetV2_x0_33', 'MobileNetV3_small_x0_75', 'ResNeXt152_vd_32x4d', 'ResNeXt101_32x32d_wsl', 'ResNet18', 'MobileNetV3_large_x0_35', 'Res2Net50_26w_4s', 'MobileNetV2_x0_5', 'EfficientNetB0_small', 'ResNet101_vd_ssld', 'EfficientNetB6', 'EfficientNetB1', 'EfficientNetB7', 'ResNeSt50', 'ShuffleNetV2_x1_0', 'MobileNetV3_small_x1_0_ssld', 'InceptionV4', 'GhostNet_x0_5', 'SE_HRNet_W64_C_ssld', 'ResNet50_ACNet_deploy', 'Xception41', 'ResNet50', 'Res2Net200_vd_26w_4s_ssld', 'Xception41_deeplab', 'SE_ResNet18_vd', 'SE_ResNeXt50_vd_32x4d', 'HRNet_W30_C', 'HRNet_W40_C', 'VGG19', 'Res2Net200_vd_26w_4s', 'ResNeXt101_32x8d_wsl', 'ResNet50_vd', 'ResNeXt152_64x4d', 'DarkNet53', 'ResNet50_vd_ssld', 'ResNeXt101_64x4d', 'MobileNetV1_x0_25', 'Xception65_deeplab', 'AlexNet', 'ResNet101', 'DenseNet121', 'ResNet50_vd_v2', 'Res2Net50_vd_26w_4s', 'ResNeXt101_32x48d_wsl', 'MobileNetV3_large_x0_5', 'MobileNetV2_x0_25', 'DPN92', 'ResNet101_vd', 'MobileNetV2_x1_5', 'DPN131', 'ResNeXt50_vd_64x4d', 'ShuffleNetV2_x1_5', 'ResNet34_vd', 'MobileNetV1', 'ResNeXt152_vd_64x4d', 'DPN107', 'VGG16', 'ResNeXt50_64x4d', 'RegNetX_4GF', 'DenseNet161', 'GhostNet_x1_3', 'HRNet_W32_C', 'Fix_ResNet50_vd_ssld_v2', 'Res2Net101_vd_26w_4s_ssld', 'DenseNet201', 'DPN68', 'EfficientNetB4', 'ResNeXt152_32x4d', 'InceptionV3', 'ShuffleNetV2_swish', 'GoogLeNet', 'ResNet50_vd_ssld_v2', 'SE_ResNet50_vd', 'MobileNetV2', 'ResNeXt101_vd_32x4d', 'MobileNetV3_large_x0_75', 'MobileNetV3_small_x0_5', 'DenseNet169', 'EfficientNetB5', 'DeiT_base_distilled_patch16_224', 'DeiT_base_distilled_patch16_384', 'DeiT_base_patch16_224', 'DeiT_base_patch16_384', 'DeiT_small_distilled_patch16_224', 'DeiT_small_patch16_224', 'DeiT_tiny_distilled_patch16_224', 'DeiT_tiny_patch16_224', 'ViT_base_patch16_224', 'ViT_base_patch16_384', 'ViT_base_patch32_384', 'ViT_large_patch16_224', 'ViT_large_patch16_384', 'ViT_large_patch32_384', 'ViT_small_patch16_224' } def download_with_progressbar(url, save_path): response = requests.get(url, stream=True) total_size_in_bytes = int(response.headers.get('content-length', 0)) block_size = 1024 # 1 Kibibyte progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True) with open(save_path, 'wb') as file: for data in response.iter_content(block_size): progress_bar.update(len(data)) file.write(data) progress_bar.close() if total_size_in_bytes == 0 or progress_bar.n != total_size_in_bytes: raise Exception( "Something went wrong while downloading model/image from {}". format(url)) def maybe_download(model_storage_directory, url): # using custom model tar_file_name_list = [ 'inference.pdiparams', 'inference.pdiparams.info', 'inference.pdmodel' ] if not os.path.exists( os.path.join(model_storage_directory, 'inference.pdiparams') ) or not os.path.exists( os.path.join(model_storage_directory, 'inference.pdmodel')): tmp_path = os.path.join(model_storage_directory, url.split('/')[-1]) print('download {} to {}'.format(url, tmp_path)) os.makedirs(model_storage_directory, exist_ok=True) download_with_progressbar(url, tmp_path) with tarfile.open(tmp_path, 'r') as tarObj: for member in tarObj.getmembers(): filename = None for tar_file_name in tar_file_name_list: if tar_file_name in member.name: filename = tar_file_name if filename is None: continue file = tarObj.extractfile(member) with open( os.path.join(model_storage_directory, filename), 'wb') as f: f.write(file.read()) os.remove(tmp_path) def load_label_name_dict(path): if not os.path.exists(path): print( "Warning: If want to use your own label_dict, please input legal path!\nOtherwise label_names will be empty!" ) return None else: result = {} for line in open(path, 'r'): partition = line.split('\n')[0].partition(' ') try: result[int(partition[0])] = str(partition[-1]) except: result = {} break return result def parse_args(mMain=True, add_help=True): def str2bool(v): return v.lower() in ("true", "t", "1") if mMain == True: # general params parser = argparse.ArgumentParser(add_help=add_help) parser.add_argument("--model_name", type=str) parser.add_argument("-i", "--image_file", type=str) parser.add_argument("--use_gpu", type=str2bool, default=False) # params for preprocess parser.add_argument("--resize_short", type=int, default=256) parser.add_argument("--resize", type=int, default=224) parser.add_argument("--normalize", type=str2bool, default=True) parser.add_argument("-b", "--batch_size", type=int, default=1) # params for predict parser.add_argument( "--model_file", type=str, default='') ## inference.pdmodel parser.add_argument( "--params_file", type=str, default='') ## inference.pdiparams parser.add_argument("--ir_optim", type=str2bool, default=True) parser.add_argument("--use_fp16", type=str2bool, default=False) parser.add_argument("--use_tensorrt", type=str2bool, default=False) parser.add_argument("--gpu_mem", type=int, default=8000) parser.add_argument("--enable_profile", type=str2bool, default=False) parser.add_argument("--top_k", type=int, default=1) parser.add_argument("--enable_mkldnn", type=str2bool, default=False) parser.add_argument("--cpu_num_threads", type=int, default=10) # parameters for pre-label the images parser.add_argument("--label_name_path", type=str, default='') parser.add_argument( "--pre_label_image", type=str2bool, default=False, help="Whether to pre-label the images using the loaded weights") parser.add_argument("--pre_label_out_idr", type=str, default=None) return parser.parse_args() else: return argparse.Namespace( model_name='', image_file='', use_gpu=False, use_fp16=False, use_tensorrt=False, is_preprocessed=False, resize_short=256, resize=224, normalize=True, batch_size=1, model_file='', params_file='', ir_optim=True, gpu_mem=8000, enable_profile=False, top_k=1, enable_mkldnn=False, cpu_num_threads=10, label_name_path='', pre_label_image=False, pre_label_out_idr=None) class PaddleClas(object): print('Inference models that Paddle provides are listed as follows:\n\n{}'. format(model_names), '\n') def __init__(self, **kwargs): process_params = parse_args(mMain=False, add_help=False) process_params.__dict__.update(**kwargs) if not os.path.exists(process_params.model_file): if process_params.model_name is None: raise Exception( 'Please input model name that you want to use!') if process_params.model_name in model_names: url = 'https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/inference/{}_infer.tar'.format( process_params.model_name) if not os.path.exists( os.path.join(BASE_INFERENCE_MODEL_DIR, process_params.model_name)): os.makedirs( os.path.join(BASE_INFERENCE_MODEL_DIR, process_params.model_name)) download_path = os.path.join(BASE_INFERENCE_MODEL_DIR, process_params.model_name) maybe_download(model_storage_directory=download_path, url=url) process_params.model_file = os.path.join(download_path, 'inference.pdmodel') process_params.params_file = os.path.join( download_path, 'inference.pdiparams') process_params.label_name_path = os.path.join( __dir__, 'ppcls/utils/imagenet1k_label_list.txt') else: raise Exception( 'The model inputed is {}, not provided by PaddleClas. If you want to use your own model, please input model_file as model path!'. format(process_params.model_name)) else: print('Using user-specified model and params!') print("process params are as follows: \n{}".format(process_params)) self.label_name_dict = load_label_name_dict( process_params.label_name_path) self.args = process_params self.predictor = Predictor(process_params) def postprocess(self, output): output = output.flatten() classes = np.argpartition(output, -self.args.top_k)[-self.args.top_k:] class_ids = classes[np.argsort(-output[classes])] scores = output[class_ids] label_names = [self.label_name_dict[c] for c in class_ids] if self.label_name_dict else [] return { "class_ids": class_ids, "scores": scores, "label_names": label_names } def predict(self, input_data): """ predict label of img with paddleclas Args: input_data(string, NumPy.ndarray): image to be classified, support: string: local path of image file, internet URL, directory containing series of images; NumPy.ndarray: preprocessed image data that has 3 channels and accords with [C, H, W], or raw image data that has 3 channels and accords with [H, W, C] Returns: dict: {image_name: "", class_id: [], scores: [], label_names: []}，if label name path == None，label_names will be empty. """ if isinstance(input_data, np.ndarray): if not self.args.is_preprocessed: input_data = input_data[:, :, ::-1] input_data = preprocess(input_data, self.args) input_data = np.expand_dims(input_data, axis=0) batch_outputs = self.predictor.predict(input_data) result = {"filename": "image"} result.update(self.postprocess(batch_outputs[0])) return result elif isinstance(input_data, str): input_path = input_data # download internet image if input_path.startswith('http'): if not os.path.exists(BASE_IMAGES_DIR): os.makedirs(BASE_IMAGES_DIR) file_path = os.path.join(BASE_IMAGES_DIR, 'tmp.jpg') download_with_progressbar(input_path, file_path) print("Current using image from Internet:{}, renamed as: {}". format(input_path, file_path)) input_path = file_path image_list = get_image_list(input_path) total_result = [] batch_input_list = [] img_path_list = [] cnt = 0 for idx, img_path in enumerate(image_list): img = cv2.imread(img_path) if img is None: print( "Warning: Image file failed to read and has been skipped. The path: {}". format(img_path)) continue else: img = img[:, :, ::-1] data = preprocess(img, self.args) batch_input_list.append(data) img_path_list.append(img_path) cnt += 1 if cnt % self.args.batch_size == 0 or (idx + 1 ) == len(image_list): batch_outputs = self.predictor.predict( np.array(batch_input_list)) for number, output in enumerate(batch_outputs): result = {"filename": img_path_list[number]} result.update(self.postprocess(output)) result_str = "top-{} result: {}".format( self.args.top_k, result) print(result_str) total_result.append(result) if self.args.pre_label_image: save_prelabel_results(result["class_ids"][0], img_path_list[number], self.args.pre_label_out_idr) batch_input_list = [] img_path_list = [] return total_result else: print( "Error: Please input legal image! The type of image supported by PaddleClas are: NumPy.ndarray and string of local path or Ineternet URL" ) return [] def main(): # for cmd args = parse_args(mMain=True) clas_engine = PaddleClas(**(args.__dict__)) print('{}{}{}'.format('*' * 10, args.image_file, '*' * 10)) total_result = clas_engine.predict(args.image_file) print("Predict complete!") if __name__ == '__main__': main()