update models

model.py

@@ -1,22 +1,31 @@
-from transformers import AutoFeatureExtractor, AutoModelForObjectDetection
+from transformers import AutoConfig, AutoModelForObjectDetection
+from PIL import Image
 import torch
-from config import MODEL_NAME
 
 
 class RadarDetectionModel:
     def __init__(self):
-        self.feature_extractor = AutoFeatureExtractor.from_pretrained(
-            MODEL_NAME)
-        self.model = AutoModelForObjectDetection.from_pretrained(MODEL_NAME)
+        self.config = AutoConfig.from_pretrained(
+            "Extremely4606/paligemma_9_19")
+        self.model = AutoModelForObjectDetection.from_pretrained(
+            "Extremely4606/paligemma_9_19")
         self.model.eval()
 
+    def preprocess_image(self, image):
+        # 这里需要根据模型的具体要求来处理图像
+        # 这只是一个示例，可能需要调整
+        image = image.resize((224, 224))
+        image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.0
+        return image.unsqueeze(0)
+
     @torch.no_grad()
     def detect(self, image):
-        inputs = self.feature_extractor(images=image, return_tensors="pt")
-        outputs = self.model(**inputs)
+        inputs = self.preprocess_image(image)
+        outputs = self.model(inputs)
 
-        target_sizes = torch.tensor([image.size[::-1]])
-        results = self.feature_extractor.post_process_object_detection(
-            outputs, threshold=0.5, target_sizes=target_sizes)[0]
+        # 这里可能需要根据模型的输出格式进行调整
+        boxes = outputs.pred_boxes[0]
+        scores = outputs.scores[0]
+        labels = outputs.labels[0]
 
-        return results
+        return {"boxes": boxes, "scores": scores, "labels": labels}