| --- |
| license: apache-2.0 |
| --- |
| |
|
|
| For details about the models, please see: https://github.com/roboflow/rf-detr |
|
|
|
|
| The models have been exported to executorch without lowering. |
|
|
| To run: |
|
|
| ```python |
| |
| from PIL import Image, ImageDraw |
| |
| from executorch.runtime import Runtime |
| import torch |
| import torch.nn.functional as F |
| from torchvision import transforms |
| |
| |
| IMG_SIZE = (512, 512) |
| # change to (384, 384) for RFDETRNano |
| # change to (512, 512) for RFDETRSmall |
| # change to (576, 576) for RFDETRMedium |
| # change to (704, 704) for RFDETRLarge |
| |
| |
| def visualize_output(image, output): |
| draw = ImageDraw.Draw(image) |
| |
| for box,logits in zip(output[0][0], output[1][0]): |
| probs = F.softmax(logits, dim=0) |
| pred_class = torch.argmax(probs, dim=0) |
| |
| if probs[pred_class] > 0.7: # only draw if confidence is greater than 0.7 |
| cx, cy, w, h = box |
| x1 = int((cx - w / 2) * img.width) |
| y1 = int((cy - h / 2) * img.height) |
| x2 = int((cx + w / 2) * img.width) |
| y2 = int((cy + h / 2) * img.height) |
| |
| draw.rectangle([(x1, y1), (x2, y2)], fill=None, outline="black", width=3) |
| |
| img = Image.open("./cats_coco.jpg").convert("RGB") |
| transform = transforms.Compose([ |
| transforms.Resize(IMG_SIZE), |
| transforms.ToTensor(), |
| ]) |
| |
| tensor = transform(img) |
| tensor = tensor.unsqueeze(0) |
| |
| |
| runtime = Runtime.get() |
| method = runtime.load_program("model_small.pte").load_method("forward") |
| outputs = method.execute([tensor]) |
| |
| visualize_output(img, outputs) |
| img.save("output.png") |
| img.show() |
| ``` |
|
|
| Example output: |
|  |
|
|
|
|
