| import os | |
| import sys | |
| from .grit_src.image_dense_captions import image_caption_api, init_demo, dense_pred_to_caption, dense_pred_to_caption_only_name,dense_pred_to_caption_tuple | |
| from detectron2.data.detection_utils import read_image | |
| class DenseCaptioning(): | |
| def __init__(self, device): | |
| self.device = device | |
| self.demo = None | |
| def initialize_model(self, model_weight): | |
| self.demo = init_demo(self.device, model_weight=model_weight) | |
| def initialize_model_det(self, model_weight): | |
| self.demo = init_demo(self.device, model_weight = model_weight, task="ObjectDet") | |
| def image_dense_caption(self, image_src): | |
| dense_caption = image_caption_api(image_src, self.device) | |
| print('\033[1;35m' + '*' * 100 + '\033[0m') | |
| print("Step2, Dense Caption:\n") | |
| print(dense_caption) | |
| print('\033[1;35m' + '*' * 100 + '\033[0m') | |
| return dense_caption | |
| def run_caption_api(self,image_src): | |
| img = read_image(image_src, format="BGR") | |
| print(img.shape) | |
| predictions, visualized_output = self.demo.run_on_image(img) | |
| new_caption = dense_pred_to_caption_only_name(predictions) | |
| return new_caption | |
| def run_caption_tensor(self,img): | |
| predictions, visualized_output = self.demo.run_on_image(img) | |
| new_caption = dense_pred_to_caption_tuple(predictions) | |
| return new_caption, visualized_output | |
| def run_det_tensor(self,img): | |
| predictions, visualized_output = self.demo.run_on_image(img) | |
| return predictions, visualized_output | |