p

app.py

@@ -21,6 +21,7 @@ from flask_cors import CORS
 import torch
 from groundingdino.util.inference import load_model, predict
 from segment_anything import sam_model_registry, SamPredictor
+import groundingdino.datasets.transforms as T
 
 # ─── Load models once ───────────────────────────────────────────────────────────
 device = torch.device("cpu")
@@ -39,20 +40,24 @@ app = Flask(__name__)
 CORS(app)
 
 def segment(image_pil: Image.Image, prompt: str):
-    # Convert PIL image to numpy array and normalize
-    image_np = np.array(image_pil).astype(np.float32) / 255.0  # Normalize to [0, 1]
-
-    # Convert numpy array to torch tensor
-    image_tensor = torch.tensor(image_np).permute(2, 0, 1).unsqueeze(0).to(device)  # Convert to CHW format
+    # Use the proper image preprocessing for GroundingDINO
+    transform = T.Compose([
+        T.RandomResize([800], max_size=1333),
+        T.ToTensor(),
+        T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    ])
+    
+    image_transformed, _ = transform(image_pil, None)
+    image_transformed = image_transformed.to(device)
 
     # Run GroundingDINO to get boxes for the prompt
     boxes, _, _ = predict(
         model=grounder,
-        image=image_tensor,  # Pass normalized tensor
+        image=image_transformed,
         caption=prompt,
         box_threshold=0.3,
         text_threshold=0.25,
-        device="cpu"  # Explicitly set device to CPU
+        device="cpu"
     )
     if boxes.size == 0:
         raise ValueError("No boxes found for prompt.")