Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,219 Bytes
ced4fcf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Batch Inference example using Rex-Omni
"""
import torch
from PIL import Image
from rex_omni import RexOmniVisualize, RexOmniWrapper
def main():
# Model path - replace with your actual model path
model_path = "IDEA-Research/Rex-Omni"
# Create wrapper with custom parameters
rex_model = RexOmniWrapper(
model_path=model_path,
backend="transformers", # or "vllm" for faster inference
max_tokens=4096,
temperature=0.0,
top_p=0.05,
top_k=1,
repetition_penalty=1.05,
)
# Load imag
image_paths = [
"tutorials/detection_example/test_images/cafe.jpg",
"tutorials/detection_example/test_images/boys.jpg",
]
images = [Image.open(image_path).convert("RGB") for image_path in image_paths]
# Object detection
categories = [
[
"man",
"woman",
"yellow flower",
"sofa",
"robot-shope light",
"blanket",
"microwave",
"laptop",
"cup",
"white chair",
"lamp",
],
[
"boys holding microphone",
"boy playing piano",
"the four guitars on the wall",
"the guitar in someone's hand",
],
]
results = rex_model.inference(
images=images, task=["detection", "detection"], categories=categories
)
# Print results
batch_idx = 0
for result, image in zip(results, images):
if result["success"]:
predictions = result["extracted_predictions"]
vis_image = RexOmniVisualize(
image=image,
predictions=predictions,
font_size=20,
draw_width=5,
show_labels=True,
)
# Save visualization
output_path = f"tutorials/other_example/batch_inference_{batch_idx}.jpg"
vis_image.save(output_path)
print(f"Visualization saved to: {output_path}")
else:
print(f"Inference failed: {result['error']}")
batch_idx += 1
if __name__ == "__main__":
main()
|