import time import cv2 import numpy as np from keras.layers import Input from PIL import Image from centernet import CenterNet centernet = CenterNet() # 调用摄像头 capture = cv2.VideoCapture(0) # capture=cv2.VideoCapture("1.mp4") fps = 0.0 while (True): t1 = time.time() # 读取某一帧 ref, frame = capture.read() # 格式转变,BGRtoRGB frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) # 转变成Image frame = Image.fromarray(np.uint8(frame)) # 进行检测 frame = np.array(centernet.detect_image(frame)) fps = (fps + (1. / (time.time() - t1))) / 2 print("fps= %.2f" % (fps)) frame = cv2.putText(frame, "fps= %.2f" % (fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
image_ids = open(os.path.join(VOCdevkit_path, "VOC2007/ImageSets/Main/test.txt")).read().strip().split() if not os.path.exists(map_out_path): os.makedirs(map_out_path) if not os.path.exists(os.path.join(map_out_path, 'ground-truth')): os.makedirs(os.path.join(map_out_path, 'ground-truth')) if not os.path.exists(os.path.join(map_out_path, 'detection-results')): os.makedirs(os.path.join(map_out_path, 'detection-results')) if not os.path.exists(os.path.join(map_out_path, 'images-optional')): os.makedirs(os.path.join(map_out_path, 'images-optional')) class_names, _ = get_classes(classes_path) if map_mode == 0 or map_mode == 1: print("Load model.") centernet = CenterNet(confidence = 0.01, nms_iou = 0.5) print("Load model done.") print("Get predict result.") for image_id in tqdm(image_ids): image_path = os.path.join(VOCdevkit_path, "VOC2007/JPEGImages/"+image_id+".jpg") image = Image.open(image_path) if map_vis: image.save(os.path.join(map_out_path, "images-optional/" + image_id + ".jpg")) centernet.get_map_txt(image_id, image, class_names, map_out_path) print("Get predict result done.") if map_mode == 0 or map_mode == 2: print("Get ground truth result.") for image_id in tqdm(image_ids): with open(os.path.join(map_out_path, "ground-truth/"+image_id+".txt"), "w") as new_f:
#-------------------------------------------------------------------------# # dir_origin_path 指定了用于检测的图片的文件夹路径 # dir_save_path 指定了检测完图片的保存路径 # # dir_origin_path和dir_save_path仅在mode='dir_predict'时有效 #-------------------------------------------------------------------------# dir_origin_path = "img/" dir_save_path = "img_out/" #-------------------------------------------------------------------------# # heatmap_save_path 热力图的保存路径,默认保存在model_data下 # # heatmap_save_path仅在mode='heatmap'有效 #-------------------------------------------------------------------------# heatmap_save_path = "model_data/heatmap_vision.png" centernet = CenterNet(heatmap=True if mode == "heatmap" else False) if mode == "predict": ''' 1、如果想要进行检测完的图片的保存,利用r_image.save("img.jpg")即可保存,直接在predict.py里进行修改即可。 2、如果想要获得预测框的坐标,可以进入centernet.detect_image函数,在绘图部分读取top,left,bottom,right这四个值。 3、如果想要利用预测框截取下目标,可以进入centernet.detect_image函数,在绘图部分利用获取到的top,left,bottom,right这四个值 在原图上利用矩阵的方式进行截取。 4、如果想要在预测图上写额外的字,比如检测到的特定目标的数量,可以进入centernet.detect_image函数,在绘图部分对predicted_class进行判断, 比如判断if predicted_class == 'car': 即可判断当前目标是否为车,然后记录数量即可。利用draw.text即可写字。 ''' while True: img = input('Input image filename:') try: image = Image.open(img) except:
''' predict.py有几个注意点 1、无法进行批量预测,如果想要批量预测,可以利用os.listdir()遍历文件夹,利用Image.open打开图片文件进行预测。 2、如果想要保存,利用r_image.save("img.jpg")即可保存。 3、如果想要获得框的坐标,可以进入detect_image函数,读取top,left,bottom,right这四个值。 4、如果想要截取下目标,可以利用获取到的top,left,bottom,right这四个值在原图上利用矩阵的方式进行截取。 ''' from centernet import CenterNet from PIL import Image centernet = CenterNet() while True: img = input('Input image filename:') try: image = Image.open(img) except: print('Open Error! Try again!') continue else: r_image = centernet.detect_image(image) # r_image.save("img.jpg") r_image.show()
#-----------------------------------------------------------------------# # predict.py将单张图片预测、摄像头检测、FPS测试和目录遍历检测等功能 # 整合到了一个py文件中,通过指定mode进行模式的修改。 #-----------------------------------------------------------------------# import time import cv2 import numpy as np from PIL import Image from centernet import CenterNet if __name__ == "__main__": centernet = CenterNet() #----------------------------------------------------------------------------------------------------------# # mode用于指定测试的模式: # 'predict' 表示单张图片预测,如果想对预测过程进行修改,如保存图片,截取对象等,可以先看下方详细的注释 # 'video' 表示视频检测,可调用摄像头或者视频进行检测,详情查看下方注释。 # 'fps' 表示测试fps,使用的图片是img里面的street.jpg,详情查看下方注释。 # 'dir_predict' 表示遍历文件夹进行检测并保存。默认遍历img文件夹,保存img_out文件夹,详情查看下方注释。 # 'heatmap' 表示进行预测结果的热力图可视化,详情查看下方注释。 # 'export_onnx' 表示将模型导出为onnx,需要pytorch1.7.1以上。 #----------------------------------------------------------------------------------------------------------# mode = "predict" #-------------------------------------------------------------------------# # crop 指定了是否在单张图片预测后对目标进行截取 # count 指定了是否进行目标的计数 # crop、count仅在mode='predict'时有效 #-------------------------------------------------------------------------# crop = False count = False
cls_heads, reg_heads, center_heads, batch_positions = net( torch.autograd.Variable(torch.randn(3, 3, image_h, image_w))) annotations = torch.FloatTensor([[[113, 120, 183, 255, 5], [13, 45, 175, 210, 2]], [[11, 18, 223, 225, 1], [-1, -1, -1, -1, -1]], [[-1, -1, -1, -1, -1], [-1, -1, -1, -1, -1]]]) decode = FCOSDecoder(image_w, image_h) batch_scores2, batch_classes2, batch_pred_bboxes2 = decode( cls_heads, reg_heads, center_heads, batch_positions) print("2222", batch_scores2.shape, batch_classes2.shape, batch_pred_bboxes2.shape) from centernet import CenterNet net = CenterNet(resnet_type="resnet50") image_h, image_w = 640, 640 heatmap_output, offset_output, wh_output = net( torch.autograd.Variable(torch.randn(3, 3, image_h, image_w))) annotations = torch.FloatTensor([[[113, 120, 183, 255, 5], [13, 45, 175, 210, 2]], [[11, 18, 223, 225, 1], [-1, -1, -1, -1, -1]], [[-1, -1, -1, -1, -1], [-1, -1, -1, -1, -1]]]) decode = CenterNetDecoder(image_w, image_h) batch_scores3, batch_classes3, batch_pred_bboxes3 = decode( heatmap_output, offset_output, wh_output) print("3333", batch_scores3.shape, batch_classes3.shape, batch_pred_bboxes3.shape)
''' predict.py有几个注意点 1、无法进行批量预测,如果想要批量预测,可以利用os.listdir()遍历文件夹,利用Image.open打开图片文件进行预测。 2、如果想要保存,利用r_image.save("img.jpg")即可保存。 3、如果想要获得框的坐标,可以进入detect_image函数,读取top,left,bottom,right这四个值。 4、如果想要截取下目标,可以利用获取到的top,left,bottom,right这四个值在原图上利用矩阵的方式进行截取。 ''' from keras.layers import Input from PIL import Image from centernet import CenterNet centernet = CenterNet() while True: img = input('Input image filename:') try: image = Image.open(img) except: print('Open Error! Try again!') continue else: r_image = centernet.detect_image(image) r_image.show() centernet.close_session()