Spaces:
Runtime error
Runtime error
Commit
•
37ce9d7
1
Parent(s):
3b57595
Update image_caption.py
Browse files- image_caption.py +54 -16
image_caption.py
CHANGED
@@ -4,26 +4,31 @@ import os
|
|
4 |
from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
|
5 |
import torch
|
6 |
from PIL import Image
|
7 |
-
|
|
|
8 |
|
9 |
class Caption:
|
10 |
def __init__(self):
|
11 |
-
self.model = VisionEncoderDecoderModel.from_pretrained(
|
12 |
-
"nlpconnect/vit-gpt2-image-captioning"
|
13 |
-
)
|
14 |
-
self.feature_extractor = ViTImageProcessor.from_pretrained(
|
15 |
-
"nlpconnect/vit-gpt2-image-captioning"
|
16 |
-
)
|
17 |
-
self.tokenizer = AutoTokenizer.from_pretrained(
|
18 |
-
"nlpconnect/vit-gpt2-image-captioning"
|
19 |
-
)
|
20 |
|
21 |
-
|
22 |
-
|
23 |
-
self.model.
|
24 |
-
|
25 |
-
|
26 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
27 |
|
28 |
|
29 |
def predict_step(self,image_paths):
|
@@ -45,6 +50,39 @@ class Caption:
|
|
45 |
preds = [pred.strip() for pred in preds]
|
46 |
return preds
|
47 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
48 |
def get_args(self):
|
49 |
parser = argparse.ArgumentParser()
|
50 |
parser.add_argument( "-i",
|
|
|
4 |
from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
|
5 |
import torch
|
6 |
from PIL import Image
|
7 |
+
import io
|
8 |
+
import google.generativeai as genai
|
9 |
|
10 |
class Caption:
|
11 |
def __init__(self):
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
12 |
|
13 |
+
self.api_key = 'AIzaSyAFG94rVbm9eWepO5uPGsMha8XJ-sHbMdA'
|
14 |
+
genai.configure(api_key=self.api_key)
|
15 |
+
self.model = genai.GenerativeModel(model_name="gemini-pro-vision")
|
16 |
+
# self.model = VisionEncoderDecoderModel.from_pretrained(
|
17 |
+
# "nlpconnect/vit-gpt2-image-captioning"
|
18 |
+
# )
|
19 |
+
# self.feature_extractor = ViTImageProcessor.from_pretrained(
|
20 |
+
# "nlpconnect/vit-gpt2-image-captioning"
|
21 |
+
# )
|
22 |
+
# self.tokenizer = AutoTokenizer.from_pretrained(
|
23 |
+
# "nlpconnect/vit-gpt2-image-captioning"
|
24 |
+
# )
|
25 |
+
|
26 |
+
# # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
27 |
+
# self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
28 |
+
# self.model.to(self.device)
|
29 |
+
# self.max_length = 16
|
30 |
+
# self.num_beams = 4
|
31 |
+
# self.gen_kwargs = {"max_length": self.max_length, "num_beams": self.num_beams}
|
32 |
|
33 |
|
34 |
def predict_step(self,image_paths):
|
|
|
50 |
preds = [pred.strip() for pred in preds]
|
51 |
return preds
|
52 |
|
53 |
+
def predict_from_memory(self, image_buffers):
|
54 |
+
images = []
|
55 |
+
|
56 |
+
for image_buffer in image_buffers:
|
57 |
+
# Ensure the buffer is positioned at the start
|
58 |
+
if isinstance(image_buffer, io.BytesIO):
|
59 |
+
image_buffer.seek(0)
|
60 |
+
try:
|
61 |
+
i_image = Image.open(image_buffer)
|
62 |
+
if i_image.mode != "RGB":
|
63 |
+
i_image = i_image.convert("RGB")
|
64 |
+
images.append(i_image)
|
65 |
+
except Exception as e:
|
66 |
+
print(f"Failed to process image buffer: {str(e)}")
|
67 |
+
continue
|
68 |
+
|
69 |
+
return self.process_images(images)
|
70 |
+
|
71 |
+
def process_images(self, images):
|
72 |
+
pixel_values = self.feature_extractor(images=images, return_tensors="pt").pixel_values
|
73 |
+
pixel_values = pixel_values.to(self.device)
|
74 |
+
output_ids = self.model.generate(pixel_values, **self.gen_kwargs)
|
75 |
+
preds = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)
|
76 |
+
preds = [pred.strip() for pred in preds]
|
77 |
+
return preds
|
78 |
+
|
79 |
+
def predict_image_caption_gemini(self,img):
|
80 |
+
prompt = "Describe the main focus of this image in detail."
|
81 |
+
response = self.model.generate_content([prompt, img], stream=True)
|
82 |
+
response.resolve()
|
83 |
+
print("Derived data",response.text)
|
84 |
+
return response.text
|
85 |
+
|
86 |
def get_args(self):
|
87 |
parser = argparse.ArgumentParser()
|
88 |
parser.add_argument( "-i",
|