Views
No views yet
1import requests
2from PIL import Image
3from transformers import AutoImageProcessor, AutoModelForObjectDetection
4
5model_path = "PaddlePaddle/PP-DocLayoutV2_safetensors"
6model = AutoModelForObjectDetection.from_pretrained(model_path)
7image_processor = AutoImageProcessor.from_pretrained(model_path)
8image = Image.open(requests.get("https://paddle-model-ecology.bj.bcebos.com/paddlex/imgs/demo_image/layout_demo.jpg", stream=True).raw)
9inputs = image_processor(images=image, return_tensors="pt")
10
11outputs = model(**inputs)
12results = image_processor.post_process_object_detection(outputs, target_sizes=[image.size[::-1]])
13for result in results:
14 print(result["scores"])
15 print(result["labels"])
16 print(result["boxes"])
17 for idx, (score, label_id, box) in enumerate(zip(result["scores"], result["labels"], result["boxes"])):
18 score, label = score.item(), label_id.item()
19 box = [round(i, 2) for i in box.tolist()]
20 print(f"Order {idx + 1}: {model.config.id2label[label]}: {score:.2f} {box}")