python pyplot用法_matplotlib中plt用法实例

import torch

from models.models import Model

import cv2

from PIL import Image

import numpy as np

from matplotlib.animation import FFMpegWriter

import time

import matplotlib.pyplot as plt

from torchvision.transforms import functional

exp_name = './xxxx_results'

dataRoot = 'xxxx.mp4'

model_path = './checkpoint_best.pth'

def pre_image(image):

image = Image.fromarray(cv2.cvtColor(image,cv2.COLOR_BGR2RGB))

input_image = image.copy()

# image.show()

height, width = image.size[1], image.size[0]

height = round(height / 16) * 16

width = round(width / 16) * 16

image = image.resize((width, height), Image.BILINEAR)

image = functional.to_tensor(image)

image = functional.normalize(image, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225])

return input_image,torch.unsqueeze(image,0)

if __name__ == '__main__':

device = torch.device('cuda:0')

# load model

model=Model()

checkpoint = torch.load(model_path)

model.load_state_dict(checkpoint['model'])

model.cuda()

model.eval()

# input video

video = cv2.VideoCapture(dataRoot)

fps = video.get(cv2.CAP_PROP_FPS)

print(fps)

frameCount = video.get(cv2.CAP_PROP_FRAME_COUNT)

print(frameCount)

size = (int(video.get(cv2.CAP_PROP_FRAME_WIDTH)), int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)))

# metadata = dict(title='Video Test', artist='Matplotlib', comment='Movie support!')

# writer = FFMpegWriter(fps=25, metadata=metadata)

# videoWriter = cv2.VideoWriter('trans.mp4', cv2.VideoWriter_fourcc(*'MP4V'), fps, size)

success, frame = video.read()

index = 1

figure = plt.figure()

while success:

# time1=time.time()

src_image,frame = pre_image(frame)

images = frame.to(device)

# time1 = time.time()

# ground truth

# gt_path = dataRoot + '/den/' + filename_no_ext + '.csv'

# predict

dense_map,atten_map = model(images)

# test = time.time() - time1

dense_map = dense_map.cpu().data.numpy()[0,0,:,:]

# test=time.time()-time1

dense_pred_count = np.sum(dense_map)

dense_map = dense_map/np.max(dense_map+1e-20)

# cv2.imshow("image", dense_map)

# cv2.waitKey(0)

plt.subplot(121)

plt.imshow(src_image)

# plt.title('original image')

plt.axis('off')

plt.subplot(122)

plt.imshow(dense_map)

# plt.title('dense map')

plt.text(25, 25, 'pred crowd count:%.4f ' % dense_pred_count, fontdict={'size': 10, 'color': 'red'})

plt.axis('off')

plt.tight_layout(pad=0.3, w_pad=0, h_pad=1)

# anni=animation.FuncAnimation(fig, animate, init_func=init,frames=200, interval=20, blit=True)

# anim.save('sin.gif', fps=75, writer='imagemagick')

plt.savefig(exp_name + '/'+ str('%05d' % index) + '_' + str(int(dense_pred_count)) + '.png', bbox_inches='tight', pad_inches=0, dpi=150)

# plt.show()

plt.clf()

success, frame = video.read()

index += 1

video.release()