myfasterrcnn-detector / inference.py
lkk688's picture
Upload folder using huggingface_hub
e8e4cb7 verified
from PIL import Image
import torch
import numpy as np
import onnxruntime
import os
class FasterRCNNInference:
def __init__(self, model_path):
# Load ONNX model
self.ort_session = onnxruntime.InferenceSession(f"{model_path}/model.onnx")
# Try to load Hugging Face image processor
try:
from transformers import AutoImageProcessor
self.image_processor = AutoImageProcessor.from_pretrained(model_path)
self.use_hf_processor = True
print("Using Hugging Face image processor")
except Exception as e:
print(f"Could not load Hugging Face image processor: {e}")
print("Falling back to custom processor")
self.processor = torch.load(f"{model_path}/processor.bin")
self.use_hf_processor = False
def predict(self, images, threshold=0.5):
# Preprocess
if self.use_hf_processor:
inputs = self.image_processor(images=images, return_tensors="pt")
pixel_values = inputs["pixel_values"].numpy()
else:
inputs = self.processor(images)
pixel_values = inputs["pixel_values"].numpy()
# Run inference
outputs = self.ort_session.run(None, {"pixel_values": pixel_values})
# Post-process
results = []
for boxes, scores, labels in zip(outputs[0], outputs[1], outputs[2]):
keep = scores > threshold
results.append({
"boxes": boxes[keep],
"scores": scores[keep],
"labels": labels[keep]
})
return results