Update image captions

config/rf-detr.yml

@@ -25,7 +25,7 @@ tracker:
 
 db:
   db_uri: sqlite:///data/hainan.db
-  use_gpt4v_captions: True
+  create_image_captions: True
 
 names:
   1: "person"

config/yolo.yml

@@ -20,7 +20,7 @@ tracker:
 
 db:
   db_uri: sqlite:///data/nba_Trim.db
-  use_gpt4v_captions: True
+  create_image_captions: True
 
 
 names: 

fast_track/databases/schemas/frame.py

@@ -14,6 +14,6 @@ class Frame(Base):
     frame_number = Column(Integer, nullable=False)
     time_created = Column(String, nullable=False)
     frame_base64 = Column(String, nullable=False)
-    gpt4v_caption = Column(String, nullable=True)
+    image_caption = Column(String, nullable=True)
     job_id = Column(Integer, ForeignKey("jobs.job_id"))
     job: Mapped["Job"] = relationship(back_populates="frames")

fast_track/databases/sql_database.py

@@ -18,16 +18,16 @@
 class SQLDatabase:
     """Database class to store information about tracks and detections."""
 
-    def __init__(self, db_uri: str, class_names: List[str], use_gpt4v_captions: bool = False):
+    def __init__(self, db_uri: str, class_names: List[str], create_image_captions: bool = False):
         """Inits Database class with a given database URI.
 
         Args:
             db_uri: database URI.
-            use_gpt4v_captions: bool to use GPT-4 Vision captions.
+            create_image_captions: bool to use VLM to generate image captions (OpenAI).
         """
         self.db_uri = db_uri
         self.class_names = class_names
-        self.use_gpt4v_captions = use_gpt4v_captions
+        self.create_image_captions = create_image_captions
 
         self.db = None
         self.job_id = None
@@ -96,7 +96,7 @@ def add_frame(self, frame: np.ndarray, frame_number: int) -> None:
                 frame_number=frame_number,
                 frame_base64=frame_base64,
                 time_created=datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-                gpt4v_caption=generate_frame_caption(frame_base64) if self.use_gpt4v_captions else None,
+                image_caption=generate_frame_caption(frame_base64) if self.create_image_captions else None,
                 job_id=self.job_id,
             )
         )

fast_track/databases/utils.py

@@ -28,8 +28,8 @@ def encode_image(image: np.ndarray) -> str:
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
 
-def generate_frame_caption(frame_base64: str, model: str = "gpt-4-vision-preview") -> np.ndarray:
-    """Generates a caption for a frame using GPT-4 Vision.
+def generate_frame_caption(frame_base64: str, model: str = "gpt-4o-mini") -> np.ndarray:
+    """Generates a caption for a frame using OpenAI.
 
     Args:
         frame: frame to generate a caption for

fast_track/detectors/util.py

@@ -11,6 +11,8 @@
 
 
 MODELS = {
+    "RF-DETR Large": "rfdetr_large",
+    "RF-DETR Base": "rfdetr_base",
     "YOLO-NAS L": "yolo_nas_l",
     "YOLO-NAS M": "yolo_nas_m",
     "YOLO-NAS S": "yolo_nas_s",
@@ -56,6 +58,6 @@ def get_detector(
     elif detector_type.startswith("yolov9"):
         return YOLOv9ONNX(weights_path=weights_path, names=names, image_shape=image_shape, **detector_params)
     elif detector_type.startswith("rfdetr"):
-        return RFDETR(weights_path=weights_path, names=names, image_shape=image_shape, **detector_params)
+        return RFDETR(weights_path=weights_path, model_name=detector_type, names=names, resolution=max(image_shape), **detector_params)
     else:
         raise ValueError("Detector name not found.")

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fast_track"
-version = "1.2.0"
+version = "1.2.1"
 description = "Object detection and tracking pipeline"
 readme = "README.md"
 keywords = [