Update README.md

README.md

@@ -24,16 +24,19 @@ Here is how to use this model in PyTorch:
 ```python
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 from PIL import Image
-import requests
 
-# load image
-url = 'https://huggingface.co/dsupa/mangaocr-hoogberta-v1/to_test_hf_model/test1.jpg'
-image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
+processor = TrOCRProcessor.from_pretrained('dsupa/mangaocr-hoogberta-v2')
+model = VisionEncoderDecoderModel.from_pretrained('dsupa/mangaocr-hoogberta-v2')
 
-processor = TrOCRProcessor.from_pretrained('dsupa/mangaocr-hoogberta-v1')
-model = VisionEncoderDecoderModel.from_pretrained('dsupa/mangaocr-hoogberta-v1')
-pixel_values = processor(images=image, return_tensors="pt").pixel_values
+def predict(image_path):
+    image = Image.open(image_path).convert("RGB")
+    pixel_values = processor(images=image, return_tensors="pt").pixel_values
+    generated_ids = model.generate(pixel_values)
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-generated_ids = model.generate(pixel_values)
-generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    return generated_text
+
+image_path = "your_img.jpg"
+pred = predit(image_path)
+print(pred)
 ```