Update modules/ai_model.py

modules/ai_model.py

@@ -1,15 +1,19 @@
-from transformers import DetrImageProcessor, DetrForObjectDetection
+from transformers import DetrForObjectDetection, AutoFeatureExtractor
 import torch
 from PIL import Image
 
-# Load pre-trained DETR model
-processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
+# Load the pre-trained DETR model and feature extractor
 model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
+feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/detr-resnet-50")
 
-def run_inference(image_path):
-    image = Image.open(image_path)  # Replace with actual image path
-    inputs = processor(images=image, return_tensors="pt")
+def run_inference(image):
+    # Preprocess the image using AutoFeatureExtractor
+    inputs = feature_extractor(images=image, return_tensors="pt")
     
     # Run inference
     outputs = model(**inputs)
-    return outputs
+    
+    # Post-process the output (get bounding boxes)
+    results = feature_extractor.post_process_object_detection(outputs, target_sizes=[image.size[::-1]], threshold=0.9)[0]
+    
+    return results