import os import torch import torch.nn as nn import torch.optim as optim import logging from torch.utils.data import DataLoader, TensorDataset from torch.nn.utils.rnn import pad_sequence from torchmetrics.classification import ( MultilabelPrecision, MultilabelRecall, MultilabelF1Score, MultilabelAccuracy, ) from typing import List, Union from joblib import Memory logger = logging.getLogger(__name__) memory = Memory("./cache_dir", verbose=0) # Set up disk-based caching for model results _models = {} @memory.cache(ignore=["yolo_model"]) def cached_yolo_predict(yolo_model, video_path, cache_params): """Predict bounding boxes and labels using YOLO model and cache the results using joblib. Args: yolo_model (YOLO): YOLO model instance video_path (str): Path to the video file cache_params (str): Parameters for caching the results, they are used in @memory.cache decorator """ frames = [] generator = yolo_model.predict(video_path, stream=True) for frame in generator: frame.orig_img = None # remove image from cache to reduce size frames.append(frame) return frames @memory.cache(ignore=["yolo_model"]) def cached_feature_extraction(yolo_model, video_path, cache_params): """Extract features from the last layer of the YOLO model and cache them using joblib. Args: yolo_model (YOLO): YOLO model instance video_path (str): Path to the video file cache_params (str): Parameters for caching the results, they are used in @memory.cache decorator """ layer_output = [None] def get_last_layer_output(module, input, output): layer_output[0] = input # Register the hook on the last layer of the model layer = yolo_model.model.model[-1].linear hook_handle = layer.register_forward_hook(get_last_layer_output) # Run model prediction, use stream to avoid out of memory generator = yolo_model.predict(video_path, stream=True) # Replace probs with last layer outputs frames = [] for frame in generator: frame.orig_img = None # remove image from cache to reduce size frame.probs = layer_output[0][0][0] # => tensor, 1280 floats for yolov8n-cls frames.append(frame) # Remove the hook hook_handle.remove() return frames class BaseNN(nn.Module): def __init__(self, **kwargs): super(BaseNN, self).__init__() self.label_map = None def set_label_map(self, label_map): self.label_map = label_map def get_label_map(self): return self.label_map def save(self, path): # ultralytics yolo11 patches torch.save to use dill, # however it leads to serialization errors, # so let's check for use_dill and disable it if 'use_dill' in torch.save.__code__.co_varnames: torch.save(self, path, use_dill=False) else: torch.save(self, path) logger.info(f"Model saved to {path}") @classmethod def load(cls, path) -> "BaseNN": model = torch.load(path) model.eval() # Set the model to evaluation mode logger.info(f"Model loaded from {path}") return model @classmethod def load_cached_model(cls, model_path: str) -> Union["BaseNN", None]: global _models if not os.path.exists(model_path): return None # Load per-project classifier if model_path not in _models: _models[model_path] = BaseNN.load(model_path) return _models[model_path] def save_and_cache(self, path): self.save(path) _models[path] = self class MultiLabelLSTM(BaseNN): def __init__( self, input_size, output_size, fc_size=128, hidden_size=16, num_layers=1, sequence_size=16, learning_rate=1e-4, weight_decay=1e-5, dropout_rate=0.2, device=None, **kwargs, ): """Initialize the MultiLabelLSTM model. Args: input_size (int): Number of features in the input data output_size (int): Number of labels in the output data fc_size (int): Size of the fully connected layer hidden_size (int): Size of the hidden state in the LSTM num_layers (int): Number of layers in the LSTM sequence_size (int): Size of the input sequence, used for chunking learning_rate (float): Learning rate for the optimizer weight_decay (float): Weight decay for the optimizer for L2 regularization dropout_rate (float): Dropout rate for the fully connected layer device (torch.device): Device to run the model on (CPU or GPU) """ super(MultiLabelLSTM, self).__init__() self.input_size = input_size self.output_size = output_size self.fc_size = fc_size self.hidden_size = hidden_size self.num_layers = num_layers self.sequence_size = sequence_size self.learning_rate = learning_rate self.weight_decay = weight_decay self.dropout_rate = dropout_rate # Reduce dimensionality of input data self.fc_input = nn.Linear(input_size, fc_size) self.layer_norm = nn.LayerNorm(fc_size) self.dropout = nn.Dropout(self.dropout_rate) # LSTM layer for handling sequential data self.lstm = nn.LSTM( fc_size, hidden_size, num_layers, batch_first=True, bidirectional=True ) # Fully connected layer for classification at each time step # 2 because of bidirectional LSTM self.fc = nn.Linear(2 * hidden_size, output_size) # Initialize the loss function and optimizer self.criterion = ( nn.BCEWithLogitsLoss() ) # Binary cross-entropy for multi-label classification self.optimizer = optim.Adam( self.parameters(), lr=learning_rate, weight_decay=weight_decay ) # Initialize device (CPU or GPU) target_device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.device = device if device else target_device self.to(self.device) def forward(self, x): # Reduce the dimensionality of the input data x = torch.relu(self.fc_input(x)) x = self.layer_norm(x) x = self.dropout(x) # x shape: (batch_size, seq_len, input_size) # lstm_out contains outputs for all time steps lstm_out, (_, _) = self.lstm(x) # Apply fully connected layer at each time step to get output with final label number out = self.fc(lstm_out) # Output shape: (batch_size, seq_len, output_size) return out def preprocess_sequence(self, sequence: List[torch.Tensor], labels=None, overlap=2): sequence = torch.stack(sequence) if isinstance(sequence, list) else sequence sequence_size = self.sequence_size # Split the data into small sequences by sequence_size with overlap chunks = [ sequence[i : i + sequence_size] for i in range(0, len(sequence), sequence_size // overlap) ] chunks = pad_sequence(chunks, batch_first=True, padding_value=0) if labels is not None: labels = torch.tensor(labels, dtype=torch.float32) labels = [ labels[i : i + sequence_size] for i in range(0, len(labels), sequence_size // overlap) ] labels = pad_sequence(labels, batch_first=True, padding_value=0) return chunks, labels def evaluate_metrics(self, dataloader, threshold=0.5): self.eval() params = { "num_labels": self.output_size, "average": "macro", "threshold": threshold, "zero_division": 1, } precision_metric = MultilabelPrecision(**params).to(self.device) recall_metric = MultilabelRecall(**params).to(self.device) f1_metric = MultilabelF1Score(**params).to(self.device) accuracy_metric = MultilabelAccuracy(**params).to(self.device) with torch.no_grad(): for data, labels in dataloader: data = data.to(self.device) labels = labels.to(self.device) outputs = torch.sigmoid(self(data)) # Reshape outputs and labels from (batch_size, seq_len, num_labels) # to (batch_size * seq_len, num_labels) outputs = outputs.view(-1, self.output_size) labels = labels.view(-1, self.output_size) # No need to threshold manually; the metrics handle it precision_metric.update(outputs, labels) recall_metric.update(outputs, labels) f1_metric.update(outputs, labels) accuracy_metric.update(outputs, labels) return { "precision": precision_metric.compute().item(), "recall": recall_metric.compute().item(), "f1_score": f1_metric.compute().item(), "accuracy": accuracy_metric.compute().item(), } def partial_fit( self, sequence, labels, batch_size=32, epochs=1000, accuracy_threshold=1.0, f1_threshold=1.0, ): """Train the model on the given sequence data. Args: sequence (List[torch.Tensor]): List of tensors containing the input data labels (List[List[int]]): List of lists containing the labels for each time step batch_size (int): Batch size for training epochs (int): Number of training epochs accuracy_threshold (float): Stop training if accuracy exceeds this threshold f1_threshold (float): Stop training if F1 score exceeds this threshold """ batches, label_batches = self.preprocess_sequence(sequence, labels) # Create a DataLoader for batching the input data metrics = {} dataset = TensorDataset( batches, torch.tensor(label_batches, dtype=torch.float32) ) dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True) for epoch in range(epochs): self.train() # Set the model to training mode epoch_loss = 0 for batch_data, batch_labels in dataloader: # Move batch data and labels to the appropriate device batch_data = batch_data.to(self.device) batch_labels = batch_labels.to(self.device) self.optimizer.zero_grad() outputs = self(batch_data) # Forward pass loss = self.criterion(outputs, batch_labels) # Calculate loss loss.backward() # Back propagation self.optimizer.step() # Update model parameters epoch_loss += loss.item() # metrics and threshold stops to avoid overfitting metrics = self.evaluate_metrics(dataloader) metrics["loss"] = epoch_loss / len(dataloader) metrics["epoch"] = epoch + 1 logger.info( f"Epoch {epoch + 1}, Loss: {epoch_loss / len(dataloader)}, {metrics}" ) if metrics["accuracy"] >= accuracy_threshold: logger.info( f"Accuracy >= {accuracy_threshold} threshold, model training stopped." ) break if metrics["f1_score"] >= f1_threshold: logger.info( f"F1 score >= {f1_threshold} threshold, model training stopped." ) break return metrics def predict(self, sequence): """Split sequence into chunks with sequence_size and predict by chunks. Then concatenate all predictions into one sequence of labels """ length = len(sequence) if length == 0: return torch.tensor([]) batches, _ = self.preprocess_sequence(sequence, overlap=1) self.eval() logits = torch.sigmoid(self(batches)) # Concatenate batches to sequence back shape = logits.shape logits = torch.reshape(logits, [shape[0] * shape[1], shape[2]]) return logits[0:length]