|
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 |
|
|
|
|