зеркало из
https://github.com/ssciwr/AMMICO.git
synced 2025-10-30 05:26:05 +02:00
64 строки
2.6 KiB
Python
64 строки
2.6 KiB
Python
from misinformation.utils import AnalysisMethod
|
|
import torch
|
|
from PIL import Image
|
|
from lavis.models import load_model_and_preprocess
|
|
|
|
|
|
class SummaryDetector(AnalysisMethod):
|
|
def __init__(
|
|
self, subdict: dict
|
|
) -> None:
|
|
super().__init__(subdict)
|
|
self.subdict.update(self.set_keys())
|
|
self.image_summary = {
|
|
"const_image_summary": None,
|
|
"3_non-deterministic summary": None,
|
|
}
|
|
|
|
summary_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
summary_model, summary_vis_processors, _ = load_model_and_preprocess(
|
|
name="blip_caption", model_type="base_coco", is_eval=True, device=summary_device
|
|
)
|
|
|
|
def set_keys(self) -> dict:
|
|
params = {
|
|
"const_image_summary": None,
|
|
"3_non-deterministic summary": None,
|
|
}
|
|
return params
|
|
|
|
def analyse_image(self):
|
|
|
|
path = self.subdict["filename"]
|
|
raw_image = Image.open(path).convert("RGB")
|
|
image = self.summary_vis_processors["eval"](raw_image).unsqueeze(0).to(self.summary_device)
|
|
self.image_summary["const_image_summary"] = self.summary_model.generate({"image": image})[0]
|
|
self.image_summary["3_non-deterministic summary"] = self.summary_model.generate({"image": image}, use_nucleus_sampling=True, num_captions=3)
|
|
for key in self.image_summary:
|
|
self.subdict[key] = self.image_summary[key]
|
|
return self.subdict
|
|
|
|
summary_VQA_model, summary_VQA_vis_processors, summary_VQA_txt_processors = load_model_and_preprocess(name="blip_vqa", model_type="vqav2", is_eval=True, device=summary_device)
|
|
|
|
def analyse_questions(self, list_of_questions):
|
|
if (len(list_of_questions)>0):
|
|
path = self.subdict["filename"]
|
|
raw_image = Image.open(path).convert("RGB")
|
|
image = self.summary_VQA_vis_processors["eval"](raw_image).unsqueeze(0).to(self.summary_device)
|
|
question_batch =[]
|
|
for quest in list_of_questions:
|
|
question_batch.append(self.summary_VQA_txt_processors["eval"](quest))
|
|
batch_size = len(list_of_questions)
|
|
image_batch = image.repeat(batch_size, 1, 1, 1)
|
|
|
|
answers_batch = self.summary_VQA_model.predict_answers(samples={"image": image_batch, "text_input": question_batch}, inference_method="generate")
|
|
|
|
for q,a in zip(question_batch,answers_batch):
|
|
self.image_summary[q] = a
|
|
|
|
for key in self.image_summary:
|
|
self.subdict[key] = self.image_summary[key]
|
|
else:
|
|
print("Please, enter list of questions")
|
|
return self.subdict |