Add python files

scripts/model_factory.py

@@ -0,0 +1,23 @@
+from inference_rfdetr import RFDETRInference
+
+
+def get_model(model_name, version, pretrain_weights):
+    """
+    Factory method to return the correct model inference class based on name and version.
+
+    Args:
+        model_name (str): Name of the model (e.g., 'rfdetr').
+        version (str): Version string (e.g., 'small', 'nano').
+        pretrain_weights (str): Path to model weights.
+
+    Returns:
+        BaseInference: A model inference object.
+
+    Raises:
+        ValueError: If model_name is unsupported.
+    """
+    if model_name == 'rfdetr':
+        return RFDETRInference(version, pretrain_weights)
+    else:
+        raise ValueError(f"Unsupported model: {model_name}")
+    

scripts/plot_bboxes_save_images_and_yolo_predictions.py

@@ -0,0 +1,131 @@
+import supervision as sv
+import cv2
+import os
+
+
+def convert_to_yolo(x1, y1, x2, y2, img_width, img_height):
+    """
+    Converts bounding box coordinates to YOLO format (relative center-x, center-y, width, height).
+
+    Returns:
+        tuple: (x_center, y_center, width, height)
+    """
+    x_center = (x1 + x2) / 2 / img_width
+    y_center = (y1 + y2) / 2 / img_height
+    width = (x2 - x1) / img_width
+    height = (y2 - y1) / img_height
+    return x_center, y_center, width, height
+
+
+def save_yolo_labels(save_path, detections, image):
+    """
+    Saves detection boxes in YOLO format to a .txt file.
+
+    Args:
+        save_path (str): Base path to match image name.
+        detections (sv.Detections): Detection results.
+        image (np.ndarray): Image to get original dimensions.
+    """
+    boxes = detections.xyxy
+    scores = detections.confidence
+    labels_ids = detections.class_id
+
+    img_height, img_width = image.shape[:2]
+    yolo_boxes = [convert_to_yolo(x1, y1, x2, y2, img_width, img_height) for x1, y1, x2, y2 in boxes]
+
+    base_name = os.path.splitext(os.path.basename(save_path))[0]
+    preds_dir = os.path.join(os.path.dirname(save_path), "../predictions")
+    os.makedirs(preds_dir, exist_ok=True)
+    txt_path = os.path.join(preds_dir, base_name + ".txt")
+
+    with open(txt_path, "w") as f:
+        for box, label, conf in zip(yolo_boxes, labels_ids, scores):
+            x_center, y_center, width, height = box
+            f.write(f"{label} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f} {conf:.6f}\n")
+
+    # print(f"[LABELS SAVED] {os.path.basename(txt_path)} → {os.path.abspath(txt_path)}")
+
+
+def annotate_image(image, detections, class_names):
+    """
+    Draws bounding boxes and class labels on the image.
+
+    Args:
+        image (np.ndarray): The image to annotate.
+        detections (sv.Detections): Detection results.
+        class_names (dict): Class ID to name mapping.
+
+    Returns:
+        np.ndarray: Annotated image.
+    """
+    if len(detections) > 0:
+        for i in range(len(detections)):
+            class_id = detections.class_id[i]
+            conf = detections.confidence[i]
+            label = f"{class_names[class_id]} {conf:.2f}"
+
+            # Draw box
+            box_annotator = sv.BoxAnnotator(thickness=6)
+            image = box_annotator.annotate(scene=image, detections=detections[i:i+1])
+
+            # Draw label
+            label_annotator = sv.LabelAnnotator(text_scale=2.0, text_thickness=4)
+            image = label_annotator.annotate(scene=image, detections=detections[i:i+1], labels=[label])
+
+    return image
+
+
+def process_image_frame(image, detections, class_names, save_path, plot_dets=True, save_preds=True, show=False):
+    """
+    Handles image output: annotation, saving image, saving labels, and optional GUI display.
+
+    Args:
+        image (np.ndarray): Input image.
+        detections (sv.Detections): Detection results.
+        class_names (dict): Class ID to name mapping.
+        save_path (str): Path to save image and labels.
+        plot_dets (bool): Save annotated image.
+        save_preds (bool): Save YOLO labels.
+        show (bool): Show OpenCV window.
+    """
+    annotated_image = annotate_image(image, detections, class_names)
+
+    if plot_dets:
+        cv2.imwrite(save_path, annotated_image)
+        # print(f"[IMAGE SAVED] {os.path.basename(save_path)} → {os.path.abspath(save_path)}")
+
+    if save_preds:
+        save_yolo_labels(save_path, detections, image)
+
+    if show:
+        h, w = image.shape[:2]
+        scale = min(750 / w, 750 / h)
+        resized = cv2.resize(annotated_image, (int(w * scale), int(h * scale)))
+        cv2.imshow("Detection", resized)
+        cv2.waitKey(1500)
+        cv2.destroyAllWindows()
+
+
+def process_video_frame(frame, detections, class_names, plot_dets=True, show=False, video_writer=None):
+    """
+    Handles each frame of video for annotation, display, and saving.
+
+    Args:
+        frame (np.ndarray): Video frame.
+        detections (sv.Detections): Detection results.
+        class_names (dict): Class ID to name mapping.
+        plot_dets (bool): Save annotated video frame.
+        show (bool): Display annotated frame in a window.
+        video_writer (cv2.VideoWriter): OpenCV video writer object.
+    """
+    annotated_frame = annotate_image(frame, detections, class_names)
+
+    if plot_dets and video_writer is not None:
+        video_writer.write(annotated_frame)
+
+    if show:
+        h, w = frame.shape[:2]
+        scale = min(750 / w, 750 / h)
+        resized = cv2.resize(annotated_frame, (int(w * scale), int(h * scale)))
+        cv2.imshow("Detection", resized)
+        cv2.waitKey(1)