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()
