Views
No views yet
1import requests
2from transformers import AutoProcessor, LlavaForConditionalGeneration, Pipeline
3from transformers.utils import PushToHubMixin
4from PIL import Image
5import random
6import torch
7
8class ImageToQuestionPipeline():
9 def __init__(self, llava_model):
10 self.processor = AutoProcessor.from_pretrained(llava_model)
11 self.llava_model = LlavaForConditionalGeneration.from_pretrained(llava_model, torch_dtype=torch.float16).to(0)
12
13 def __call__(self, image_path):
14 commands = ["Generate a simple question\n","Suggest 1 correct answer\n","Suggest 3 incorrect answers\n", ""]
15 prompt ='''
16 <|user|>\n<image>\nDescribe this image in a passage\n<|end|>\n
17 <|assistant|>\n
18 '''
19 image_file = image_path
20 raw_image = Image.open(image_file)
21 inputs = self.processor(prompt, raw_image, return_tensors='pt').to(0)
22 artifacts = []
23 while commands:
24 inputs = self.processor(prompt, raw_image, return_tensors='pt').to(0)
25 output = self.llava_model.generate(**inputs, eos_token_id=32007, max_new_tokens=500, do_sample=False)
26 index = torch.where(output[0]==32001)[0][-1].item()
27 text = self.processor.decode(output[0][index:], skip_special_tokens=True)
28 artifacts.append(text)
29 prompt += "{}<|end|>\n<|user|>\n{}<|end|>\n<|assistant|>\n".format(text,commands.pop(0))
30
31 distractors = artifacts.pop(-1)
32 a = distractors.split("\n")
33 a = [x[3:] for x in a]
34
35 correct_answer = random.randint(0,3)
36 a.insert(correct_answer, artifacts[2])
37
38 a = ["{}) {}".format(i+1, a[i]) for i in range(len(a))]
39 answer = "Correct Answer: {}".format(correct_answer+1, a[correct_answer])
40 result = {}
41 result.update({"questions":artifacts[1]})
42 result.update({"choices":a})
43 result.update({"answer":answer})
44 result.update({"desc":artifacts[0]})
45 return result
46
47
48pipe = ImageToQuestionPipeline("Clyine1/phi3_image_question_generator")
49output = pipe(<image_file_path>)
50print(json.dumps(output, indent=4))
51
52"""
53Generated output:
54{
55 "questions": "What is the color of the shirt the girl in the center of the image is wearing?",
56 "choices": [
57 "1) The girl in the center of the image is wearing a pink shirt.",
58 "2) The girl in the center of the image is wearing a blue shirt.",
59 "3) The girl in the center of the image is wearing a red shirt.",
60 "4) The girl in the center of the image is wearing a green shirt."
61 ],
62 "answer": "Correct Answer: 1) The girl in the center of the image is wearing a pink shirt.",
63 "desc": "The image depicts a lively scene at a playground. In the foreground, a young girl is sitting on a green slide, her face reflecting a sense of surprise or shock. She is dressed in a pink shirt and a red hat. Behind her, a boy is standing on the same slide, his arms crossed in a defensive posture. He is wearing a red shirt and a blue hat.\n\nIn the background, a girl is sitting on a swing, her legs swinging back and forth. She is wearing a pink shirt and a red hat. Another girl is standing on a blue slide, her arms crossed in a similar defensive posture as the boy on the green slide. She is wearing a blue shirt and a red hat.\n\nIn the distance, a boy is standing on a yellow slide, his arms crossed in a defensive posture. He is wearing a yellow shirt and a red hat. Another boy is standing on a blue slide, his arms crossed in a similar defensive posture as the boy on the yellow slide. He is wearing a blue shirt and a red hat.\n\nThe playground is surrounded by trees and buildings, providing a natural and urban backdrop to the scene. The colors of the playground equipment and the clothing of the children are vibrant and contrasting, adding to the lively atmosphere of the image."
64}
65"""
66