Add vision.py

vision.py

@@ -0,0 +1,376 @@
+# -*- coding: utf-8 -*-
+"""vision.ipynb
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/drive/1JriMvbXyr0_2BXST58NUljv9sWWmgbHC
+"""
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+from torchvision import models, transforms
+from PIL import Image
+import os
+import numpy as np
+import time
+from tqdm import tqdm
+
+class Config:
+    seed = 42
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    batch_size = 32
+    num_workers = 4
+    learning_rate = 1e-4
+    num_epochs = 10
+    num_classes = 2
+    img_size = 224
+
+def seed_everything(seed):
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+
+seed_everything(Config.seed)
+print(f"Using device: {Config.device}")
+
+# Стандартные статистики ImageNet
+NORM_MEAN = [0.485, 0.456, 0.406]
+NORM_STD = [0.229, 0.224, 0.225]
+
+def get_transforms(phase='train'):
+    if phase == 'train':
+        return transforms.Compose([
+            transforms.Resize((256, 256)),                        # Сначала приводим к общему размеру
+            transforms.RandomResizedCrop(Config.img_size),        # Случайный кроп
+            transforms.RandomHorizontalFlip(p=0.5),               # Отражение
+            transforms.RandomRotation(degrees=15),                # Поворот
+            transforms.ColorJitter(brightness=0.2, contrast=0.2), # Изменение цвета
+            transforms.ToTensor(),
+            transforms.Normalize(NORM_MEAN, NORM_STD)
+        ])
+    else:
+        return transforms.Compose([
+            transforms.Resize((256, 256)),
+            transforms.CenterCrop(Config.img_size),
+            transforms.ToTensor(),
+            transforms.Normalize(NORM_MEAN, NORM_STD)
+        ])
+
+class CustomDataset(Dataset):
+    def __init__(self, file_paths, labels, transform=None):
+        self.file_paths = file_paths
+        self.labels = labels
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.file_paths)
+
+    def __getitem__(self, idx):
+        img_path = self.file_paths[idx]
+        image = Image.open(img_path).convert("RGB")
+        label = self.labels[idx]
+
+        if self.transform:
+            image = self.transform(image)
+
+        return image, torch.tensor(label, dtype=torch.long)
+
+def build_model(num_classes, pretrained=True):
+    # 1. Загружаем предобученный ResNet18
+    model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1 if pretrained else None)
+
+    # 2. (Опционально) Замораживаем веса бэкбона
+    # Это нужно, если данных мало. Если данных много, можно обучать всё (fine-tuning).
+    for param in model.parameters():
+        param.requires_grad = False
+
+    # 3. Заменяем "голову" (полносвязный слой)
+    # model.fc.in_features - это количество входов в оригинальном слое (512 для ResNet18)
+    num_ftrs = model.fc.in_features
+
+    model.fc = nn.Sequential(
+        nn.Linear(num_ftrs, 256),
+        nn.ReLU(),
+        nn.Dropout(0.5), # Для предотвращения переобучения
+        nn.Linear(256, num_classes)
+    )
+
+    return model
+
+model = build_model(Config.num_classes).to(Config.device)
+
+def train_one_epoch(model, loader, criterion, optimizer, device):
+    model.train()
+    running_loss = 0.0
+    correct = 0
+    total = 0
+
+    loop = tqdm(loader, leave=True) # Прогресс-бар
+
+    for images, labels in loop:
+        images, labels = images.to(device), labels.to(device)
+
+        optimizer.zero_grad()
+        outputs = model(images)
+        loss = criterion(outputs, labels)
+        loss.backward()
+
+        optimizer.step()
+
+        running_loss += loss.item()
+        _, predicted = torch.max(outputs.data, 1)
+        total += labels.size(0)
+        correct += (predicted == labels).sum().item()
+
+        loop.set_description(f"Train Loss: {loss.item():.4f}")
+
+    epoch_loss = running_loss / len(loader)
+    epoch_acc = 100 * correct / total
+    return epoch_loss, epoch_acc
+
+def validate(model, loader, criterion, device):
+    model.eval()
+    running_loss = 0.0
+    correct = 0
+    total = 0
+
+    with torch.no_grad():
+        for images, labels in loader:
+            images, labels = images.to(device), labels.to(device)
+
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+
+            running_loss += loss.item()
+            _, predicted = torch.max(outputs.data, 1)
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+
+    epoch_loss = running_loss / len(loader)
+    epoch_acc = 100 * correct / total
+    return epoch_loss, epoch_acc
+
+import tempfile
+fake_data_len = 100
+fake_paths = [tempfile.NamedTemporaryFile(suffix='.jpg').name for _ in range(fake_data_len)]
+for p in fake_paths:
+    Image.new('RGB', (300, 300)).save(p)
+fake_labels = np.random.randint(0, 2, fake_data_len)
+
+# Инициализация датасетов
+train_dataset = CustomDataset(fake_paths, fake_labels, transform=get_transforms('train'))
+val_dataset = CustomDataset(fake_paths, fake_labels, transform=get_transforms('val'))
+
+# DataLoader'ы
+train_loader = DataLoader(train_dataset, batch_size=Config.batch_size, shuffle=True, num_workers=0) # num_workers=0 для примера
+val_loader = DataLoader(val_dataset, batch_size=Config.batch_size, shuffle=False, num_workers=0)
+
+# Оптимизатор и Лосс
+# Обучаем только параметры fc (головы), если заморозили бэкбон.
+# Если не замораживали, передавайте model.parameters()
+optimizer = optim.Adam(model.fc.parameters(), lr=Config.learning_rate)
+criterion = nn.CrossEntropyLoss()
+
+# Основной цикл
+best_acc = 0.0
+
+print("Start Training...")
+for epoch in range(Config.num_epochs):
+    print(f"\nEpoch {epoch+1}/{Config.num_epochs}")
+
+    train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, Config.device)
+    val_loss, val_acc = validate(model, val_loader, criterion, Config.device)
+
+    print(f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%")
+    print(f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2f}%")
+
+    if val_acc > best_acc:
+        best_acc = val_acc
+        torch.save(model.state_dict(), "best_model.pth")
+        print("Model Saved!")
+
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+import cv2
+import torch
+import torch.nn as nn
+from torchvision import models
+
+class AugmentationFactory:
+    """Класс для создания пайплайна аугментаций"""
+    def __init__(self, img_size=224):
+        self.img_size = img_size
+
+        # Mean и Std для ImageNet (стандарт для предобученных моделей)
+        self.mean = (0.485, 0.456, 0.406)
+        self.std = (0.229, 0.224, 0.225)
+
+    def get_train_transforms(self):
+        return A.Compose([
+            A.Resize(height=256, width=256),
+            A.RandomCrop(height=self.img_size, width=self.img_size),
+
+            # Геометрические аугментации
+            A.HorizontalFlip(p=0.5),
+            A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.05, rotate_limit=15, p=0.5),
+
+            # Цветовые и шумовые аугментации (Albumentations тут очень силен)
+            A.OneOf([
+                A.GaussNoise(var_limit=(10.0, 50.0)),
+                A.GaussianBlur(),
+                A.MotionBlur(),
+            ], p=0.3),
+            A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1, p=0.3),
+
+            # Обязательные шаги в конце
+            A.Normalize(mean=self.mean, std=self.std),
+            ToTensorV2() # Конвертирует в torch.Tensor (C, H, W)
+        ])
+
+    def get_val_transforms(self):
+        return A.Compose([
+            A.Resize(height=self.img_size, width=self.img_size), # Или Resize -> CenterCrop
+            A.Normalize(mean=self.mean, std=self.std),
+            ToTensorV2()
+        ])
+
+# Обновленный Dataset под Albumentations
+class Cv2Dataset(torch.utils.data.Dataset):
+    def __init__(self, file_paths, labels, transforms=None):
+        self.file_paths = file_paths
+        self.labels = labels
+        self.transforms = transforms
+
+    def __len__(self):
+        return len(self.file_paths)
+
+    def __getitem__(self, idx):
+        path = self.file_paths[idx]
+
+        # 1. Читаем через OpenCV (BGR формат по умолчанию)
+        image = cv2.imread(path)
+        # 2. Конвертируем в RGB !!! Очень важно !!!
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        # 3. Применяем аугментации
+        if self.transforms:
+            # Albumentations возвращает словарь
+            augmented = self.transforms(image=image)
+            image = augmented['image']
+
+        label = torch.tensor(self.labels[idx], dtype=torch.long)
+        return image, label
+
+import torch.nn as nn
+
+class UniversalClassifier(nn.Module):
+    def __init__(self, model_name, num_classes, pretrained=True, freeze_backbone=False):
+        super().__init__()
+
+        if model_name not in AVAILABLE_BACKBONES:
+             raise ValueError(f"Model {model_name} not found.")
+
+        full_model = AVAILABLE_BACKBONES[model_name](weights="DEFAULT" if pretrained else None)
+        self.encoder = full_model
+
+        if freeze_backbone:
+            for param in self.encoder.parameters():
+                param.requires_grad = False
+
+        self.head_layer_name = ""
+
+        if "resnet" in model_name:
+            self.emb_dim = self.encoder.fc.in_features
+            self.encoder.fc = nn.Identity()
+
+        elif "efficientnet" in model_name:
+            self.emb_dim = self.encoder.classifier[-1].in_features
+            self.encoder.classifier[-1] = nn.Identity()
+
+        elif "vit" in model_name:
+            self.emb_dim = self.encoder.heads.head.in_features
+            self.encoder.heads.head = nn.Identity()
+
+        self.head = nn.Sequential(
+            nn.Dropout(p=0.3),
+            nn.Linear(self.emb_dim, num_classes)
+        )
+
+    def forward(self, x):
+        features = self.encoder(x)
+        output = self.head(features)
+        return output
+
+    def get_features(self, x):
+        """Метод специально для получения только эмбеддингов"""
+        return self.encoder(x)
+
+AVAILABLE_BACKBONES = {
+    # Тяжелые и точные
+    "resnet50": models.resnet50,
+    "efficientnet_b0": models.efficientnet_b0, # Хороший баланс
+    "efficientnet_b4": models.efficientnet_b4, # Мощнее
+
+    # Легкие (для мобилок/быстрого инференса)
+    "resnet18": models.resnet18,
+    "mobilenet_v3_large": models.mobilenet_v3_large,
+
+    # Современные (Transformers)
+    "vit_b_16": models.vit_b_16, # Требует img_size=224
+}
+
+"""# Пример"""
+
+# --- КОНФИГУРАЦИЯ ---
+class Config:
+    model_name = "efficientnet_b0"
+    num_classes = 2
+    img_size = 224 # EfficientNet_B0 любит 224, B4 любит 380
+    batch_size = 32
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# 1. Аугментации
+aug_factory = AugmentationFactory(img_size=Config.img_size)
+train_transforms = aug_factory.get_train_transforms()
+
+# 2. Создание датасета (пример путей)
+# train_dataset = Cv2Dataset(train_paths, train_labels, transforms=train_transforms)
+# train_loader = DataLoader(train_dataset, batch_size=Config.batch_size, shuffle=True)
+
+# 3. Инициализация модели
+model = UniversalClassifier(
+    model_name=Config.model_name,
+    num_classes=Config.num_classes,
+    pretrained=True,
+    freeze_backbone=False
+).to(Config.device)
+
+print(f"Model {Config.model_name} initialized successfully.")
+dummy_input = torch.randn(2, 3, Config.img_size, Config.img_size).to(Config.device)
+output = model(dummy_input)
+print(f"Output shape: {output.shape}")
+
+"""# Достать эмбединг"""
+
+model = UniversalClassifier("resnet18", num_classes=2).to(Config.device)
+
+def get_embeddings_clean(model, loader, device):
+    model.eval()
+    embeddings_list = []
+
+    with torch.no_grad():
+        for images, _ in tqdm(loader):
+            images = images.to(device)
+            features = model.get_features(images)
+            embeddings_list.append(features.cpu().numpy())
+
+    return np.vstack(embeddings_list)
+
+embs = get_embeddings_clean(model, val_loader, Config.device)
+
+embs
+