kartik91 commited on
Commit
6dc5be7
1 Parent(s): 7fe949f

image to story

Browse files
Files changed (1) hide show
  1. app.py +87 -0
app.py ADDED
@@ -0,0 +1,87 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline
2
+ from langchain_core.runnables import RunnableLambda
3
+ from langchain_huggingface import HuggingFacePipeline
4
+ from PIL import Image
5
+
6
+ pipe1 = pipeline("object-detection", model="facebook/detr-resnet-50")
7
+ pipe2 = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
8
+ repo_id = "mistralai/Mistral-7B-Instruct-v0.2"
9
+
10
+
11
+ llm = HuggingFacePipeline.from_model_id(
12
+ model_id=repo_id,
13
+ task="text-generation",
14
+ pipeline_kwargs={"max_length": 100,"temperature":0.5},
15
+ )
16
+ def reduce_add(a):
17
+ ll=dict()
18
+ for i in a:
19
+ if i['score'] > 0.89:
20
+ if i['label'] not in ll.keys():
21
+ ll[i['label']] = 1
22
+ else:
23
+ ll[i['label']]+=1
24
+ return "there are \n"+', \n'.join([str(i[1])+' '+i[0] for i in ll.items() ])
25
+
26
+ def image_segmentation_tool(image: str):
27
+ # image = Image.open(image_path)
28
+ segmentation_results = pipe1(image)
29
+ if reduce_add(segmentation_results) == "there are \n":
30
+ raise Passs()
31
+ return reduce_add(segmentation_results)
32
+
33
+ def image_caption_tool(image: str):
34
+ # image = Image.open(image_path)
35
+ segmentation_results = pipe2(image)
36
+ if segmentation_results[0]["generated_text"] == "":
37
+ raise Passs("no result found use different image to create story")
38
+ return segmentation_results[0]["generated_text"]
39
+
40
+ from langchain_core.prompts import PromptTemplate
41
+
42
+
43
+ def story_generation_tool(segmentation_results):
44
+ prompt_template = """
45
+ You are a storyteller. Based on the following segmentation results, create a story:
46
+ {segmentation_results}
47
+
48
+ Story:
49
+ """
50
+ prompt = PromptTemplate.from_template(prompt_template)
51
+ story = prompt | llm
52
+ return story.invoke(input={"segmentation_results":segmentation_results})
53
+
54
+ # def translation_tool(english_text):
55
+ # prompt_template = """
56
+ # You are a translator. Translate the following English text to Hindi:
57
+ # {english_text}
58
+
59
+ # Translation:
60
+ # """
61
+ # prompt = PromptTemplate.from_template(prompt_template)
62
+ # translation = prompt | llm
63
+ # return translation.invoke(input={"english_text": english_text})
64
+
65
+
66
+ runnable = RunnableLambda(image_segmentation_tool).with_fallbacks([RunnableLambda(image_caption_tool)])
67
+ runnable2 = RunnableLambda(story_generation_tool)
68
+ # runnable3 = RunnableLambda(translation_tool)
69
+
70
+ chain = runnable | runnable2
71
+
72
+ import gradio as gr
73
+
74
+ title = "Image to short Story Generator"
75
+ description = """
76
+ Upload an image, and this app will generate a short story based on the image.
77
+ """
78
+
79
+ def sepia(input_img):
80
+ sepia_img=chain.invoke(input_img)
81
+ return sepia_img
82
+
83
+ demo = gr.Interface(sepia, gr.Image(type='pil'), "textarea",title=title,
84
+ description=description,live=True
85
+ )
86
+ if __name__ == "__main__":
87
+ demo.launch()