Update modules/ai_model.py

modules/ai_model.py

@@ -1,19 +1,47 @@
-from transformers import DetrForObjectDetection, AutoFeatureExtractor
+from transformers import AutoImageProcessor, AutoModelForObjectDetection
 import torch
 from PIL import Image
+import logging
 
-# 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")
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-def run_inference(image):
-    # Preprocess the image using AutoFeatureExtractor
-    inputs = feature_extractor(images=image, return_tensors="pt")
-    
-    # Run inference
-    outputs = model(**inputs)
+try:
+    # Load the pre-trained DETR model and image processor
+    logger.info("Loading facebook/detr-resnet-50 model and processor...")
+    processor = AutoImageProcessor.from_pretrained("facebook/detr-resnet-50")
+    model = AutoModelForObjectDetection.from_pretrained("facebook/detr-resnet-50")
+    logger.info("Model and processor loaded successfully.")
+except Exception as e:
+    logger.error(f"Failed to load model or processor: {str(e)}")
+    raise
+
+def run_inference(image: Image.Image) -> dict:
+    """
+    Run object detection inference on the input image.
     
-    # Post-process the output (get bounding boxes)
-    results = feature_extractor.post_process_object_detection(outputs, target_sizes=[image.size[::-1]], threshold=0.9)[0]
+    Args:
+        image (PIL.Image.Image): Input image for object detection.
     
-    return results
+    Returns:
+        dict: Processed results containing bounding boxes, scores, and labels.
+    """
+    try:
+        # Preprocess the image using AutoImageProcessor
+        inputs = processor(images=image, return_tensors="pt")
+        
+        # Run inference
+        with torch.no_grad():  # Disable gradient calculation for inference
+            outputs = model(**inputs)
+        
+        # Post-process the output (get bounding boxes)
+        target_sizes = torch.tensor([image.size[::-1]])  # Format: [height, width]
+        results = processor.post_process_object_detection(
+            outputs, target_sizes=target_sizes, threshold=0.9
+        )[0]
+        
+        return results
+    except Exception as e:
+        logger.error(f"Error during inference: {str(e)}")
+        raise