erayman09 commited on
Commit
3d3d03f
·
verified ·
1 Parent(s): ad2921a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +67 -0
app.py ADDED
@@ -0,0 +1,67 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
3
+ from PIL import Image
4
+
5
+ # Load the Image-to-Text (OCR) model
6
+ ocr_model = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
7
+
8
+ # Load the Text Generation model
9
+ story_model_name = "EleutherAI/gpt-neo-2.7B"
10
+ story_tokenizer = AutoTokenizer.from_pretrained(story_model_name)
11
+ story_model = AutoModelForCausalLM.from_pretrained(story_model_name)
12
+
13
+ # Function to extract text description from an image
14
+ def extract_description(image_path):
15
+ try:
16
+ # Use the OCR model to extract a caption/description from the image
17
+ result = ocr_model(Image.open(image_path))
18
+ return result[0]["generated_text"]
19
+ except Exception as e:
20
+ return f"Error extracting description: {e}"
21
+
22
+ # Function to generate a story based on the extracted description
23
+ def generate_story(description):
24
+ try:
25
+ # Format the input prompt for the story
26
+ prompt = f"Create a creative story based on this description: {description}"
27
+
28
+ # Use the story model to generate text
29
+ inputs = story_tokenizer(prompt, return_tensors="pt", truncation=True)
30
+ outputs = story_model.generate(inputs["input_ids"], max_length=300, num_return_sequences=1, temperature=0.8)
31
+ story = story_tokenizer.decode(outputs[0], skip_special_tokens=True)
32
+ return story
33
+ except Exception as e:
34
+ return f"Error generating story: {e}"
35
+
36
+ # Main function to process the image and generate a story
37
+ def create_story(image):
38
+ try:
39
+ # Step 1: Extract a description from the image
40
+ description = extract_description(image)
41
+ if not description or "Error" in description:
42
+ return description, None
43
+
44
+ # Step 2: Generate a story from the extracted description
45
+ story = generate_story(description)
46
+
47
+ # Combine the description and story for the output
48
+ output = f"📷 Extracted Description:\n{description}\n\n📖 Generated Story:\n{story}"
49
+ return output
50
+ except Exception as e:
51
+ return f"Error processing the image: {e}"
52
+
53
+ # Gradio interface
54
+ interface = gr.Interface(
55
+ fn=create_story,
56
+ inputs=gr.Image(label="Upload an Image (PNG, JPG, JPEG)"),
57
+ outputs=gr.Textbox(label="Generated Story"),
58
+ title="Text-Based Story Creator",
59
+ description=(
60
+ "Upload an image, and this app will generate a creative story based on the description of the image. "
61
+ "It uses advanced AI models for image-to-text conversion and story generation."
62
+ ),
63
+ allow_flagging="never"
64
+ )
65
+
66
+ if __name__ == "__main__":
67
+ interface.launch()