Update README.md

README.md

@@ -21,48 +21,52 @@ from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoToken
 import torch
 from PIL import Image
 import re
+import requests
 
 def prepare(text):
-    text = text.replace('. ', '.').replace(' .', '.')
-    text = text.replace('( ', '(').replace(' (', '(')
-    text = text.replace(') ', ')').replace(' )', ')')
-    text = text.replace(': ', ':').replace(' :', ':')
-    text = text.replace('_ ', '_').replace(' _', '_')
-    text = text.replace(',(())', '').replace('(()),', '')
-    for i in range(10):
-        text = text.replace(')))', '))').replace('(((', '((')
     text = re.sub(r'<[^>]*>', '', text)
+    text = ','.join(list(set(text.split(',')))[:-1])
+    for i in range(5):
+        if text[0]==',' or  text[0]==' ':
+            text=text[1:]
+    
     return text
 
-model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
-feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
-tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+path_to_model = "ifmain/vit-gpt2-image2promt-stable-diffusion"
+
+model = VisionEncoderDecoderModel.from_pretrained(path_to_model)
+feature_extractor = ViTImageProcessor.from_pretrained(path_to_model)
+tokenizer = AutoTokenizer.from_pretrained(path_to_model)
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-max_length = 16
+max_length = 256
 num_beams = 4
 gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
-def predict_step(image_paths):
-  images = []
-  for image_path in image_paths:
-    i_image = Image.open(image_path)
-    if i_image.mode != "RGB":
-      i_image = i_image.convert(mode="RGB")
-
-    images.append(i_image)
 
-  pixel_values = feature_extractor(images=images, return_tensors="pt").pixel_values
-  pixel_values = pixel_values.to(device)
 
-  output_ids = model.generate(pixel_values, **gen_kwargs)
-
-  preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
-  preds = [prepare(pred).strip() for pred in preds]
-  return preds
-
-predict_step(['doctor.e16ba4e4.jpg']) # ['a woman in a hospital bed with a woman in a hospital bed']
+def predict_step(image_paths):
+    images = []
+    for image_path in image_paths:
+        if 'http' in image_path:
+                i_image = Image.open(requests.get(image_path, stream=True).raw).convert('RGB')
+        else:
+            i_image = Image.open(image_path).convert('RGB')
+        images.append(i_image)
+
+    pixel_values = feature_extractor(images=images, return_tensors="pt").pixel_values
+    pixel_values = pixel_values.to(device)
+
+    output_ids = model.generate(pixel_values, **gen_kwargs)
+
+    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+    preds = [prepare(pred).strip() for pred in preds]
+    return preds
+
+img_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg' 
+result = predict_step([img_url]) # ['red shirt, chromatic aberration, light emitting object, barefoot, best quality, ocean background, 1girl, 8k wallpaper, intricate details, chromatic light, light, ocean, backpack, ultra-detailed, ocean light,masterpiece']
+print(result)
 ```
 
 ## Additional Information