support hf

README.md

@@ -12,59 +12,26 @@ Kosmos-2.5 is a multimodal literate model for machine reading of text-intensive
 
 [Kosmos-2.5: A Multimodal Literate Model](https://arxiv.org/abs/2309.11419)
 
-## NOTE:
+## NOTE
 Since this is a generative model, there is a risk of **hallucination** during the generation process, and it **CAN NOT** guarantee the accuracy of all OCR/Markdown results in the images.
 
-## Use with transformers：
-```python
-from PIL import Image
-import requests
-import torch
-from transformers import AutoProcessor, Kosmos2_5ForConditionalGeneration
-import re
-repo = "microsoft/kosmos-2.5"
-device = "cuda:0"
-dtype = torch.bfloat16
-model = Kosmos2_5ForConditionalGeneration.from_pretrained(repo, device_map=device, torch_dtype=dtype)
-processor = AutoProcessor.from_pretrained(repo)
-url = "https://huggingface.co/kirp/kosmos2_5/resolve/main/receipt_00008.png"
-image = Image.open(requests.get(url, stream=True).raw)
-prompt = "<ocr>" # <md>
-inputs = processor(text=prompt, images=image, return_tensors="pt")
-height, width = inputs.pop("height"), inputs.pop("width")
-raw_width, raw_height = image.size
-scale_height = raw_height / height
-scale_width = raw_width / width
-inputs = {k: v.to(device) if v is not None else None for k, v in inputs.items()}
-inputs["flattened_patches"] = inputs["flattened_patches"].to(dtype)
-generated_ids = model.generate(
-    **inputs,
-    max_new_tokens=1024,
-)
-generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
-def postprocess(y, scale_height, scale_width):
-    y = y.replace(prompt, "")
-    if "<md>" in prompt:
-        return y
-    pattern = r"<bbox><x_\d+><y_\d+><x_\d+><y_\d+></bbox>"
-    bboxs_raw = re.findall(pattern, y)
-    lines = re.split(pattern, y)[1:]
-    bboxs = [re.findall(r"\d+", i) for i in bboxs_raw]
-    bboxs = [[int(j) for j in i] for i in bboxs]
-    info = ""
-    for i in range(len(lines)):
-        box = bboxs[i]
-        x0, y0, x1, y1 = box
-        if not (x0 >= x1 or y0 >= y1):
-            x0 = int(x0 * scale_width)
-            y0 = int(y0 * scale_height)
-            x1 = int(x1 * scale_width)
-            y1 = int(y1 * scale_height)
-            info += f"{x0},{y0},{x1},{y0},{x1},{y1},{x0},{y1},{lines[i]}"
-    return info
-output_text = postprocess(generated_text[0], scale_height, scale_width)
-print(output_text)
+## Usage
+### Markdown Task
+Run with [md.py](md.py).
+```text
+- **1 \[REG\] BLACK SAKURA** 45,455
+- **1 COOKIE DOH SAUCES** 0
+- **1 NATA DE COCO** 0
+- **Sub Total** 45,455
+- **PB1 (10%)** 4,545
+- **Rounding** 0
+- **Total** **50,000**
+
+Card Payment 50,000
 ```
+
+### OCR Task
+Run with [ocr.py](ocr.py).
 ```text
 55,595,71,595,71,629,55,629,1
 82,595,481,595,481,635,82,635,[REG] BLACK SAKURA
@@ -81,7 +48,7 @@ print(output_text)
 24,905,858,905,858,956,24,956,Total 50,000
 17,1096,868,1096,868,1150,17,1150,Card Payment 50,000
 ```
-
+![output](output.png)
 
 
 ## Citation
@@ -103,3 +70,4 @@ The content of this project itself is licensed under the [MIT](https://github.co
 [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct)
 
 
+

config.json

@@ -148,4 +148,4 @@
     "typical_p": 1.0,
     "use_bfloat16": false
   }
-}
+}

md.py

@@ -0,0 +1,33 @@
+import re
+import torch
+import requests
+from PIL import Image, ImageDraw
+from transformers import AutoProcessor, Kosmos2_5ForConditionalGeneration
+
+repo = "microsoft/kosmos-2.5"
+device = "cuda:0"
+dtype = torch.bfloat16
+model = Kosmos2_5ForConditionalGeneration.from_pretrained(repo, device_map=device, torch_dtype=dtype)
+processor = AutoProcessor.from_pretrained(repo)
+
+# sample image
+url = "https://huggingface.co/microsoft/kosmos-2.5/blob/main/receipt_00008.png"
+image = Image.open(requests.get(url, stream=True).raw)
+
+prompt = "<md>"
+inputs = processor(text=prompt, images=image, return_tensors="pt")
+
+height, width = inputs.pop("height"), inputs.pop("width")
+raw_width, raw_height = image.size
+scale_height = raw_height / height
+scale_width = raw_width / width
+
+inputs = {k: v.to(device) if v is not None else None for k, v in inputs.items()}
+inputs["flattened_patches"] = inputs["flattened_patches"].to(dtype)
+generated_ids = model.generate(
+    **inputs,
+    max_new_tokens=1024,
+)
+
+generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
+print(generated_text[0])

ocr.py

@@ -0,0 +1,73 @@
+import re
+import torch
+import requests
+from PIL import Image, ImageDraw
+from transformers import AutoProcessor, Kosmos2_5ForConditionalGeneration
+
+repo = "microsoft/kosmos-2.5"
+device = "cuda:0"
+dtype = torch.bfloat16
+model = Kosmos2_5ForConditionalGeneration.from_pretrained(repo, device_map=device, torch_dtype=dtype)
+processor = AutoProcessor.from_pretrained(repo)
+
+# sample image
+url = "https://huggingface.co/microsoft/kosmos-2.5/blob/main/receipt_00008.png"
+image = Image.open(requests.get(url, stream=True).raw)
+
+# bs = 1
+prompt = "<ocr>"
+inputs = processor(text=prompt, images=image, return_tensors="pt")
+height, width = inputs.pop("height"), inputs.pop("width")
+raw_width, raw_height = image.size
+scale_height = raw_height / height
+scale_width = raw_width / width
+
+# bs > 1, batch generation
+# inputs = processor(text=[prompt, prompt], images=[image,image], return_tensors="pt")
+# height, width = inputs.pop("height"), inputs.pop("width")
+# raw_width, raw_height = image.size
+# scale_height = raw_height / height[0]
+# scale_width = raw_width / width[0]
+
+inputs = {k: v.to(device) if v is not None else None for k, v in inputs.items()}
+inputs["flattened_patches"] = inputs["flattened_patches"].to(dtype)
+generated_ids = model.generate(
+    **inputs,
+    max_new_tokens=1024,
+)
+
+generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
+def post_process(y, scale_height, scale_width):
+    y = y.replace(prompt, "")
+    if "<md>" in prompt:
+        return y
+    pattern = r"<bbox><x_\d+><y_\d+><x_\d+><y_\d+></bbox>"
+    bboxs_raw = re.findall(pattern, y)
+    lines = re.split(pattern, y)[1:]
+    bboxs = [re.findall(r"\d+", i) for i in bboxs_raw]
+    bboxs = [[int(j) for j in i] for i in bboxs]
+    info = ""
+    for i in range(len(lines)):
+        box = bboxs[i]
+        x0, y0, x1, y1 = box
+        if not (x0 >= x1 or y0 >= y1):
+            x0 = int(x0 * scale_width)
+            y0 = int(y0 * scale_height)
+            x1 = int(x1 * scale_width)
+            y1 = int(y1 * scale_height)
+            info += f"{x0},{y0},{x1},{y0},{x1},{y1},{x0},{y1},{lines[i]}"
+    return info
+
+output_text = post_process(generated_text[0], scale_height, scale_width)
+print(output_text)
+
+draw = ImageDraw.Draw(image)
+lines = output_text.split("\n")
+for line in lines:
+    # draw the bounding box
+    line = list(line.split(","))
+    if len(line) < 8:
+        continue
+    line = list(map(int, line[:8]))
+    draw.polygon(line, outline="red")
+image.save("output.png")

tokenizer.json

@@ -91,7 +91,7 @@
       "lstrip": true,
       "rstrip": false,
       "normalized": false,
-      "special": false
+      "special": true
     },
     {
       "id": 100283,
@@ -145,7 +145,7 @@
       "lstrip": true,
       "rstrip": false,
       "normalized": false,
-      "special": false
+      "special": true
     },
     {
       "id": 100289,