Update app.py

app.py

@@ -1,16 +1,45 @@
-from transformers import ViTImageProcessor, ViTForImageClassification
-from PIL import Image
 import requests
 
-url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
+from PIL import Image
+from transformers import AutoProcessor, AutoModelForVision2Seq
+
+
+model = AutoModelForVision2Seq.from_pretrained("microsoft/kosmos-2-patch14-224")
+processor = AutoProcessor.from_pretrained("microsoft/kosmos-2-patch14-224")
+
+prompt = "<grounding>An image of"
+
+url = "https://huggingface.co/microsoft/kosmos-2-patch14-224/resolve/main/snowman.png"
 image = Image.open(requests.get(url, stream=True).raw)
 
-processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')
-model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
+# The original Kosmos-2 demo saves the image first then reload it. For some images, this will give slightly different image input and change the generation outputs.
+image.save("new_image.jpg")
+image = Image.open("new_image.jpg")
+
+inputs = processor(text=prompt, images=image, return_tensors="pt")
+
+generated_ids = model.generate(
+    pixel_values=inputs["pixel_values"],
+    input_ids=inputs["input_ids"],
+    attention_mask=inputs["attention_mask"],
+    image_embeds=None,
+    image_embeds_position_mask=inputs["image_embeds_position_mask"],
+    use_cache=True,
+    max_new_tokens=128,
+)
+generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+# Specify `cleanup_and_extract=False` in order to see the raw model generation.
+processed_text = processor.post_process_generation(generated_text, cleanup_and_extract=False)
+
+print(processed_text)
+# `<grounding> An image of<phrase> a snowman</phrase><object><patch_index_0044><patch_index_0863></object> warming himself by<phrase> a fire</phrase><object><patch_index_0005><patch_index_0911></object>.`
+
+# By default, the generated  text is cleanup and the entities are extracted.
+processed_text, entities = processor.post_process_generation(generated_text)
+
+print(processed_text)
+# `An image of a snowman warming himself by a fire.`
 
-inputs = processor(images=image, return_tensors="pt")
-outputs = model(**inputs)
-logits = outputs.logits
-# model predicts one of the 1000 ImageNet classes
-predicted_class_idx = logits.argmax(-1).item()
-print("Predicted class:", model.config.id2label[predicted_class_idx])
+print(entities)
+# `[('a snowman', (12, 21), [(0.390625, 0.046875, 0.984375, 0.828125)]), ('a fire', (41, 47), [(0.171875, 0.015625, 0.484375, 0.890625)])]`