feat

app.py

@@ -243,8 +243,8 @@ def count_objects(image, pos_caption, neg_caption, box_threshold, point_radius,
     neg_queries = outputs["neg_queries"].squeeze(0).float()
     pos_queries = pos_queries[-1].squeeze(0)
     neg_queries = neg_queries[-1].squeeze(0)
-    pos_queries = pos_queries.cpu().numpy()
-    neg_queries = neg_queries.cpu().numpy()
+    pos_queries = pos_queries.cpu()
+    neg_queries = neg_queries.cpu()
     results = post_process_grounded_object_detection_with_queries(outputs, pos_queries, box_threshold=threshold)[0]
 
     boxes = results["boxes"]
@@ -266,6 +266,8 @@ def count_objects(image, pos_caption, neg_caption, box_threshold, point_radius,
     
     pos_queries = results["queries"]
     neg_queries = neg_results["queries"]
+    pos_queries = pos_queries.numpy()
+    neg_queries = neg_queries.numpy()
     
     img_size = image.size
     # filtered_points, kept_indices = filter_points_by_negative(
@@ -323,7 +325,6 @@ def create_demo():
         gr.Markdown("""
         # CountEx: Fine-Grained Counting via Exemplars and Exclusion
         Count specific objects in images using positive and negative text prompts.
-        **Important Note: Both the Positive and Negative prompts must end with a period (.) for the model to correctly interpret the instruction.**
         """)
 
         with gr.Row():
@@ -377,6 +378,7 @@ def create_demo():
                 ["examples/apples.png", "Green apple.", "Red apple."],
                 ["examples/black_beans.jpg", "Black bean.", "Soy bean."],
                 ["examples/candy.jpg", "Brown coffee candy.", "Black coffee candy."],
+                ["examples/strawberry.jpg", "strawberry.", "None."],
             ],
             inputs=[input_image, pos_caption, neg_caption],
             outputs=[output_image, count_output],

utils.py

@@ -55,10 +55,7 @@ def post_process_grounded_object_detection_with_queries(
     Now also returns the query embeddings for each kept prediction.
     """
     logits, boxes = outputs.logits, outputs.pred_boxes
-    print("logits: ", logits.shape)
-    print("boxes: ", boxes.shape)
-    print("queries: ", queries.shape)
-    assert len(logits[0]) == queries.shape[0], "logits and queries must have the same batch size, but got {} and {}".format(len(logits), queries.shape[0])
+    assert logits.shape == queries.shape, "logits and queries must have the same batch size, but got {} and {}".format(logits.shape[0], queries.shape[0])
 
     probs = torch.sigmoid(logits)  # (batch_size, num_queries, 256)
     scores = torch.max(probs, dim=-1)[0]  # (batch_size, num_queries)
@@ -69,15 +66,11 @@ def post_process_grounded_object_detection_with_queries(
         score = s[mask]
         box = b[mask]
         prob = p[mask]
-        
-        result = {"scores": score, "boxes": box}
-        
-        # 保存对应的 query embeddings
-        if queries is not None:
-            result["queries"] = queries[idx][mask]  # (num_kept, D)
-        
+        queries = queries[mask]
+        result = {"scores": score, "boxes": box, "queries": queries}
+        print('scores: ', score.shape, 'boxes: ', box.shape, 'queries: ', queries.shape)
         results.append(result)
-    assert len(results['scores']) == len(results['boxes']) == results['queries'].shape[0], "scores, boxes and queries must have the same length"
+    assert results['scores'].shape == results['boxes'].shape == results['queries'].shape, "scores, boxes and queries must have the same shape, but got {} and {}".format(results['scores'].shape, results['boxes'].shape, results['queries'].shape)
     return results