Update symbol_detection.py

symbol_detection.py

@@ -1,11 +1,16 @@
-from transformers import DetrFeatureExtractor, DetrForObjectDetection
+from transformers import DetrImageProcessor, DetrForObjectDetection
+import torch
 
-extractor = DetrFeatureExtractor.from_pretrained("facebook/detr-resnet-50")
+# Use the new DetrImageProcessor
+processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
 model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
 
 def detect_elements(image):
     """Detects elements within the diagram."""
-    inputs = extractor(images=image, return_tensors="pt")
+    inputs = processor(images=image, return_tensors="pt")
     outputs = model(**inputs)
-    labels = outputs.logits.argmax(-1).numpy()
-    return labels
+
+    # Get the labels of the detected objects
+    logits = outputs.logits.argmax(-1).numpy()
+    return logits
+