1Noura commited on
Commit
47a9da5
·
verified ·
1 Parent(s): 544f7f8

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +25 -11
app.py CHANGED
@@ -8,9 +8,12 @@ import wget
8
  device = "cuda" if torch.cuda.is_available() else "cpu"
9
 
10
  # Load the models
 
11
  caption_image = pipeline("image-to-text", model="Salesforce/blip-image-captioning-large", device=device)
 
12
  sd_pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5").to(device)
13
 
 
14
  translator = pipeline(
15
  task="translation",
16
  model="facebook/nllb-200-distilled-600M",
@@ -18,37 +21,48 @@ translator = pipeline(
18
  device=device
19
  )
20
 
 
 
 
21
  # Function to generate images based on the image's caption
22
  def generate_image_and_translate(image, num_images=1):
 
23
  caption_en = caption_image(image)[0]['generated_text']
 
 
24
  caption_ar = translator(caption_en, src_lang="eng_Latn", tgt_lang="arb_Arab")[0]['translation_text']
25
 
26
  generated_images = []
 
 
27
  for _ in range(num_images):
28
  generated_image = sd_pipeline(prompt=caption_en).images[0]
29
  generated_images.append(generated_image)
30
 
 
31
  return generated_images, caption_en, caption_ar
32
 
33
  # Set up the Gradio interface
34
  interface = gr.Interface(
35
- fn=generate_image_and_translate,
36
  inputs=[
37
- gr.Image(type="pil", label="Upload Image"),
38
- gr.Slider(minimum=1, maximum=10, label="Number of Images", value=1, step=1)
39
  ],
40
  outputs=[
41
- gr.Gallery(label="Generated Images"),
42
- gr.Textbox(label="Generated Caption (English)", interactive=False),
43
- gr.Textbox(label="Translated Caption (Arabic)", interactive=False)
 
44
  ],
45
- title="Image Generation and Translation",
46
- description="Upload an image to generate new images based on its caption and translate the caption into Arabic.",
47
- examples=[
48
  ["sea.jpg", 3]
 
 
49
  ]
50
  )
51
 
52
  # Launch the Gradio application
53
- if __name__ == "__main__":
54
- interface.launch()
 
8
  device = "cuda" if torch.cuda.is_available() else "cpu"
9
 
10
  # Load the models
11
+ # Image captioning model to generate captions from uploaded images
12
  caption_image = pipeline("image-to-text", model="Salesforce/blip-image-captioning-large", device=device)
13
+ # Stable Diffusion model for generating new images based on captions
14
  sd_pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5").to(device)
15
 
16
+ # Load the translation model (English to Arabic)
17
  translator = pipeline(
18
  task="translation",
19
  model="facebook/nllb-200-distilled-600M",
 
21
  device=device
22
  )
23
 
24
+ # Download the image
25
+ url1 = "https://github.com/Shahad-b/Image-database/blob/main/sea.jpg?raw=true"
26
+ sea = wget.download(url1)
27
  # Function to generate images based on the image's caption
28
  def generate_image_and_translate(image, num_images=1):
29
+ # Generate caption in English from the uploaded image
30
  caption_en = caption_image(image)[0]['generated_text']
31
+
32
+ # Translate the English caption to Arabic
33
  caption_ar = translator(caption_en, src_lang="eng_Latn", tgt_lang="arb_Arab")[0]['translation_text']
34
 
35
  generated_images = []
36
+
37
+ # Generate the specified number of images based on the English caption
38
  for _ in range(num_images):
39
  generated_image = sd_pipeline(prompt=caption_en).images[0]
40
  generated_images.append(generated_image)
41
 
42
+ # Return the generated images along with both captions
43
  return generated_images, caption_en, caption_ar
44
 
45
  # Set up the Gradio interface
46
  interface = gr.Interface(
47
+ fn=generate_image_and_translate, # Function to call when processing input
48
  inputs=[
49
+ gr.Image(type="pil", label="Upload Image"), # Input for image upload
50
+ gr.Slider(minimum=1, maximum=10, label="Number of Images", value=1, step=1) # Slider to select number of images
51
  ],
52
  outputs=[
53
+ gr.Gallery(label="Generated Images"), # Output for displaying generated images
54
+ gr.Textbox(label="Generated Caption (English)", interactive=False), # Output for English caption
55
+ gr.Textbox(label="Translated Caption (Arabic)", interactive=False)# Output for Arabic caption
56
+
57
  ],
58
+ title="Image Generation and Translation", # Title of the interface
59
+ description="Upload an image to generate new images based on its caption and translate the caption into Arabic.", # Description
60
+ examples=[ # Example input
61
  ["sea.jpg", 3]
62
+ # ["Cat.jpeg", 4],
63
+ # ["Car.jpeg", 2]
64
  ]
65
  )
66
 
67
  # Launch the Gradio application
68
+ interface.launch()