def rnms_gpu(det_boxes, iou_threshold, device_id): if det_boxes.shape[0] == 0: return np.array([], np.int64) else: keep = rotate_gpu_nms(det_boxes, iou_threshold, device_id) keep = np.reshape(keep, [-1]) return np.array(keep, np.int64)
def rnms_gpu(det_boxes, iou_threshold, device_id): if det_boxes.shape[0] == 0: return np.array([], np.int64) else: assert det_boxes.shape[ 1] == 6, 'shape of det_boxes is not 6, {}'.format(det_boxes) keep = rotate_gpu_nms(det_boxes, iou_threshold, device_id) keep = np.reshape(keep, [-1]) return np.array(keep, np.int64)
def worker(self, gpu_id, images, det_net, result_queue): os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3]) # is RGB. not BGR img_batch = tf.cast(img_plac, tf.float32) if self.cfgs.NET_NAME in [ 'resnet152_v1d', 'resnet101_v1d', 'resnet50_v1d' ]: img_batch = (img_batch / 255 - tf.constant( self.cfgs.PIXEL_MEAN_)) / tf.constant(self.cfgs.PIXEL_STD) else: img_batch = img_batch - tf.constant(self.cfgs.PIXEL_MEAN) img_batch = tf.expand_dims(img_batch, axis=0) detection_boxes, detection_scores, detection_category = det_net.build_whole_detection_network( input_img_batch=img_batch) init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) restorer, restore_ckpt = det_net.get_restorer() config = tf.ConfigProto() config.gpu_options.allow_growth = True with tf.Session(config=config) as sess: sess.run(init_op) if not restorer is None: restorer.restore(sess, restore_ckpt) print('restore model %d ...' % gpu_id) for img_path in images: # if 'P0006' not in img_path: # continue img = cv2.imread(img_path) box_res_rotate = [] label_res_rotate = [] score_res_rotate = [] imgH = img.shape[0] imgW = img.shape[1] img_short_side_len_list = self.cfgs.IMG_SHORT_SIDE_LEN if isinstance( self.cfgs.IMG_SHORT_SIDE_LEN, list) else [self.cfgs.IMG_SHORT_SIDE_LEN] img_short_side_len_list = [ img_short_side_len_list[0] ] if not self.args.multi_scale else img_short_side_len_list if imgH < self.args.h_len: temp = np.zeros([self.args.h_len, imgW, 3], np.float32) temp[0:imgH, :, :] = img img = temp imgH = self.args.h_len if imgW < self.args.w_len: temp = np.zeros([imgH, self.args.w_len, 3], np.float32) temp[:, 0:imgW, :] = img img = temp imgW = self.args.w_len for hh in range(0, imgH, self.args.h_len - self.args.h_overlap): if imgH - hh - 1 < self.args.h_len: hh_ = imgH - self.args.h_len else: hh_ = hh for ww in range(0, imgW, self.args.w_len - self.args.w_overlap): if imgW - ww - 1 < self.args.w_len: ww_ = imgW - self.args.w_len else: ww_ = ww src_img = img[hh_:(hh_ + self.args.h_len), ww_:(ww_ + self.args.w_len), :] for short_size in img_short_side_len_list: max_len = self.cfgs.IMG_MAX_LENGTH if self.args.h_len < self.args.w_len: new_h, new_w = short_size, min( int(short_size * float(self.args.w_len) / self.args.h_len), max_len) else: new_h, new_w = min( int(short_size * float(self.args.h_len) / self.args.w_len), max_len), short_size img_resize = cv2.resize(src_img, (new_w, new_h)) resized_img, det_boxes_r_, det_scores_r_, det_category_r_ = \ sess.run( [img_batch, detection_boxes, detection_scores, detection_category], feed_dict={img_plac: img_resize[:, :, ::-1]} ) resized_h, resized_w = resized_img.shape[ 1], resized_img.shape[2] src_h, src_w = src_img.shape[0], src_img.shape[1] if len(det_boxes_r_) > 0: det_boxes_r_ = forward_convert( det_boxes_r_, False) det_boxes_r_[:, 0::2] *= (src_w / resized_w) det_boxes_r_[:, 1::2] *= (src_h / resized_h) for ii in range(len(det_boxes_r_)): box_rotate = det_boxes_r_[ii] box_rotate[0::2] = box_rotate[0::2] + ww_ box_rotate[1::2] = box_rotate[1::2] + hh_ box_res_rotate.append(box_rotate) label_res_rotate.append( det_category_r_[ii]) score_res_rotate.append(det_scores_r_[ii]) if self.args.flip_img: det_boxes_r_flip, det_scores_r_flip, det_category_r_flip = \ sess.run( [detection_boxes, detection_scores, detection_category], feed_dict={img_plac: cv2.flip(img_resize, flipCode=1)[:, :, ::-1]} ) if len(det_boxes_r_flip) > 0: det_boxes_r_flip = forward_convert( det_boxes_r_flip, False) det_boxes_r_flip[:, 0::2] *= (src_w / resized_w) det_boxes_r_flip[:, 1::2] *= (src_h / resized_h) for ii in range(len(det_boxes_r_flip)): box_rotate = det_boxes_r_flip[ii] box_rotate[0::2] = ( src_w - box_rotate[0::2]) + ww_ box_rotate[ 1::2] = box_rotate[1::2] + hh_ box_res_rotate.append(box_rotate) label_res_rotate.append( det_category_r_flip[ii]) score_res_rotate.append( det_scores_r_flip[ii]) det_boxes_r_flip, det_scores_r_flip, det_category_r_flip = \ sess.run( [detection_boxes, detection_scores, detection_category], feed_dict={img_plac: cv2.flip(img_resize, flipCode=0)[:, :, ::-1]} ) if len(det_boxes_r_flip) > 0: det_boxes_r_flip = forward_convert( det_boxes_r_flip, False) det_boxes_r_flip[:, 0::2] *= (src_w / resized_w) det_boxes_r_flip[:, 1::2] *= (src_h / resized_h) for ii in range(len(det_boxes_r_flip)): box_rotate = det_boxes_r_flip[ii] box_rotate[ 0::2] = box_rotate[0::2] + ww_ box_rotate[1::2] = ( src_h - box_rotate[1::2]) + hh_ box_res_rotate.append(box_rotate) label_res_rotate.append( det_category_r_flip[ii]) score_res_rotate.append( det_scores_r_flip[ii]) box_res_rotate = np.array(box_res_rotate) label_res_rotate = np.array(label_res_rotate) score_res_rotate = np.array(score_res_rotate) box_res_rotate_ = [] label_res_rotate_ = [] score_res_rotate_ = [] threshold = { 'roundabout': 0.1, 'tennis-court': 0.3, 'swimming-pool': 0.1, 'storage-tank': 0.2, 'soccer-ball-field': 0.3, 'small-vehicle': 0.2, 'ship': 0.2, 'plane': 0.3, 'large-vehicle': 0.1, 'helicopter': 0.2, 'harbor': 0.0001, 'ground-track-field': 0.3, 'bridge': 0.0001, 'basketball-court': 0.3, 'baseball-diamond': 0.3 } for sub_class in range(1, self.cfgs.CLASS_NUM + 1): index = np.where(label_res_rotate == sub_class)[0] if len(index) == 0: continue tmp_boxes_r = box_res_rotate[index] tmp_label_r = label_res_rotate[index] tmp_score_r = score_res_rotate[index] tmp_boxes_r_ = backward_convert(tmp_boxes_r, False) try: inx = nms_rotate.nms_rotate_cpu( boxes=np.array(tmp_boxes_r_), scores=np.array(tmp_score_r), iou_threshold=threshold[ self.label_name_map[sub_class]], max_output_size=5000) except: tmp_boxes_r_ = np.array(tmp_boxes_r_) tmp = np.zeros( [tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1]) tmp[:, 0:-1] = tmp_boxes_r_ tmp[:, -1] = np.array(tmp_score_r) # Note: the IoU of two same rectangles is 0, which is calculated by rotate_gpu_nms jitter = np.zeros( [tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1]) jitter[:, 0] += np.random.rand( tmp_boxes_r_.shape[0], ) / 1000 inx = rotate_gpu_nms( np.array(tmp, np.float32) + np.array(jitter, np.float32), float(threshold[self.label_name_map[sub_class]]), 0) box_res_rotate_.extend(np.array(tmp_boxes_r)[inx]) score_res_rotate_.extend(np.array(tmp_score_r)[inx]) label_res_rotate_.extend(np.array(tmp_label_r)[inx]) result_dict = { 'boxes': np.array(box_res_rotate_), 'scores': np.array(score_res_rotate_), 'labels': np.array(label_res_rotate_), 'image_id': img_path } result_queue.put_nowait(result_dict)
def eval_with_plac(self, img_dir, det_net, image_ext): os.environ["CUDA_VISIBLE_DEVICES"] = self.args.gpu # 1. preprocess img img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3]) # is RGB. not BGR img_batch = tf.cast(img_plac, tf.float32) pretrain_zoo = PretrainModelZoo() if self.cfgs.NET_NAME in pretrain_zoo.pth_zoo or self.cfgs.NET_NAME in pretrain_zoo.mxnet_zoo: img_batch = (img_batch / 255 - tf.constant( self.cfgs.PIXEL_MEAN_)) / tf.constant(self.cfgs.PIXEL_STD) else: img_batch = img_batch - tf.constant(self.cfgs.PIXEL_MEAN) img_batch = tf.expand_dims(img_batch, axis=0) detection_boxes, detection_scores, detection_category = det_net.build_whole_detection_network( input_img_batch=img_batch) init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) restorer, restore_ckpt = det_net.get_restorer() config = tf.ConfigProto() config.gpu_options.allow_growth = True with tf.Session(config=config) as sess: sess.run(init_op) if not restorer is None: restorer.restore(sess, restore_ckpt) print('restore model') all_boxes_r = [] imgs = os.listdir(img_dir) pbar = tqdm(imgs) for a_img_name in pbar: a_img_name = a_img_name.split(image_ext)[0] raw_img = cv2.imread( os.path.join(img_dir, a_img_name + image_ext)) raw_h, raw_w = raw_img.shape[0], raw_img.shape[1] det_boxes_r_all, det_scores_r_all, det_category_r_all = [], [], [] img_short_side_len_list = self.cfgs.IMG_SHORT_SIDE_LEN if isinstance( self.cfgs.IMG_SHORT_SIDE_LEN, list) else [self.cfgs.IMG_SHORT_SIDE_LEN] img_short_side_len_list = [ img_short_side_len_list[0] ] if not self.args.multi_scale else img_short_side_len_list for short_size in img_short_side_len_list: max_len = self.cfgs.IMG_MAX_LENGTH if raw_h < raw_w: new_h, new_w = short_size, min( int(short_size * float(raw_w) / raw_h), max_len) else: new_h, new_w = min( int(short_size * float(raw_h) / raw_w), max_len), short_size img_resize = cv2.resize(raw_img, (new_w, new_h)) resized_img, detected_boxes, detected_scores, detected_categories = \ sess.run( [img_batch, detection_boxes, detection_scores, detection_category], feed_dict={img_plac: img_resize[:, :, ::-1]} ) if detected_boxes.shape[0] == 0: continue resized_h, resized_w = resized_img.shape[ 1], resized_img.shape[2] detected_boxes = forward_convert(detected_boxes, False) detected_boxes[:, 0::2] *= (raw_w / resized_w) detected_boxes[:, 1::2] *= (raw_h / resized_h) det_boxes_r_all.extend(detected_boxes) det_scores_r_all.extend(detected_scores) det_category_r_all.extend(detected_categories) det_boxes_r_all = np.array(det_boxes_r_all) det_scores_r_all = np.array(det_scores_r_all) det_category_r_all = np.array(det_category_r_all) box_res_rotate_ = [] label_res_rotate_ = [] score_res_rotate_ = [] if det_scores_r_all.shape[0] != 0: for sub_class in range(1, self.cfgs.CLASS_NUM + 1): index = np.where(det_category_r_all == sub_class)[0] if len(index) == 0: continue tmp_boxes_r = det_boxes_r_all[index] tmp_label_r = det_category_r_all[index] tmp_score_r = det_scores_r_all[index] if self.args.multi_scale: tmp_boxes_r_ = backward_convert(tmp_boxes_r, False) # try: # inx = nms_rotate.nms_rotate_cpu(boxes=np.array(tmp_boxes_r_), # scores=np.array(tmp_score_r), # iou_threshold=self.cfgs.NMS_IOU_THRESHOLD, # max_output_size=5000) # except: tmp_boxes_r_ = np.array(tmp_boxes_r_) tmp = np.zeros([ tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1 ]) tmp[:, 0:-1] = tmp_boxes_r_ tmp[:, -1] = np.array(tmp_score_r) # Note: the IoU of two same rectangles is 0, which is calculated by rotate_gpu_nms jitter = np.zeros([ tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1 ]) jitter[:, 0] += np.random.rand( tmp_boxes_r_.shape[0], ) / 1000 inx = rotate_gpu_nms( np.array(tmp, np.float32) + np.array(jitter, np.float32), float(self.cfgs.NMS_IOU_THRESHOLD), 0) else: inx = np.arange(0, tmp_score_r.shape[0]) box_res_rotate_.extend(np.array(tmp_boxes_r)[inx]) score_res_rotate_.extend(np.array(tmp_score_r)[inx]) label_res_rotate_.extend(np.array(tmp_label_r)[inx]) if len(box_res_rotate_) == 0: all_boxes_r.append(np.array([])) continue det_boxes_r_ = np.array(box_res_rotate_) det_scores_r_ = np.array(score_res_rotate_) det_category_r_ = np.array(label_res_rotate_) if self.args.draw_imgs: detected_indices = det_scores_r_ >= self.cfgs.VIS_SCORE detected_scores = det_scores_r_[detected_indices] detected_boxes = det_boxes_r_[detected_indices] detected_categories = det_category_r_[detected_indices] detected_boxes = backward_convert(detected_boxes, False) drawer = DrawBox(self.cfgs) det_detections_r = drawer.draw_boxes_with_label_and_scores( raw_img[:, :, ::-1], boxes=detected_boxes, labels=detected_categories, scores=detected_scores, method=1, in_graph=True) save_dir = os.path.join('test_hrsc', self.cfgs.VERSION, 'hrsc2016_img_vis') tools.makedirs(save_dir) cv2.imwrite(save_dir + '/{}.jpg'.format(a_img_name), det_detections_r[:, :, ::-1]) det_boxes_r_ = backward_convert(det_boxes_r_, False) x_c, y_c, w, h, theta = det_boxes_r_[:, 0], det_boxes_r_[:, 1], det_boxes_r_[:, 2], \ det_boxes_r_[:, 3], det_boxes_r_[:, 4] boxes_r = np.transpose(np.stack([x_c, y_c, w, h, theta])) dets_r = np.hstack((det_category_r_.reshape(-1, 1), det_scores_r_.reshape(-1, 1), boxes_r)) all_boxes_r.append(dets_r) pbar.set_description("Eval image %s" % a_img_name) # fw1 = open(cfgs.VERSION + '_detections_r.pkl', 'wb') # pickle.dump(all_boxes_r, fw1) return all_boxes_r
def worker(self, gpu_id, images, det_net, result_queue): os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) # 1. preprocess img img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3]) # is RGB. not BGR img_batch = tf.cast(img_plac, tf.float32) if self.cfgs.NET_NAME in ['resnet152_v1d', 'resnet101_v1d', 'resnet50_v1d', 'resnet152_v1b', 'resnet101_v1b', 'resnet50_v1b', 'resnet34_v1b', 'resnet18_v1b']: img_batch = (img_batch / 255 - tf.constant(self.cfgs.PIXEL_MEAN_)) / tf.constant(self.cfgs.PIXEL_STD) else: img_batch = img_batch - tf.constant(self.cfgs.PIXEL_MEAN) img_batch = tf.expand_dims(img_batch, axis=0) detection_boxes, detection_scores, detection_category = det_net.build_whole_detection_network( input_img_batch=img_batch, gtboxes_batch_h=None, gtboxes_batch_r=None, gpu_id=0) init_op = tf.group( tf.global_variables_initializer(), tf.local_variables_initializer() ) restorer, restore_ckpt = det_net.get_restorer() config = tf.ConfigProto() config.gpu_options.allow_growth = True with tf.Session(config=config) as sess: sess.run(init_op) if not restorer is None: restorer.restore(sess, restore_ckpt) print('restore model %d ...' % gpu_id) for a_img in images: raw_img = cv2.imread(a_img) raw_h, raw_w = raw_img.shape[0], raw_img.shape[1] det_boxes_r_all, det_scores_r_all, det_category_r_all = [], [], [] img_short_side_len_list = self.cfgs.IMG_SHORT_SIDE_LEN if isinstance(self.cfgs.IMG_SHORT_SIDE_LEN, list) else [ self.cfgs.IMG_SHORT_SIDE_LEN] img_short_side_len_list = [img_short_side_len_list[0]] if not self.args.multi_scale else img_short_side_len_list for short_size in img_short_side_len_list: max_len = self.cfgs.IMG_MAX_LENGTH if raw_h < raw_w: new_h, new_w = short_size, min(int(short_size * float(raw_w) / raw_h), max_len) else: new_h, new_w = min(int(short_size * float(raw_h) / raw_w), max_len), short_size img_resize = cv2.resize(raw_img, (new_w, new_h)) resized_img, detected_boxes, detected_scores, detected_categories = \ sess.run( [img_batch, detection_boxes, detection_scores, detection_category], feed_dict={img_plac: img_resize[:, :, ::-1]} ) detected_indices = detected_scores >= self.cfgs.VIS_SCORE detected_scores = detected_scores[detected_indices] detected_boxes = detected_boxes[detected_indices] detected_categories = detected_categories[detected_indices] if detected_boxes.shape[0] == 0: continue resized_h, resized_w = resized_img.shape[1], resized_img.shape[2] detected_boxes = forward_convert(detected_boxes, False) detected_boxes[:, 0::2] *= (raw_w / resized_w) detected_boxes[:, 1::2] *= (raw_h / resized_h) det_boxes_r_all.extend(detected_boxes) det_scores_r_all.extend(detected_scores) det_category_r_all.extend(detected_categories) if self.args.flip_img: detected_boxes, detected_scores, detected_categories = \ sess.run( [detection_boxes, detection_scores, detection_category], feed_dict={img_plac: cv2.flip(img_resize, flipCode=1)[:, :, ::-1]} ) detected_indices = detected_scores >= self.cfgs.VIS_SCORE detected_scores = detected_scores[detected_indices] detected_boxes = detected_boxes[detected_indices] detected_categories = detected_categories[detected_indices] if detected_boxes.shape[0] == 0: continue resized_h, resized_w = resized_img.shape[1], resized_img.shape[2] detected_boxes = forward_convert(detected_boxes, False) detected_boxes[:, 0::2] *= (raw_w / resized_w) detected_boxes[:, 0::2] = (raw_w - detected_boxes[:, 0::2]) detected_boxes[:, 1::2] *= (raw_h / resized_h) det_boxes_r_all.extend(sort_corners(detected_boxes)) det_scores_r_all.extend(detected_scores) det_category_r_all.extend(detected_categories) detected_boxes, detected_scores, detected_categories = \ sess.run( [detection_boxes, detection_scores, detection_category], feed_dict={img_plac: cv2.flip(img_resize, flipCode=0)[:, :, ::-1]} ) detected_indices = detected_scores >= self.cfgs.VIS_SCORE detected_scores = detected_scores[detected_indices] detected_boxes = detected_boxes[detected_indices] detected_categories = detected_categories[detected_indices] if detected_boxes.shape[0] == 0: continue resized_h, resized_w = resized_img.shape[1], resized_img.shape[2] detected_boxes = forward_convert(detected_boxes, False) detected_boxes[:, 0::2] *= (raw_w / resized_w) detected_boxes[:, 1::2] *= (raw_h / resized_h) detected_boxes[:, 1::2] = (raw_h - detected_boxes[:, 1::2]) det_boxes_r_all.extend(sort_corners(detected_boxes)) det_scores_r_all.extend(detected_scores) det_category_r_all.extend(detected_categories) det_boxes_r_all = np.array(det_boxes_r_all) det_scores_r_all = np.array(det_scores_r_all) det_category_r_all = np.array(det_category_r_all) box_res_rotate_ = [] label_res_rotate_ = [] score_res_rotate_ = [] if det_scores_r_all.shape[0] != 0: for sub_class in range(1, self.cfgs.CLASS_NUM + 1): index = np.where(det_category_r_all == sub_class)[0] if len(index) == 0: continue tmp_boxes_r = det_boxes_r_all[index] tmp_label_r = det_category_r_all[index] tmp_score_r = det_scores_r_all[index] if self.args.multi_scale: tmp_boxes_r_ = backward_convert(tmp_boxes_r, False) # try: # inx = nms_rotate.nms_rotate_cpu(boxes=np.array(tmp_boxes_r_), # scores=np.array(tmp_score_r), # iou_threshold=self.cfgs.NMS_IOU_THRESHOLD, # max_output_size=5000) # except: tmp_boxes_r_ = np.array(tmp_boxes_r_) tmp = np.zeros([tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1]) tmp[:, 0:-1] = tmp_boxes_r_ tmp[:, -1] = np.array(tmp_score_r) # Note: the IoU of two same rectangles is 0, which is calculated by rotate_gpu_nms jitter = np.zeros([tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1]) jitter[:, 0] += np.random.rand(tmp_boxes_r_.shape[0], ) / 1000 inx = rotate_gpu_nms(np.array(tmp, np.float32) + np.array(jitter, np.float32), float(self.cfgs.NMS_IOU_THRESHOLD), 0) else: inx = np.arange(0, tmp_score_r.shape[0]) box_res_rotate_.extend(np.array(tmp_boxes_r)[inx]) score_res_rotate_.extend(np.array(tmp_score_r)[inx]) label_res_rotate_.extend(np.array(tmp_label_r)[inx]) box_res_rotate_ = np.array(box_res_rotate_) score_res_rotate_ = np.array(score_res_rotate_) label_res_rotate_ = np.array(label_res_rotate_) result_dict = {'scales': [1, 1], 'boxes': box_res_rotate_, 'scores': score_res_rotate_, 'labels': label_res_rotate_, 'image_id': a_img} result_queue.put_nowait(result_dict)