welcometoFightclub commited on
Commit
f0240fe
1 Parent(s): f932c6a

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +35 -0
model.py CHANGED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import requests
2
+ from PIL import Image
3
+ from transformers import BlipProcessor, BlipForConditionalGeneration
4
+
5
+ # Initialize processor and model
6
+ processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
7
+ model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large").to("cuda")
8
+
9
+ # Function to process and caption an image from a URL
10
+ def caption_image(image_url):
11
+ try:
12
+ # Load image from the provided URL
13
+ raw_image = Image.open(requests.get(image_url, stream=True).raw).convert('RGB')
14
+
15
+ # Conditional image captioning
16
+ text = "a photography of"
17
+ inputs = processor(raw_image, text, return_tensors="pt").to("cuda")
18
+ out = model.generate(**inputs)
19
+ conditional_caption = processor.decode(out[0], skip_special_tokens=True)
20
+
21
+ # Unconditional image captioning
22
+ inputs = processor(raw_image, return_tensors="pt").to("cuda")
23
+ out = model.generate(**inputs)
24
+ unconditional_caption = processor.decode(out[0], skip_special_tokens=True)
25
+
26
+ # Print the results
27
+ print("Conditional Caption:", conditional_caption)
28
+ print("Unconditional Caption:", unconditional_caption)
29
+
30
+ except Exception as e:
31
+ print(f"Error occurred: {e}")
32
+
33
+ # Get image URL from user input
34
+ image_url = input("Enter the image URL: ")
35
+ caption_image(image_url)