wirthual commited on
Commit
362c76c
·
verified ·
1 Parent(s): f60ba78

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +67 -3
README.md CHANGED
@@ -1,3 +1,67 @@
1
- ---
2
- license: apache-2.0
3
- ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ ---
4
+
5
+
6
+ For details about the models, please see: https://github.com/roboflow/rf-detr
7
+
8
+
9
+ The models have been exported to executorch without lowering.
10
+
11
+ To run:
12
+
13
+ ```python
14
+
15
+ from PIL import Image, ImageDraw
16
+
17
+ from executorch.runtime import Runtime
18
+ import torch
19
+ import torch.nn.functional as F
20
+ from torchvision import transforms
21
+
22
+
23
+ IMG_SIZE = (512, 512)
24
+ # change to (384, 384) for RFDETRNano
25
+ # change to (512, 512) for RFDETRSmall
26
+ # change to (576, 576) for RFDETRMedium
27
+ # change to (704, 704) for RFDETRLarge
28
+
29
+
30
+ def visualize_output(image, output):
31
+ draw = ImageDraw.Draw(image)
32
+
33
+ for box,logits in zip(output[0][0], output[1][0]):
34
+ probs = F.softmax(logits, dim=0)
35
+ pred_class = torch.argmax(probs, dim=0)
36
+
37
+ if probs[pred_class] > 0.7: # only draw if confidence is greater than 0.7
38
+ cx, cy, w, h = box
39
+ x1 = int((cx - w / 2) * img.width)
40
+ y1 = int((cy - h / 2) * img.height)
41
+ x2 = int((cx + w / 2) * img.width)
42
+ y2 = int((cy + h / 2) * img.height)
43
+
44
+ draw.rectangle([(x1, y1), (x2, y2)], fill=None, outline="black", width=3)
45
+
46
+ img = Image.open("./cats_coco.jpg").convert("RGB")
47
+ transform = transforms.Compose([
48
+ transforms.Resize(IMG_SIZE),
49
+ transforms.ToTensor(),
50
+ ])
51
+
52
+ tensor = transform(img)
53
+ tensor = tensor.unsqueeze(0)
54
+
55
+
56
+ runtime = Runtime.get()
57
+ method = runtime.load_program("model_small.pte").load_method("forward")
58
+ outputs = method.execute([tensor])
59
+
60
+ visualize_output(img, outputs)
61
+ img.save("output.png")
62
+ img.show()
63
+ ```
64
+
65
+ Example output:
66
+ ![Example output](./output.png)
67
+