I am using the CLI version of the yolo model yolov8l.pt (accessed through the WEIGHTS parameter below):

!yolo detect train model={WEIGHTS} data='data/tvt3_data_v8.yaml' single_cls imgsz={IMG_SIZE} batch={BATCH_SIZE} epochs={3}

How it currently works

During training, at each epoch:

• The training loss is calculated
• The validation loss and validation Recall, Precision and mAP are calculated

At the end of the training, the best.pt model is selected and evaluated on the validation set.

What I need from the script

At each epoch, I want to extract the predictions made on the VALIDATION dataset

How do I do this?

I am using the CLI version of the yolo model yolov8l.pt (accessed through the WEIGHTS parameter below):

!yolo detect train model={WEIGHTS} data='data/tvt3_data_v8.yaml' single_cls imgsz={IMG_SIZE} batch={BATCH_SIZE} epochs={3}

How it currently works

During training, at each epoch:

• The training loss is calculated
• The validation loss and validation Recall, Precision and mAP are calculated

At the end of the training, the best.pt model is selected and evaluated on the validation set.

What I need from the script

At each epoch, I want to extract the predictions made on the VALIDATION dataset

How do I do this?

• python
• machine-learning
• pytorch
• object-detection
• yolo

1 Answer 1

I am not sure that you have checked or now, but this script is easily available, here is a script that I made for my work and modified according to your need:

from ultralytics import YOLO
import torch
from pathlib import Path
import json
from datetime import datetime

class ValidationPredictionCallback:
    def __init__(self, save_dir='validation_predictions'):
        self.save_dir = Path(save_dir)
        self.save_dir.mkdir(exist_ok=True, parents=True)
    
    def save_predictions(self, validator, epoch):
        predictions = []
        
        # Iterate through validation set predictions
        for batch_idx, (batch_images, batch_outputs) in enumerate(validator.dataloader):
            # Get predictions for this batch
            results = validator.model(batch_images)
            
            for img_idx, result in enumerate(results):
                # Get image path
                img_path = validator.dataloader.dataset.im_files[batch_idx * validator.dataloader.batch_size + img_idx]
                
                # Extract boxes, scores, and classes
                boxes = result.boxes.xyxy.cpu().numpy()
                scores = result.boxes.conf.cpu().numpy()
                classes = result.boxes.cls.cpu().numpy()
                
                # Store predictions for this image
                img_predictions = {
                    'image_path': img_path,
                    'predictions': [{
                        'bbox': box.tolist(),
                        'confidence': float(score),
                        'class': int(cls)
                    } for box, score, cls in zip(boxes, scores, classes)]
                }
                predictions.append(img_predictions)
        
        # Save predictions for this epoch
        output_file = self.save_dir / f'epoch_{epoch}_predictions.json'
        with open(output_file, 'w') as f:
            json.dump(predictions, f, indent=2)
        
        print(f"Saved validation predictions to {output_file}")

def train_with_validation_predictions(
    weights='yolov8l.pt',
    data='data/tvt3_data_v8.yaml',
    img_size=640,
    batch_size=16,
    epochs=3,
    single_cls=True
):
    # Initialize model
    model = YOLO(weights)
    
    # Create callback
    val_callback = ValidationPredictionCallback(
        save_dir=f'validation_predictions_{datetime.now().strftime("%Y%m%d_%H%M%S")}'
    )
    
    # Custom training loop
    model.train(
        data=data,
        imgsz=img_size,
        batch=batch_size,
        epochs=epochs,
        single_cls=single_cls,
        callbacks={
            'on_val_end': lambda validator: val_callback.save_predictions(
                validator, 
                validator.epoch
            )
        }
    )
    
    return model

# Example usage
if __name__ == "__main__":
    model = train_with_validation_predictions(
        weights='yolov8l.pt',
        data='data/tvt3_data_v8.yaml',
        img_size=640,
        batch_size=16,
        epochs=3,
        single_cls=True
    )

To run:

from train_yolo import train_with_validation_predictions

model = train_with_validation_predictions(
    weights='yolov8l.pt',
    data='data/tvt3_data_v8.yaml'
)

Ouptut will look like:

[
  {
    "image_path": "path/to/image.jpg",
    "predictions": [
      {
        "bbox": [x1, y1, x2, y2],
        "confidence": 0.95,
        "class": 0
      }
      // ... more predictions
    ]
  }
  // ... more images
]