def main(): im_height = 224 im_width = 224 num_classes = 2 cap = cv2.VideoCapture(0) # read class_indict json_path = './class_indices.json' assert os.path.exists(json_path), "file: '{}' dose not exist.".format( json_path) json_file = open(json_path, "r") class_indict = json.load(json_file) # create model feature = MobileNetV2(include_top=False) model = tf.keras.Sequential([ feature, tf.keras.layers.GlobalAvgPool2D(), tf.keras.layers.Dropout(rate=0.5), tf.keras.layers.Dense(num_classes), tf.keras.layers.Softmax() ]) weights_path = './save_weights/resMobileNetV2.ckpt' assert len(glob.glob(weights_path + "*")), "cannot find {}".format(weights_path) model.load_weights(weights_path) while True: rval, cap_img = cap.read() cap_img = cv2.flip(cap_img, 1, 1) img = Image.fromarray(cap_img) # resize image to 224x224 img = img.resize((im_width, im_height)) # scaling pixel value to (-1,1) img = np.array(img).astype(np.float32) img = ((img / 255.) - 0.5) * 2.0 # Add the image to a batch where it's the only member. img = (np.expand_dims(img, 0)) result = np.squeeze(model.predict(img)) predict_class = np.argmax(result) result_class_text = "Predict Result: " + class_indict[str( predict_class)] result_prob_text = "Accuracy: " + str(result[predict_class]) fps = cap.get(cv2.CAP_PROP_FPS) # 读取帧率 fps_text = "fps: " + str(fps) cv2.putText(cap_img, fps_text, (10, 30), cv2.FONT_HERSHEY_PLAIN, 1.0, (255, 255, 255), 2) cv2.putText(cap_img, result_class_text, (120, 30), cv2.FONT_HERSHEY_PLAIN, 1.0, (255, 255, 255), 2) cv2.putText(cap_img, result_prob_text, (360, 30), cv2.FONT_HERSHEY_PLAIN, 1.0, (255, 255, 255), 2) cv2.imshow('im', cap_img) if cv2.waitKey(100) & 0xFF == ord('q'): break cap.release() cv2.destroyAllWindows()
def main(): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") data_transform = transforms.Compose( # 预处理部分 [ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # load image img_path = "../tulip.jpg" img_path = ".\\pytorch_classification\\Test6_mobilenet\\tulip.jpg" assert os.path.exists(img_path), "file: '{}' dose not exist.".format( img_path) img = Image.open(img_path) plt.imshow(img) # [N, C, H, W] img = data_transform(img) # expand batch dimension img = torch.unsqueeze(img, dim=0) # * 添加一个batch维度 # read class_indict json_path = './class_indices.json' assert os.path.exists(json_path), "file: '{}' dose not exist.".format( json_path) json_file = open(json_path, "r") class_indict = json.load(json_file) # * create model 实例化模型 model = MobileNetV2(num_classes=5).to(device) # * load model weights 载入模型权重 # model_weight_path = "./MobileNetV2.pth" model_weight_path = ".\\pytorch_classification\\Test6_mobilenet\\mobilenet_v2.pth" model.load_state_dict(torch.load(model_weight_path, map_location=device)) model.eval() with torch.no_grad(): # * 禁止更新梯度 # predict class output = torch.squeeze(model( img.to(device))).cpu() # * squeeze压缩batch维度 predict = torch.softmax(output, dim=0) # * softmax将输出转化为概率分布 predict_cla = torch.argmax(predict).numpy() # 获得最大的预测值所对应的索引 print_res = "class: {} prob: {:.3}".format( class_indict[str(predict_cla)], predict[predict_cla].numpy()) plt.title(print_res) print(print_res) plt.show()
def main(): im_height = 224 im_width = 224 num_classes = 2 # load image img_path = "unMaskedFace_135.jpg" assert os.path.exists(img_path), "file: '{}' dose not exist.".format( img_path) img = Image.open(img_path) # resize image to 224x224 img = img.resize((im_width, im_height)) plt.imshow(img) # scaling pixel value to (-1,1) img = np.array(img).astype(np.float32) img = ((img / 255.) - 0.5) * 2.0 # Add the image to a batch where it's the only member. img = (np.expand_dims(img, 0)) # read class_indict json_path = './class_indices.json' assert os.path.exists(json_path), "file: '{}' dose not exist.".format( json_path) json_file = open(json_path, "r") class_indict = json.load(json_file) # create model feature = MobileNetV2(include_top=False) model = tf.keras.Sequential([ feature, tf.keras.layers.GlobalAvgPool2D(), tf.keras.layers.Dropout(rate=0.5), tf.keras.layers.Dense(num_classes), tf.keras.layers.Softmax() ]) weights_path = './save_weights/resMobileNetV2.ckpt' assert len(glob.glob(weights_path + "*")), "cannot find {}".format(weights_path) model.load_weights(weights_path) result = np.squeeze(model.predict(img)) predict_class = np.argmax(result) print_res = "class: {} prob: {:.3}".format( class_indict[str(predict_class)], result[predict_class]) plt.title(print_res) print(print_res) plt.show()
def main(): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") data_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # load image img_path = "../tulip.jpg" assert os.path.exists(img_path), "file: '{}' dose not exist.".format( img_path) img = Image.open(img_path) plt.imshow(img) # [N, C, H, W] img = data_transform(img) # expand batch dimension img = torch.unsqueeze(img, dim=0) # read class_indict json_path = './class_indices.json' assert os.path.exists(json_path), "file: '{}' dose not exist.".format( json_path) json_file = open(json_path, "r") class_indict = json.load(json_file) # create model model = MobileNetV2(num_classes=5).to(device) # load model weights model_weight_path = "./MobileNetV2.pth" model.load_state_dict(torch.load(model_weight_path, map_location=device)) model.eval() with torch.no_grad(): # predict class output = torch.squeeze(model(img.to(device))).cpu() predict = torch.softmax(output, dim=0) predict_cla = torch.argmax(predict).numpy() print_res = "class: {} prob: {:.3}".format( class_indict[str(predict_cla)], predict[predict_cla].numpy()) plt.title(print_res) for i in range(len(predict)): print("class: {:10} prob: {:.3}".format(class_indict[str(i)], predict[i].numpy())) plt.show()
def main(): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") data_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # load image img_path = "../tulip.jpg" assert os.path.exists(img_path), "file: '{}' dose not exist.".format( img_path) img = Image.open(img_path) #载入图像 plt.imshow(img) # [N, C, H, W] img = data_transform(img) ##使用和预处理过程中相同的图像预处理方式 # expand batch dimension img = torch.unsqueeze(img, dim=0) #添加一个batch维度 # read class_indict json_path = './class_indices.json' assert os.path.exists(json_path), "file: '{}' dose not exist.".format( json_path) json_file = open(json_path, "r") class_indict = json.load(json_file) # create model model = MobileNetV2(num_classes=5).to(device) # load model weights model_weight_path = "./MobileNetV2.pth" model.load_state_dict(torch.load(model_weight_path, map_location=device)) #载入迁移学习之后训练好的模型权重 model.eval() #进入eval模式 with torch.no_grad(): #通过这句话,禁止运算过程中跟踪我们的误差梯度信息,从而节省内存空间 # predict class output = torch.squeeze(model(img.to(device))).cpu() #将图像传入模型当中 predict = torch.softmax(output, dim=0) #将输出转换成概率分布 predict_cla = torch.argmax(predict).numpy() #获得最大的预测概率值所对应的索引 print_res = "class: {} prob: {:.3}".format( class_indict[str(predict_cla)], predict[predict_cla].numpy()) plt.title(print_res) print(print_res) plt.show()
def main(): data_root = os.path.abspath(os.path.join(os.getcwd(), "../..")) # get data root path image_path = os.path.join(data_root, "data_set", "flower_data") # flower data set path train_dir = os.path.join(image_path, "train") validation_dir = os.path.join(image_path, "val") assert os.path.exists(train_dir), "cannot find {}".format(train_dir) assert os.path.exists(validation_dir), "cannot find {}".format(validation_dir) im_height = 224 im_width = 224 batch_size = 16 epochs = 20 num_classes = 5 def pre_function(img): # img = im.open('test.jpg') # img = np.array(img).astype(np.float32) img = img / 255. img = (img - 0.5) * 2.0 return img # data generator with data augmentation train_image_generator = ImageDataGenerator(horizontal_flip=True, preprocessing_function=pre_function) validation_image_generator = ImageDataGenerator(preprocessing_function=pre_function) train_data_gen = train_image_generator.flow_from_directory(directory=train_dir, batch_size=batch_size, shuffle=True, target_size=(im_height, im_width), class_mode='categorical') total_train = train_data_gen.n # get class dict class_indices = train_data_gen.class_indices # transform value and key of dict inverse_dict = dict((val, key) for key, val in class_indices.items()) # write dict into json file json_str = json.dumps(inverse_dict, indent=4) with open('class_indices.json', 'w') as json_file: json_file.write(json_str) val_data_gen = validation_image_generator.flow_from_directory(directory=validation_dir, batch_size=batch_size, shuffle=False, target_size=(im_height, im_width), class_mode='categorical') # img, _ = next(train_data_gen) total_val = val_data_gen.n print("using {} images for training, {} images for validation.".format(total_train, total_val)) # create model except fc layer feature = MobileNetV2(include_top=False) # download weights 链接: https://pan.baidu.com/s/1YgFoIKHqooMrTQg_IqI2hA 密码: 2qht pre_weights_path = './pretrain_weights.ckpt' assert len(glob.glob(pre_weights_path+"*")), "cannot find {}".format(pre_weights_path) feature.load_weights(pre_weights_path) feature.trainable = False feature.summary() # add last fc layer model = tf.keras.Sequential([feature, tf.keras.layers.GlobalAvgPool2D(), tf.keras.layers.Dropout(rate=0.5), tf.keras.layers.Dense(num_classes), tf.keras.layers.Softmax()]) model.summary() # using keras low level api for training loss_object = tf.keras.losses.CategoricalCrossentropy(from_logits=False) optimizer = tf.keras.optimizers.Adam(learning_rate=0.001) train_loss = tf.keras.metrics.Mean(name='train_loss') train_accuracy = tf.keras.metrics.CategoricalAccuracy(name='train_accuracy') val_loss = tf.keras.metrics.Mean(name='val_loss') val_accuracy = tf.keras.metrics.CategoricalAccuracy(name='val_accuracy') @tf.function def train_step(images, labels): with tf.GradientTape() as tape: output = model(images, training=True) loss = loss_object(labels, output) gradients = tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) train_loss(loss) train_accuracy(labels, output) @tf.function def val_step(images, labels): output = model(images, training=False) loss = loss_object(labels, output) val_loss(loss) val_accuracy(labels, output) best_val_acc = 0. for epoch in range(epochs): train_loss.reset_states() # clear history info train_accuracy.reset_states() # clear history info val_loss.reset_states() # clear history info val_accuracy.reset_states() # clear history info # train train_bar = tqdm(range(total_train // batch_size), file=sys.stdout) for step in train_bar: images, labels = next(train_data_gen) train_step(images, labels) # print train process train_bar.desc = "train epoch[{}/{}] loss:{:.3f}, acc:{:.3f}".format(epoch + 1, epochs, train_loss.result(), train_accuracy.result()) # validate val_bar = tqdm(range(total_val // batch_size), file=sys.stdout) for step in val_bar: val_images, val_labels = next(val_data_gen) val_step(val_images, val_labels) # print val process val_bar.desc = "valid epoch[{}/{}] loss:{:.3f}, acc:{:.3f}".format(epoch + 1, epochs, val_loss.result(), val_accuracy.result()) # only save best weights if val_accuracy.result() > best_val_acc: best_val_acc = val_accuracy.result() model.save_weights("./save_weights/resMobileNetV2.ckpt", save_format="tf")
def main(): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print("using {} device.".format(device)) batch_size = 16 epochs = 5 data_transform = { "train": transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), "val": transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) } data_root = os.path.abspath(os.path.join(os.getcwd(), "../..")) # get data root path image_path = os.path.join(data_root, "data_set", "flower_data") # flower data set path assert os.path.exists(image_path), "{} path does not exist.".format( image_path) train_dataset = datasets.ImageFolder(root=os.path.join( image_path, "train"), transform=data_transform["train"]) train_num = len(train_dataset) # {'daisy':0, 'dandelion':1, 'roses':2, 'sunflower':3, 'tulips':4} flower_list = train_dataset.class_to_idx cla_dict = dict((val, key) for key, val in flower_list.items()) # write dict into json file json_str = json.dumps(cla_dict, indent=4) with open('class_indices.json', 'w') as json_file: json_file.write(json_str) nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers print('Using {} dataloader workers every process'.format(nw)) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=nw) validate_dataset = datasets.ImageFolder(root=os.path.join( image_path, "val"), transform=data_transform["val"]) val_num = len(validate_dataset) validate_loader = torch.utils.data.DataLoader(validate_dataset, batch_size=batch_size, shuffle=False, num_workers=nw) print("using {} images for training, {} images for validation.".format( train_num, val_num)) # create model net = MobileNetV2(num_classes=5) # load pretrain weights # download url: https://download.pytorch.org/models/mobilenet_v2-b0353104.pth model_weight_path = "./mobilenet_v2.pth" assert os.path.exists(model_weight_path), "file {} dose not exist.".format( model_weight_path) pre_weights = torch.load(model_weight_path, map_location='cpu') # delete classifier weights pre_dict = { k: v for k, v in pre_weights.items() if net.state_dict()[k].numel() == v.numel() } missing_keys, unexpected_keys = net.load_state_dict(pre_dict, strict=False) # freeze features weights for param in net.features.parameters(): param.requires_grad = False net.to(device) # define loss function loss_function = nn.CrossEntropyLoss() # construct an optimizer params = [p for p in net.parameters() if p.requires_grad] optimizer = optim.Adam(params, lr=0.0001) best_acc = 0.0 save_path = './MobileNetV2.pth' train_steps = len(train_loader) for epoch in range(epochs): # train net.train() running_loss = 0.0 train_bar = tqdm(train_loader, file=sys.stdout) for step, data in enumerate(train_bar): images, labels = data optimizer.zero_grad() logits = net(images.to(device)) loss = loss_function(logits, labels.to(device)) loss.backward() optimizer.step() # print statistics running_loss += loss.item() train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format( epoch + 1, epochs, loss) # validate net.eval() acc = 0.0 # accumulate accurate number / epoch with torch.no_grad(): val_bar = tqdm(validate_loader, file=sys.stdout) for val_data in val_bar: val_images, val_labels = val_data outputs = net(val_images.to(device)) # loss = loss_function(outputs, test_labels) predict_y = torch.max(outputs, dim=1)[1] acc += torch.eq(predict_y, val_labels.to(device)).sum().item() val_bar.desc = "valid epoch[{}/{}]".format(epoch + 1, epochs) val_accurate = acc / val_num print('[epoch %d] train_loss: %.3f val_accuracy: %.3f' % (epoch + 1, running_loss / train_steps, val_accurate)) if val_accurate > best_acc: best_acc = val_accurate torch.save(net.state_dict(), save_path) print('Finished Training')
def main(): gpus = tf.config.experimental.list_physical_devices("GPU") if gpus: try: for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) except RuntimeError as e: print(e) exit(-1) data_root = os.path.abspath(os.path.join(os.getcwd(), "../..")) # get data root path image_path = os.path.join(data_root, "data_set", "flower_data") # flower data set path train_dir = os.path.join(image_path, "train") validation_dir = os.path.join(image_path, "val") assert os.path.exists(train_dir), "cannot find {}".format(train_dir) assert os.path.exists(validation_dir), "cannot find {}".format( validation_dir) # create direction for saving weights if not os.path.exists("save_weights"): os.makedirs("save_weights") im_height = 224 im_width = 224 batch_size = 32 epochs = 30 # class dict data_class = [ cla for cla in os.listdir(train_dir) if os.path.isdir(os.path.join(train_dir, cla)) ] class_num = len(data_class) class_dict = dict((value, index) for index, value in enumerate(data_class)) # reverse value and key of dict inverse_dict = dict((val, key) for key, val in class_dict.items()) # write dict into json file json_str = json.dumps(inverse_dict, indent=4) with open('class_indices.json', 'w') as json_file: json_file.write(json_str) # load train images list train_image_list = glob.glob(train_dir + "/*/*.jpg") random.shuffle(train_image_list) train_num = len(train_image_list) assert train_num > 0, "cannot find any .jpg file in {}".format(train_dir) train_label_list = [ class_dict[path.split(os.path.sep)[-2]] for path in train_image_list ] # load validation images list val_image_list = glob.glob(validation_dir + "/*/*.jpg") random.shuffle(val_image_list) val_num = len(val_image_list) assert val_num > 0, "cannot find any .jpg file in {}".format( validation_dir) val_label_list = [ class_dict[path.split(os.path.sep)[-2]] for path in val_image_list ] print("using {} images for training, {} images for validation.".format( train_num, val_num)) def process_train_img(img_path, label): label = tf.one_hot(label, depth=class_num) image = tf.io.read_file(img_path) image = tf.image.decode_jpeg(image) image = tf.image.convert_image_dtype(image, tf.float32) image = tf.image.resize(image, [im_height, im_width]) image = tf.image.random_flip_left_right(image) # image = (image - 0.5) / 0.5 image = (image - 0.5) * 2.0 return image, label def process_val_img(img_path, label): label = tf.one_hot(label, depth=class_num) image = tf.io.read_file(img_path) image = tf.image.decode_jpeg(image) image = tf.image.convert_image_dtype(image, tf.float32) image = tf.image.resize(image, [im_height, im_width]) # image = (image - 0.5) / 0.5 image = (image - 0.5) * 2.0 return image, label AUTOTUNE = tf.data.experimental.AUTOTUNE # load train dataset train_dataset = tf.data.Dataset.from_tensor_slices( (train_image_list, train_label_list)) train_dataset = train_dataset.shuffle(buffer_size=train_num)\ .map(process_train_img, num_parallel_calls=AUTOTUNE)\ .repeat().batch(batch_size).prefetch(AUTOTUNE) # load train dataset val_dataset = tf.data.Dataset.from_tensor_slices( (val_image_list, val_label_list)) val_dataset = val_dataset.map(process_val_img, num_parallel_calls=tf.data.experimental.AUTOTUNE)\ .repeat().batch(batch_size) # 实例化模型 model = MobileNetV2(num_classes=5) pre_weights_path = './pretrain_weights.ckpt' assert len(glob.glob(pre_weights_path + "*")), "cannot find {}".format(pre_weights_path) model.load_weights(pre_weights_path) for layer_t in model.layers[:-1]: layer_t.trainable = False model.summary() # using keras low level api for training loss_object = tf.keras.losses.CategoricalCrossentropy(from_logits=True) optimizer = tf.keras.optimizers.Adam(learning_rate=0.0005) train_loss = tf.keras.metrics.Mean(name='train_loss') train_accuracy = tf.keras.metrics.CategoricalAccuracy( name='train_accuracy') test_loss = tf.keras.metrics.Mean(name='test_loss') test_accuracy = tf.keras.metrics.CategoricalAccuracy(name='test_accuracy') @tf.function def train_step(images, labels): with tf.GradientTape() as tape: output = model(images, training=True) loss = loss_object(labels, output) gradients = tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) train_loss(loss) train_accuracy(labels, output) @tf.function def test_step(images, labels): output = model(images, training=False) t_loss = loss_object(labels, output) test_loss(t_loss) test_accuracy(labels, output) best_test_loss = float('inf') train_step_num = train_num // batch_size val_step_num = val_num // batch_size for epoch in range(1, epochs + 1): train_loss.reset_states() # clear history info train_accuracy.reset_states() # clear history info test_loss.reset_states() # clear history info test_accuracy.reset_states() # clear history info t1 = time.perf_counter() for index, (images, labels) in enumerate(train_dataset): train_step(images, labels) if index + 1 == train_step_num: break print(time.perf_counter() - t1) for index, (images, labels) in enumerate(val_dataset): test_step(images, labels) if index + 1 == val_step_num: break template = 'Epoch {}, Loss: {}, Accuracy: {}, Test Loss: {}, Test Accuracy: {}' print( template.format(epoch, train_loss.result(), train_accuracy.result() * 100, test_loss.result(), test_accuracy.result() * 100)) if test_loss.result() < best_test_loss: model.save_weights("./save_weights/myMobileNet.ckpt".format(epoch), save_format='tf')
def main(): im_height = 224 im_width = 224 num_classes = 2 cap = cv2.VideoCapture(0) # read class_indict json_path = './class_indices.json' assert os.path.exists(json_path), "file: '{}' dose not exist.".format( json_path) json_file = open(json_path, "r") class_indict = json.load(json_file) # create model feature = MobileNetV2(include_top=False) model = tf.keras.Sequential([ feature, tf.keras.layers.GlobalAvgPool2D(), tf.keras.layers.Dropout(rate=0.5), tf.keras.layers.Dense(num_classes), tf.keras.layers.Softmax() ]) weights_path = './save_weights/resMobileNetV2.ckpt' assert len(glob.glob(weights_path + "*")), "cannot find {}".format(weights_path) model.load_weights(weights_path) classifier = cv2.CascadeClassifier('haarcascade_frontalface_default.xml') while True: rval, im = cap.read() im = cv2.flip(im, 1, 1) # Flip to act as a mirror # Resize the image to speed up detection mini = cv2.resize(im, (im.shape[1] // 4, im.shape[0] // 4)) # detect MultiScale / faces faces = classifier.detectMultiScale(mini) # Draw rectangles around each face for f in faces: (x, y, w, h) = [v * 4 for v in f] # Scale the shapesize backup # Save just the rectangle faces in SubRecFaces if y >= 300 and x >= 300: face_img = im[y - 300:y + h + 300, x - 300:x + w + 300] else: face_img = im[y:y + h, x:x + w] img = Image.fromarray(face_img) img = img.resize((im_width, im_height)) # scaling pixel value to (-1,1) img = np.array(img).astype(np.float32) img = ((img / 255.) - 0.5) * 2.0 # Add the image to a batch where it's the only member. img = (np.expand_dims(img, 0)) result = np.squeeze(model.predict(img)) predict_class = np.argmax(result) result_class_text = "Predict Result: " + class_indict[str( predict_class)] result_prob_text = "Accuracy: " + str(result[predict_class]) fps = cap.get(cv2.CAP_PROP_FPS) # 读取帧率 fps_text = "fps: " + str(fps) cv2.putText(im, fps_text, (10, 30), cv2.FONT_HERSHEY_PLAIN, 1.0, (255, 255, 255), 2) cv2.putText(im, result_class_text, (120, 30), cv2.FONT_HERSHEY_PLAIN, 1.0, (255, 255, 255), 2) cv2.putText(im, result_prob_text, (360, 30), cv2.FONT_HERSHEY_PLAIN, 1.0, (255, 255, 255), 2) cv2.rectangle(im, (x, y), (x + w, y + h), (0, 255, 0), 2) cv2.imshow('im', im) if cv2.waitKey(100) & 0xFF == ord('q'): break cap.release() cv2.destroyAllWindows()
def main(): device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') data_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # load image img_path = 'a.jpg' assert os.path.exists(img_path), "file: '{}' does not exit.".format( img_path) img = Image.open(img_path) plt.imshow(img) # [N, C, H, W] img = data_transform(img) # expand batch dimension img = torch.unsqueeze(img, dim=0) print(img.shape) # read class_indict json_path = 'class_indices.json' assert os.path.exists(json_path), "file: '{}' does not exist.".format( json_path) json_file = open(json_path, 'r') class_indict = json.load(json_file) # print(class_indict) # create model model = MobileNetV2(num_classes=5).to(device) # load model weights model_weight_path = 'weights/MobileNetV2.pth' model.load_state_dict(torch.load(model_weight_path, map_location=device)) model.eval() print('=================================') dummy_input = torch.randn(1, 3, 224, 224).to(device) torch.onnx.export(model, dummy_input, 'mobilenetv2_op9.onnx', dynamic_axes={ 'image': { 0: 'B' }, 'outputs': { 0: 'B' } }, input_names=['image'], output_names=['outputs'], opset_version=9) print('=================================') with torch.no_grad(): # predict class output = torch.squeeze(model(img.to(device))).cpu() predict = torch.softmax(output, dim=0) predict_cla = torch.argmax(predict).numpy() print_res = "class: {} prob: {:.3}".format(class_indict[str(predict_cla)], predict[predict_cla].numpy()) plt.title(print_res) print(print_res) plt.show()