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