Update README.md

README.md

@@ -17,10 +17,15 @@ image_processor = RTDetrImageProcessor.from_pretrained("HuggingPanda/docling-lay
 model = RTDetrForObjectDetection.from_pretrained("HuggingPanda/docling-layout")
 
 # Load an image
-image = Image.open("your-image.png")
+image = Image.open("hocr_output_page-0001.jpg")
 
 # Preprocess the image
-inputs = image_processor(images=image, return_tensors="pt")
+resize = {"height":640, "width":640}
+inputs = image_processor(
+    images=image,
+    return_tensors="pt",
+    size=resize,
+)
 
 # Perform inference
 with torch.no_grad():
@@ -29,7 +34,7 @@ with torch.no_grad():
 # Post-process results
 results = image_processor.post_process_object_detection(
     outputs, 
-    target_sizes=torch.tensor([(image.height, image.width)]), 
+    target_sizes=torch.tensor([image.size[::-1]]), 
     threshold=0.3
 )
 
@@ -38,7 +43,8 @@ for result in results:
     for score, label_id, box in zip(result["scores"], result["labels"], result["boxes"]):
         score, label = score.item(), label_id.item()
         box = [round(i, 2) for i in box.tolist()]
-        print(f"{model.config.id2label[label]}: {score:.2f} {box}")
+        print(f"{model.config.id2label[label+1]}: {score:.2f} {box}")
+
 ```