update models

model.py

@@ -1,31 +1,23 @@
-from transformers import AutoConfig, AutoModelForObjectDetection
-from PIL import Image
+from transformers import AutoFeatureExtractor, AutoModelForObjectDetection
 import torch
+from config import MODEL_NAME
 
 
 class RadarDetectionModel:
     def __init__(self):
-        self.config = AutoConfig.from_pretrained(
-            "Extremely4606/paligemma_9_19")
+        self.feature_extractor = AutoFeatureExtractor.from_pretrained(
+            "google/pali-gamma-336m")
         self.model = AutoModelForObjectDetection.from_pretrained(
-            "Extremely4606/paligemma_9_19")
+            "google/pali-gamma-336m")
         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.preprocess_image(image)
-        outputs = self.model(inputs)
+        inputs = self.feature_extractor(images=image, return_tensors="pt")
+        outputs = self.model(**inputs)
 
-        # 这里可能需要根据模型的输出格式进行调整
-        boxes = outputs.pred_boxes[0]
-        scores = outputs.scores[0]
-        labels = outputs.labels[0]
+        target_sizes = torch.tensor([image.size[::-1]])
+        results = self.feature_extractor.post_process_object_detection(
+            outputs, threshold=0.5, target_sizes=target_sizes)[0]
 
-        return {"boxes": boxes, "scores": scores, "labels": labels}
+        return results