diff --git a/FPS_test.py b/FPS_test.py index 6fbcf3c78650d5828c19f37dc538e49274fffb96..be94d19f9d4348e64ba4833fb902c984a32c1fb9 100644 --- a/FPS_test.py +++ b/FPS_test.py @@ -30,8 +30,13 @@ class FPS_YOLO(YOLO): #---------------------------------------------------------# # 给图像增加灰条,实现不失真的resize + # 也可以直接resize进行识别 #---------------------------------------------------------# - crop_img = np.array(letterbox_image(image, (self.model_image_size[1],self.model_image_size[0]))) + if self.letterbox_image: + crop_img = np.array(letterbox_image(image, (self.model_image_size[1],self.model_image_size[0]))) + else: + crop_img = image.convert('RGB') + crop_img = crop_img.resize((self.model_image_size[1],self.model_image_size[0]), Image.BICUBIC) photo = np.array(crop_img,dtype = np.float32) / 255.0 photo = np.transpose(photo, (2, 0, 1)) #---------------------------------------------------------# @@ -58,8 +63,16 @@ class FPS_YOLO(YOLO): top_label = np.array(batch_detections[top_index,-1],np.int32) top_bboxes = np.array(batch_detections[top_index,:4]) top_xmin, top_ymin, top_xmax, top_ymax = np.expand_dims(top_bboxes[:,0],-1),np.expand_dims(top_bboxes[:,1],-1),np.expand_dims(top_bboxes[:,2],-1),np.expand_dims(top_bboxes[:,3],-1) - # 去掉灰条 - boxes = yolo_correct_boxes(top_ymin,top_xmin,top_ymax,top_xmax,np.array([self.model_image_size[0],self.model_image_size[1]]),image_shape) + + if self.letterbox_image: + boxes = yolo_correct_boxes(top_ymin,top_xmin,top_ymax,top_xmax,np.array([self.model_image_size[0],self.model_image_size[1]]),image_shape) + else: + top_xmin = top_xmin / self.model_image_size[1] * image_shape[1] + top_ymin = top_ymin / self.model_image_size[0] * image_shape[0] + top_xmax = top_xmax / self.model_image_size[1] * image_shape[1] + top_ymax = top_ymax / self.model_image_size[0] * image_shape[0] + boxes = np.concatenate([top_ymin,top_xmin,top_ymax,top_xmax], axis=-1) + except: pass @@ -81,8 +94,16 @@ class FPS_YOLO(YOLO): top_label = np.array(batch_detections[top_index,-1],np.int32) top_bboxes = np.array(batch_detections[top_index,:4]) top_xmin, top_ymin, top_xmax, top_ymax = np.expand_dims(top_bboxes[:,0],-1),np.expand_dims(top_bboxes[:,1],-1),np.expand_dims(top_bboxes[:,2],-1),np.expand_dims(top_bboxes[:,3],-1) - # 去掉灰条 - boxes = yolo_correct_boxes(top_ymin,top_xmin,top_ymax,top_xmax,np.array([self.model_image_size[0],self.model_image_size[1]]),image_shape) + + if self.letterbox_image: + boxes = yolo_correct_boxes(top_ymin,top_xmin,top_ymax,top_xmax,np.array([self.model_image_size[0],self.model_image_size[1]]),image_shape) + else: + top_xmin = top_xmin / self.model_image_size[1] * image_shape[1] + top_ymin = top_ymin / self.model_image_size[0] * image_shape[0] + top_xmax = top_xmax / self.model_image_size[1] * image_shape[1] + top_ymax = top_ymax / self.model_image_size[0] * image_shape[0] + boxes = np.concatenate([top_ymin,top_xmin,top_ymax,top_xmax], axis=-1) + except: pass diff --git a/get_dr_txt.py b/get_dr_txt.py index c46f82899a6dfd2cc77b7953a2efd76b7b07ba60..c1b9e9f6b54231a9cf5d8f9857b7503f2bb86a12 100644 --- a/get_dr_txt.py +++ b/get_dr_txt.py @@ -33,8 +33,13 @@ class mAP_Yolo(YOLO): #---------------------------------------------------------# # 给图像增加灰条,实现不失真的resize + # 也可以直接resize进行识别 #---------------------------------------------------------# - crop_img = np.array(letterbox_image(image, (self.model_image_size[1],self.model_image_size[0]))) + if self.letterbox_image: + crop_img = np.array(letterbox_image(image, (self.model_image_size[1],self.model_image_size[0]))) + else: + crop_img = image.convert('RGB') + crop_img = crop_img.resize((self.model_image_size[1],self.model_image_size[0]), Image.BICUBIC) photo = np.array(crop_img,dtype = np.float32) / 255.0 photo = np.transpose(photo, (2, 0, 1)) #---------------------------------------------------------# @@ -69,7 +74,7 @@ class mAP_Yolo(YOLO): try: batch_detections = batch_detections[0].cpu().numpy() except: - return image + return #---------------------------------------------------------# # 对预测框进行得分筛选 @@ -85,8 +90,15 @@ class mAP_Yolo(YOLO): # 因此生成的top_bboxes是相对于有灰条的图像的 # 我们需要对其进行修改,去除灰条的部分。 #-----------------------------------------------------------------# - boxes = yolo_correct_boxes(top_ymin,top_xmin,top_ymax,top_xmax,np.array([self.model_image_size[0],self.model_image_size[1]]),image_shape) - + if self.letterbox_image: + boxes = yolo_correct_boxes(top_ymin,top_xmin,top_ymax,top_xmax,np.array([self.model_image_size[0],self.model_image_size[1]]),image_shape) + else: + top_xmin = top_xmin / self.model_image_size[1] * image_shape[1] + top_ymin = top_ymin / self.model_image_size[0] * image_shape[0] + top_xmax = top_xmax / self.model_image_size[1] * image_shape[1] + top_ymax = top_ymax / self.model_image_size[0] * image_shape[0] + boxes = np.concatenate([top_ymin,top_xmin,top_ymax,top_xmax], axis=-1) + for i, c in enumerate(top_label): predicted_class = self.class_names[c] score = str(top_conf[i]) diff --git a/yolo.py b/yolo.py index 7d439cbddd88b061a815a5f903e323f52efe785f..73268f6087da1f68e568c39e1d606ee804fbed8e 100644 --- a/yolo.py +++ b/yolo.py @@ -31,7 +31,12 @@ class YOLO(object): "model_image_size" : (416, 416, 3), "confidence" : 0.5, "iou" : 0.3, - "cuda" : True + "cuda" : True, + #---------------------------------------------------------------------# + # 该变量用于控制是否使用letterbox_image对输入图像进行不失真的resize, + # 在多次测试后,发现关闭letterbox_image直接resize的效果更好 + #---------------------------------------------------------------------# + "letterbox_image" : False, } @classmethod @@ -118,8 +123,13 @@ class YOLO(object): #---------------------------------------------------------# # 给图像增加灰条,实现不失真的resize + # 也可以直接resize进行识别 #---------------------------------------------------------# - crop_img = np.array(letterbox_image(image, (self.model_image_size[1],self.model_image_size[0]))) + if self.letterbox_image: + crop_img = np.array(letterbox_image(image, (self.model_image_size[1],self.model_image_size[0]))) + else: + crop_img = image.convert('RGB') + crop_img = crop_img.resize((self.model_image_size[1],self.model_image_size[0]), Image.BICUBIC) photo = np.array(crop_img,dtype = np.float32) / 255.0 photo = np.transpose(photo, (2, 0, 1)) #---------------------------------------------------------# @@ -170,8 +180,15 @@ class YOLO(object): # 因此生成的top_bboxes是相对于有灰条的图像的 # 我们需要对其进行修改,去除灰条的部分。 #-----------------------------------------------------------------# - boxes = yolo_correct_boxes(top_ymin,top_xmin,top_ymax,top_xmax,np.array([self.model_image_size[0],self.model_image_size[1]]),image_shape) - + if self.letterbox_image: + boxes = yolo_correct_boxes(top_ymin,top_xmin,top_ymax,top_xmax,np.array([self.model_image_size[0],self.model_image_size[1]]),image_shape) + else: + top_xmin = top_xmin / self.model_image_size[1] * image_shape[1] + top_ymin = top_ymin / self.model_image_size[0] * image_shape[0] + top_xmax = top_xmax / self.model_image_size[1] * image_shape[1] + top_ymax = top_ymax / self.model_image_size[0] * image_shape[0] + boxes = np.concatenate([top_ymin,top_xmin,top_ymax,top_xmax], axis=-1) + font = ImageFont.truetype(font='model_data/simhei.ttf',size=np.floor(3e-2 * np.shape(image)[1] + 0.5).astype('int32')) thickness = max((np.shape(image)[0] + np.shape(image)[1]) // self.model_image_size[0], 1)