|
from PIL import Image |
|
import torch |
|
import numpy as np |
|
import onnxruntime |
|
import os |
|
|
|
class FasterRCNNInference: |
|
def __init__(self, model_path): |
|
|
|
self.ort_session = onnxruntime.InferenceSession(f"{model_path}/model.onnx") |
|
|
|
|
|
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): |
|
|
|
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() |
|
|
|
|
|
outputs = self.ort_session.run(None, {"pixel_values": pixel_values}) |
|
|
|
|
|
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 |
|
|