Views
No views yet
1import requests
2from PIL import Image
3
4from transformers import AutoProcessor, OmDetTurboForObjectDetection
5
6processor = AutoProcessor.from_pretrained("omlab/omdet-turbo-swin-tiny-hf")
7model = OmDetTurboForObjectDetection.from_pretrained("omlab/omdet-turbo-swin-tiny-hf")
8
9url = "http://images.cocodataset.org/val2017/000000039769.jpg"
10image = Image.open(requests.get(url, stream=True).raw)
11classes = ["cat", "remote"]
12inputs = processor(image, text=classes, return_tensors="pt")
13
14outputs = model(**inputs)
15
16# convert outputs (bounding boxes and class logits)
17results = processor.post_process_grounded_object_detection(
18 outputs,
19 classes=classes,
20 target_sizes=[image.size[::-1]],
21 score_threshold=0.3,
22 nms_threshold=0.3,
23)[0]
24for score, class_name, box in zip(
25 results["scores"], results["classes"], results["boxes"]
26):
27 box = [round(i, 1) for i in box.tolist()]
28 print(
29 f"Detected {class_name} with confidence "
30 f"{round(score.item(), 2)} at location {box}"
31 )1>>> import torch
2>>> import requests
3>>> from io import BytesIO
4>>> from PIL import Image
5>>> from transformers import AutoProcessor, OmDetTurboForObjectDetection
6
7>>> processor = AutoProcessor.from_pretrained("omlab/omdet-turbo-swin-tiny-hf")
8>>> model = OmDetTurboForObjectDetection.from_pretrained("omlab/omdet-turbo-swin-tiny-hf")
9
10>>> url1 = "http://images.cocodataset.org/val2017/000000039769.jpg"
11>>> image1 = Image.open(BytesIO(requests.get(url1).content)).convert("RGB")
12>>> classes1 = ["cat", "remote"]
13>>> task1 = "Detect {}.".format(", ".join(classes1))
14
15>>> url2 = "http://images.cocodataset.org/train2017/000000257813.jpg"
16>>> image2 = Image.open(BytesIO(requests.get(url2).content)).convert("RGB")
17>>> classes2 = ["boat"]
18>>> task2 = "Detect everything that looks like a boat."
19
20>>> url3 = "https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg"
21>>> image3 = Image.open(BytesIO(requests.get(url3).content)).convert("RGB")
22>>> classes3 = ["statue", "trees"]
23>>> task3 = "Focus on the foreground, detect statue and trees."
24
25>>> inputs = processor(
26... images=[image1, image2, image3],
27... text=[classes1, classes2, classes3],
28... task=[task1, task2, task3],
29... return_tensors="pt",
30... )
31
32>>> with torch.no_grad():
33... outputs = model(**inputs)
34
35>>> # convert outputs (bounding boxes and class logits)
36>>> results = processor.post_process_grounded_object_detection(
37... outputs,
38... classes=[classes1, classes2, classes3],
39... target_sizes=[image1.size[::-1], image2.size[::-1], image3.size[::-1]],
40... score_threshold=0.2,
41... nms_threshold=0.3,
42... )
43
44>>> for i, result in enumerate(results):
45... for score, class_name, box in zip(
46... result["scores"], result["classes"], result["boxes"]
47... ):
48... box = [round(i, 1) for i in box.tolist()]
49... print(
50... f"Detected {class_name} with confidence "
51... f"{round(score.item(), 2)} at location {box} in image {i}"
52... )
53Detected remote with confidence 0.77 at location [39.9, 70.4, 176.7, 118.0] in image 0
54Detected cat with confidence 0.72 at location [11.6, 54.2, 314.8, 474.0] in image 0
55Detected remote with confidence 0.56 at location [333.4, 75.8, 370.7, 187.0] in image 0
56Detected cat with confidence 0.55 at location [345.2, 24.0, 639.8, 371.7] in image 0
57Detected boat with confidence 0.32 at location [146.9, 219.8, 209.6, 250.7] in image 1
58Detected boat with confidence 0.3 at location [319.1, 223.2, 403.2, 238.4] in image 1
59Detected boat with confidence 0.27 at location [37.7, 220.3, 84.0, 235.9] in image 1
60Detected boat with confidence 0.22 at location [407.9, 207.0, 441.7, 220.2] in image 1
61Detected statue with confidence 0.73 at location [544.7, 210.2, 651.9, 502.8] in image 2
62Detected trees with confidence 0.25 at location [3.9, 584.3, 391.4, 785.6] in image 2
63Detected trees with confidence 0.25 at location [1.4, 621.2, 118.2, 787.8] in image 2
64Detected statue with confidence 0.2 at location [428.1, 205.5, 767.3, 759.5] in image 2
65