File size: 3,644 Bytes
cbcaa4f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | from typing import Generator, Iterable, List, TypeVar
import numpy as np
import supervision as sv
import torch
import umap
from sklearn.cluster import KMeans
from tqdm import tqdm
from transformers import AutoProcessor, SiglipVisionModel
V = TypeVar("V")
SIGLIP_MODEL_PATH = 'google/siglip-base-patch16-224'
def create_batches(
sequence: Iterable[V], batch_size: int
) -> Generator[List[V], None, None]:
"""
Generate batches from a sequence with a specified batch size.
Args:
sequence (Iterable[V]): The input sequence to be batched.
batch_size (int): The size of each batch.
Yields:
Generator[List[V], None, None]: A generator yielding batches of the input
sequence.
"""
batch_size = max(batch_size, 1)
current_batch = []
for element in sequence:
if len(current_batch) == batch_size:
yield current_batch
current_batch = []
current_batch.append(element)
if current_batch:
yield current_batch
class TeamClassifier:
"""
A classifier that uses a pre-trained SiglipVisionModel for feature extraction,
UMAP for dimensionality reduction, and KMeans for clustering.
"""
def __init__(self, device: str = 'cpu', batch_size: int = 32):
"""
Initialize the TeamClassifier with device and batch size.
Args:
device (str): The device to run the model on ('cpu' or 'cuda').
batch_size (int): The batch size for processing images.
"""
self.device = device
self.batch_size = batch_size
self.features_model = SiglipVisionModel.from_pretrained(
SIGLIP_MODEL_PATH).to(device)
self.processor = AutoProcessor.from_pretrained(SIGLIP_MODEL_PATH)
self.reducer = umap.UMAP(n_components=3)
self.cluster_model = KMeans(n_clusters=2)
def extract_features(self, crops: List[np.ndarray]) -> np.ndarray:
"""
Extract features from a list of image crops using the pre-trained
SiglipVisionModel.
Args:
crops (List[np.ndarray]): List of image crops.
Returns:
np.ndarray: Extracted features as a numpy array.
"""
crops = [sv.cv2_to_pillow(crop) for crop in crops]
batches = create_batches(crops, self.batch_size)
data = []
with torch.no_grad():
for batch in tqdm(batches, desc='Embedding extraction'):
inputs = self.processor(
images=batch, return_tensors="pt").to(self.device)
outputs = self.features_model(**inputs)
embeddings = torch.mean(outputs.last_hidden_state, dim=1).cpu().numpy()
data.append(embeddings)
return np.concatenate(data)
def fit(self, crops: List[np.ndarray]) -> None:
"""
Fit the classifier model on a list of image crops.
Args:
crops (List[np.ndarray]): List of image crops.
"""
data = self.extract_features(crops)
projections = self.reducer.fit_transform(data)
self.cluster_model.fit(projections)
def predict(self, crops: List[np.ndarray]) -> np.ndarray:
"""
Predict the cluster labels for a list of image crops.
Args:
crops (List[np.ndarray]): List of image crops.
Returns:
np.ndarray: Predicted cluster labels.
"""
if len(crops) == 0:
return np.array([])
data = self.extract_features(crops)
projections = self.reducer.transform(data)
return self.cluster_model.predict(projections)
|