def merge_detections(image_name, results): # project = 'SKU_dataset' result_df = pandas.DataFrame() result_df['x1'] = results[:, 0].astype(int) result_df['y1'] = results[:, 1].astype(int) result_df['x2'] = results[:, 2].astype(int) result_df['y2'] = results[:, 3].astype(int) result_df['confidence'] = results[:, 4] result_df['hard_score'] = results[:, 5] result_df['uuid'] = 'object_label' result_df['label_type'] = 'object_label' # result_df['project'] = project result_df['image_name'] = image_name result_df.reset_index() result_df['id'] = result_df.index pixel_data = None duplicate_merger = DuplicateMerger() duplicate_merger.multiprocess = False duplicate_merger.compression_factor = 1 # project = result_df['project'].iloc[0] image_name = result_df['image_name'].iloc[0] if pixel_data is None: if 'COLAB_GPU' in os.environ: pixel_data = read_image_bgr(os.path.join('dataset', image_name)) else: pixel_data = read_image_bgr(os.path.join(root_dir(), image_name)) filtered_data = duplicate_merger.filter_duplicate_candidates( result_df, pixel_data) return filtered_data
def load_image(self, image_index): """ Load an image at the image_index. """ return read_image_bgr(self.image_path(image_index))
def infer(weights, image_dir, labels_to_name, output_file, threshold, hard_score_rate): model = models.load_model(weights, backbone_name='resnet50', convert=1, nms=False) labels_to_names = labels_to_name csv_data_lst = [] csv_data_lst.append( ['image_id', 'x1', 'y1', 'x2', 'y2', 'confidence', 'hard_score']) threshold = threshold hard_score_rate = hard_score_rate max_detections = 9999 image_dir = image_dir images = os.listdir(image_dir) if os.path.exists(str(output_file) + '/' + 'images'): shutil.rmtree(str(output_file) + '/' + 'images') os.makedirs(str(output_file) + '/' + 'images') # Run inference t0 = time.time() for img in images: image_path = image_dir + '/' + img t = time.time() image = read_image_bgr(image_path) # patches = image.unfold(0, 128, 128).unfold(1, 128, 128).unfold(2, 3, 3) draw = image.copy() draw = cv2.cvtColor(draw, cv2.COLOR_BGR2RGB) image = preprocess_image(image) image, scale = resize_image(image) boxes, hard_scores, labels, soft_scores = model.predict_on_batch( np.expand_dims(image, axis=0)) soft_scores = np.squeeze(soft_scores, axis=-1) soft_scores = hard_score_rate * hard_scores + ( 1 - hard_score_rate) * soft_scores boxes /= scale indices = np.where(hard_scores[0, :] > threshold)[0] scores = soft_scores[0][indices] hard_scores = hard_scores[0][indices] scores_sort = np.argsort(-scores)[:max_detections] image_boxes = boxes[0, indices[scores_sort], :] image_scores = scores[scores_sort] image_hard_scores = hard_scores[scores_sort] image_labels = labels[0, indices[scores_sort]] image_detections = np.concatenate([ image_boxes, np.expand_dims(image_scores, axis=1), np.expand_dims(image_labels, axis=1) ], axis=1) results = np.concatenate([ image_boxes, np.expand_dims(image_scores, axis=1), np.expand_dims(image_hard_scores, axis=1), np.expand_dims(image_labels, axis=1) ], axis=1) filtered_data = EmMerger.merge_detections(image_path, results) filtered_boxes = [] filtered_scores = [] filtered_labels = [] for ind, detection in filtered_data.iterrows(): box = np.asarray([ detection['x1'], detection['y1'], detection['x2'], detection['y2'] ]) filtered_boxes.append(box) filtered_scores.append(detection['confidence']) filtered_labels.append('{0:.2f}'.format(detection['hard_score'])) row = [ image_path, detection['x1'], detection['y1'], detection['x2'], detection['y2'], detection['confidence'], detection['hard_score'] ] csv_data_lst.append(row) for box, score, label in zip(filtered_boxes, filtered_scores, filtered_labels): # scores are sorted so we can break if score < threshold: break color = [ 31, 0, 255 ] #change the length of color array here based on the object classes, here I have hardcoded! b = box.astype(int) draw_box(draw, b, color=color) caption = "{} {:.3f}".format( labels_to_names[0], score) #hardcoded the index in the dictionary draw_caption(draw, b, caption) plt.figure(figsize=(20, 20)) plt.axis('off') plt.imshow(draw) plt.savefig( str(output_file) + '/' + 'images' + '/' + str(img.split('.')[0]) + '.png')