/
IvlevAlexander0
/
QNN_FastApi
Обзор
Документация
Войти
/
IvlevAlexander0
/
QNN_FastApi
Код
Запросы
0
Задачи
Вики
Пакеты
0
Релизы
0
CI/CD
Аналитика
Безопасность
master
segmentation/segmenter.py
556 строк
25 KB
IvlevAlexander0
ver0.9.1
23 окт 2025, 08:17
23 окт 2025, 08:17
0d124a2
Код
Авторство
О чём код?
import torch from torch import nn from torch.nn import Module, Unfold import albumentations as A from albumentations.pytorch import ToTensorV2 import torch.nn.functional as F import numpy as np from PIL import Image, ImageOps from tqdm import tqdm import pennylane as qml class Quanv2d(Module): """ Квантовый сверточно-подобный слой - Выделяет сектора изображения с помощью torch.Unfold - Вычисляет результат работы квантовой цепочки в пачках (batches) """ def __init__( self, in_channels=3, out_channels=4, kernel_size=2, stride=2, device_name="lightning.qubit", # быстрая симуляция в PennyLane n_layers=3, batch_circuit_size=4096, diff_method="adjoint", # более быстрое вычисление градиентов чем в parameter-shift rule fast_mode=False, # если True, вместо квантовой цепочки включается нелинейная "заглушка" seed=13, ): ''' :param in_channels: число каналов во входном изображений (пока что поддерживается только 3) :param out_channels: число выходных каналов изображения :param kernel_size: размер ядра свертки (n_qubits = kernel_size**2 ) :param stride: отступ свертки :param device_name: выбор симулятора PennyLane :param n_layers: число слоев :param batch_circuit_size: сколько секторов обрабатывать с помощью квантовой цепочки за один прогон :param diff_method: метод вычисления градиентов :param fast_mode: вкл/выкл нелинейную заглушку вместо квантовой цепочки :param seed: сид ''' super().__init__() if in_channels != 3: raise ValueError("Quanv2d currently only supports in_channels=3") self.in_channels = in_channels self.out_channels = out_channels self.kernel_size = kernel_size self.stride = stride self.n_layers = n_layers self.batch_circuit_size = batch_circuit_size self.fast_mode = fast_mode # фиксация сида torch.manual_seed(seed) np.random.seed(seed) # объявление параметров self.n_qubits = kernel_size * kernel_size self.rand_params = nn.Parameter(torch.rand(n_layers, self.n_qubits) * 2 * np.pi) try: self.dev = qml.device(device_name, wires=self.n_qubits) except Exception: # откат в default.qubit если симулятор не поддерживается print(f'Device failed, reverting to default.qubit...') self.dev = qml.device("default.qubit", wires=self.n_qubits) # набор параметризованных гейтов alternating_op = [qml.RX, qml.RY, qml.RZ] @qml.qnode(self.dev, interface="torch", diff_method=diff_method) def circuit(inputs, weights): op_idx = 0 # преобразуем входные данные for c, rot_gate in enumerate(alternating_op): vals = inputs[c].reshape(-1) # 4 values for i in range(self.n_qubits): rot_gate(np.pi * vals[i], wires=i) # реализуем HEA анзац for l in range(weights.shape[0]): rot_op = alternating_op[op_idx] for q in range(self.n_qubits): rot_op(weights[l, q], wires=q) # запутанность только между соседями for q in range(self.n_qubits - 1): qml.CNOT(wires=[q, q + 1]) if self.n_qubits > 1: qml.CNOT(wires=[self.n_qubits - 1, 0]) op_idx = (op_idx + 1) % 3 wires_to_measure = [j % self.n_qubits for j in range(self.out_channels)] # измеряем среднее значение сигма з return [qml.expval(qml.PauliZ(w)) for w in wires_to_measure] self.circuit = circuit # преобразуем входной тензор в тензор из секторов которые мы подадим на вход self.unfold = Unfold(kernel_size=self.kernel_size, stride=self.stride) def forward(self, x): """ x: Тензор (batch, num_channels=3, height, width) returns: Тензор (batch, out_channels, height_out, width_out) """ N, C, H, W = x.shape assert C == 3 # еще одна проверка на три канала patches = self._extract_patches(x) N, L, F12 = patches.shape # проверка что размер сектора вычислен правильно для 3 каналов assert F12 == self.in_channels * self.kernel_size * self.kernel_size outputs = [] weights = self.rand_params B = self.batch_circuit_size H_out = (H - self.kernel_size) // self.stride + 1 W_out = (W - self.kernel_size) // self.stride + 1 # совмещаем размерности batch (N) и num_sectors (L) patches = patches.contiguous() patches_qlike = patches.view(N * L, self.in_channels, self.kernel_size, self.kernel_size) for start in tqdm(range(0, N * L, B), leave=False): end = min(start + B, N * L) batch_in = torch.tensor(patches_qlike[start:end]) # (batch_circuit_size, 3, ksize, ksize) batch_out = [] for i in tqdm(range(batch_in.shape[0]), leave=False): # здесь явно можно провести доп оптимизацию batch_out.append(torch.stack(self.circuit(batch_in[i], weights))) batch_out = torch.stack(batch_out, dim=0) outputs.append(batch_out) outputs = torch.cat(outputs, dim=0) outputs = outputs.view(N, L, self.out_channels) # (N, L, outC) out = outputs.transpose(1, 2).reshape(N, self.out_channels, H_out, W_out).to(torch.float) return out def _extract_patches(self, x): ''' Torch как то странно распределяет размерности тензора при Unfold, эта функция приводит их к более интуитивно понятному :param x: Тензор (batch, num_channels=3, height, width) :return: Тензор (batch, num_sectors, num_channels*(kernel_size**2)) ''' # (N, C*ks*ks, L) -> (N, L, C*ks*ks) p = self.unfold(x) # (N, 12, L) return p.transpose(1, 2) # (N, L, 12) class DoubleConv(nn.Module): """Составная часть UNET, (conv => BN => ReLU) x2""" def __init__(self, in_ch, out_ch): super().__init__() self.net = nn.Sequential( nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), ) def forward(self, x): return self.net(x) class Down(nn.Module): """Понижение размерности через maxpool затем double conv""" def __init__(self, in_ch, out_ch): super().__init__() self.pool = nn.MaxPool2d(2) self.block = DoubleConv(in_ch, out_ch) def forward(self, x): return self.block(self.pool(x)) class Up(nn.Module): """Повышение размерности затем double conv""" def __init__(self, in_ch, out_ch, bilinear=True): super().__init__() self.bilinear = bilinear if bilinear: self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False) self.conv = DoubleConv(in_ch, out_ch) else: self.up = nn.ConvTranspose2d(in_ch // 2, in_ch // 2, kernel_size=2, stride=2) self.conv = DoubleConv(in_ch, out_ch) def forward(self, x, skip): x = self.up(x) # пэддинг если размерности входа не делятся нацело diffY = skip.size(2) - x.size(2) diffX = skip.size(3) - x.size(3) if diffY != 0 or diffX != 0: x = F.pad(x, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2]) x = torch.cat([skip, x], dim=1) return self.conv(x) class OutConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv = nn.Conv2d(in_ch, out_ch, kernel_size=1) def forward(self, x): return self.conv(x) class UNetWithQuanv(nn.Module): ''' Сверточная сегментационная модель UNET с возможность классической или гибридной реализации ''' def __init__(self, in_channels=3, base_channels=64, bilinear=True, quanv_out_channels=4, quanv_kernel_size=2, quanv_stride=2, model_type='hybrid', quanv_kwargs=None): ''' :param in_channels: число каналов во входном изображений (пока что поддерживается только 3) :param base_channels: число каналов в первых сверточных слоях (базовое значение каналов) :param bilinear: вкл/выкл билинейный апсемплинг :param quanv_out_channels: число выходных каналов в квантовом слое :param quanv_kernel_size: размер "ядра свертки" квантового слоя :param quanv_stride: отступ квантового слоя :param model_type: тип модели, hybrid/classic :param quanv_kwargs: дополнительные аргументы для проброса в класс квантового слоя ''' super().__init__() if quanv_kwargs is None: quanv_kwargs = {} # Первый слой UNET self.quanv = Quanv2d( in_channels=in_channels, out_channels=quanv_out_channels, kernel_size=quanv_kernel_size, stride=quanv_stride, **quanv_kwargs ) self.first_conv = nn.Sequential( nn.Conv2d(in_channels, base_channels, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(base_channels), nn.ReLU(inplace=True), nn.Conv2d(base_channels, base_channels, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(base_channels), nn.ReLU(inplace=True), ) # слой с параллельной обработкой данных (квантовый слой находится здесь) self.parallel_processing_quantum = nn.Sequential( nn.MaxPool2d(4), # (batch, 3, H/2, W/2) nn.BatchNorm2d(in_channels), self.quanv, nn.Conv2d(quanv_out_channels, in_channels, kernel_size=3, padding=1, bias=False), nn.ConvTranspose2d(in_channels, base_channels, kernel_size=2, stride=2, bias=False), nn.Upsample(scale_factor=4, mode='bilinear', align_corners=False), nn.BatchNorm2d(base_channels), nn.ReLU(inplace=True), ) # классическая параллельная обработка self.parallel_processing_classic = nn.Sequential( nn.MaxPool2d(4), # (batch, 3, H/2, W/2) nn.BatchNorm2d(in_channels), nn.Conv2d(in_channels, quanv_out_channels, 2, stride=2), nn.Conv2d(quanv_out_channels, in_channels, kernel_size=3, padding=1, bias=False), nn.ConvTranspose2d(in_channels, base_channels, kernel_size=2, stride=2, bias=False), nn.Upsample(scale_factor=4, mode='bilinear', align_corners=False), nn.BatchNorm2d(base_channels), nn.ReLU(inplace=True), ) # выбор модели if model_type == 'hybrid': self.parallel_processing = self.parallel_processing_quantum elif model_type == 'classic': self.parallel_processing = self.parallel_processing_classic # self.q_up = nn.Upsample(scale_factor=quanv_stride, mode='bilinear', align_corners=False) # Энкодеры self.down1 = Down(base_channels, base_channels * 2) # /2 self.down2 = Down(base_channels * 2, base_channels * 4) # /4 self.down3 = Down(base_channels * 4, base_channels * 8) # /8 factor = 2 if bilinear else 1 self.down4 = Down(base_channels * 8, base_channels * 16 // factor) # /16 # Декодеры self.up1 = Up(base_channels * 16, base_channels * 8 // factor, bilinear=bilinear) self.up2 = Up(base_channels * 8, base_channels * 4 // factor, bilinear=bilinear) self.up3 = Up(base_channels * 4, base_channels * 2 // factor, bilinear=bilinear) # Последний Up конкатенируется с САМЫМ ПЕРВЫМ БЛОКОМ (после Quanv+Conv) self.up4 = Up(base_channels * 2, base_channels, bilinear=bilinear) # Выход сегментации self.outc = OutConv(base_channels, 1) def forward(self, x): x_0 = self.first_conv(x) # (N, base, H/2, W/2) x_par = self.parallel_processing(x) x0 = x_0 + x_par # совмещаем данные # Проход через энкодеры x1 = self.down1(x0) # (batch, 2*base, H/4, W/4) x2 = self.down2(x1) # (batch, 4*base, H/8, W/8) x3 = self.down3(x2) # (batch, 8*base, H/16, W/16) x4 = self.down4(x3) # (batch,16*base/f, H/32, W/32) # Проход через декодеры u1 = self.up1(x4, x3) # -> (batch, 8*base/f, H/16, W/16) u2 = self.up2(u1, x2) # -> (batch, 4*base/f, H/8, W/8) u3 = self.up3(u2, x1) # -> (batch, 2*base/f, H/4, W/4) u4 = self.up4(u3, x0) # -> (batch, base, H/2, W/2) logits = self.outc(u4) # НЕ применяем сигмоиду/софтмакс — лоссы сами разберутся return logits class SegmentationInference: ''' Класс взаимодействия с моделью в режиме сегментации (инференс). Поддерживает обработку как одного изображения, \ так и сразу множества изображений из указанной директории ''' def __init__(self, checkpoint_path, device=None, target_size=(224, 224), transform_type='crop'): ''' :param checkpoint_path: путь к сохраненной модели :param device: PyTorch девайс :param target_size: размер изображения которое пойдет на вход в модель :param transform_type: 'crop'/'resize' - обрезать изображение до target size/уменьшить изображение \ до target size с сохранением соотношения сторон и добавлением пэддинга ''' self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu") self.target_size = target_size self.checkpoint_path = checkpoint_path self.transform_type = transform_type self.transform_types = ['crop', 'resize'] # загрузка чекпоинта checkpoint = torch.load(checkpoint_path, map_location=self.device) model_type = 'hybrid'#checkpoint.get('model_type') seed = 42#checkpoint.get('seed') self.model = UNetWithQuanv( in_channels=3, base_channels=32, bilinear=True, # или False для ConvTranspose2d quanv_out_channels=4, model_type=model_type, quanv_kernel_size=2, quanv_stride=2, quanv_kwargs=dict( device_name="lightning.qubit", n_layers=3, batch_circuit_size=4096, diff_method="adjoint", fast_mode=False, seed=seed, ) ) # Загрузка весов модели self.model.load_state_dict(checkpoint["model_state_dict"]) self.model.to(self.device) self.model.eval() def resize_image_with_padding(self, image, target_size): """ Изменяет размер изображения до нужного, сохранив соотношение сторон. При необходимости добавляет отступы, чтобы избежать искажений. image: PIL Image target_size: (height, width) Returns: resized_image: изображение другого размера scale_factor: факторы увеличения размерности (scale_h, scale_w) padding: пэддинг (top, bottom, left, right) """ original_size = image.size # (width, height) target_width, target_height = target_size[1], target_size[0] # Вычисление факторов увеличения размерности scale_w = target_width / original_size[0] scale_h = target_height / original_size[1] scale = min(scale_w, scale_h) # Вычисление нового размера изображения new_width = int(original_size[0] * scale) new_height = int(original_size[1] * scale) # Изменение размера изображения resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS) # Вычисление пэддинга pad_left = (target_width - new_width) // 2 pad_right = target_width - new_width - pad_left pad_top = (target_height - new_height) // 2 pad_bottom = target_height - new_height - pad_top # Добавление пэддинга padded_image = ImageOps.expand( resized_image, border=(pad_left, pad_top, pad_right, pad_bottom), fill=0 # черный цвет пэддинга ) return padded_image, (scale, scale), (pad_top, pad_bottom, pad_left, pad_right) def remove_padding_from_mask(self, mask, original_size, padding, scale_factor): """ Убирает пэддинг из маски и возвращает маску к ее оригинальному размеру :param mask: Prediction mask (numpy array) :param original_size: Original image size (width, height) :param padding: (top, bottom, left, right) padding values :param scale_factor: (scale_h, scale_w) scaling factors Returns: resized_mask: Mask in original image dimensions """ pad_top, pad_bottom, pad_left, pad_right = padding scale_h, scale_w = scale_factor # Вычленение пэддинга if pad_top > 0 or pad_bottom > 0 or pad_left > 0 or pad_right > 0: h, w = mask.shape mask = mask[pad_top:h - pad_bottom, pad_left:w - pad_right] # Восстановление размера к исходному original_width, original_height = original_size if mask.size > 0: # Проверка - пуста ли маска mask_img = Image.fromarray(mask.astype(np.uint8)) resized_mask = mask_img.resize( (original_width, original_height), Image.Resampling.NEAREST # ресэмплинг ближайших соседей ) return np.array(resized_mask) else: return np.zeros((original_height, original_width), dtype=np.uint8) def preprocess_image(self, image): ''' Полный цикл предпоготовки изображения :param image: изображение :return: словарь с ключевыми параметрами и изображением ''' # Загрузка оригинального изображения original_image = image.convert("RGB")# original_size = original_image.size # (width, height) if self.transform_type == self.transform_types[0]: # crop transform_crop = A.Compose([ A.CenterCrop(self.target_size[0], self.target_size[1]), # Обрезка на tsize x tsize ]) # конвертация в numpy нужна для работы Albumentations original_image = np.array(original_image) processed_image = transform_crop(image=original_image)['image'] transform_final = A.Compose([ A.Normalize(), A.ToTensorV2(), ]) image = transform_final(image=processed_image)['image'] return { 'input_tensor': image.unsqueeze(0).to(self.device), 'processed_image': np.array(processed_image), 'original_image': np.array(original_image), 'original_size': original_size, } elif self.transform_type == self.transform_types[1]: # resize # Подгонка размера с сохранением соотношения сторон resized_image, scale_factor, padding = self.resize_image_with_padding( original_image, self.target_size ) image_array = np.array(resized_image) transform_crop = A.Compose([ A.Normalize(), ToTensorV2(), ]) transformed = transform_crop(image=image_array) return { 'input_tensor': transformed["image"].unsqueeze(0).to(self.device), 'original_image': np.array(original_image), 'original_size': original_size, 'scale_factor': scale_factor, 'padding': padding, 'processed_image': image_array } else: raise ValueError(f'Unknown transform type, available transforms: {self.transform_types}') def predict(self, original_image, threshold=0.5): """ Запуск инференса на одном изображении Returns: original_image: изначальное изображение prediction_mask: предсказанная маска confidence_scores: значения вероятностей preprocessed_data: словарь с информацией о предобработке """ preprocessed_data = self.preprocess_image(original_image) input_tensor = preprocessed_data['input_tensor'] with torch.no_grad(): logits = self.model(input_tensor) probs = torch.sigmoid(logits) prediction_mask = (probs > threshold).float().cpu().numpy()[0, 0] # (H, W) confidence_scores = probs.cpu().numpy()[0, 0] # (H, W) if self.transform_type == self.transform_types[1]: # Убрать пэддинг и вернуть маску к исходному формату prediction_mask = self.remove_padding_from_mask( prediction_mask, preprocessed_data['original_size'], preprocessed_data['padding'], preprocessed_data['scale_factor'] ) confidence_scores = self.remove_padding_from_mask( confidence_scores, preprocessed_data['original_size'], preprocessed_data['padding'], preprocessed_data['scale_factor'] ) return preprocessed_data['original_image'], prediction_mask, confidence_scores, preprocessed_data def return_segmenter(config): """Возвращает лямбда функцию сегментации одного изображения""" ''' Param: config: набор внешних параметров для создания модели: 'device': Устройство исполнения (cuda или cpu) 'target_size': Размер изображения которое пойдет на вход в модель (w, h) 'checkpoint': Путь к сохраненной модели 'transform_type': 'crop'/'resize' - обрезать изображение до target size/уменьшить изображение \ до target size с сохранением соотношения сторон и добавлением пэддинга Return: Лямбда функция сегментации одного изображения ''' model = SegmentationInference( device=config['device'], checkpoint_path=config['checkpoint'], target_size=tuple(config['target_size']), transform_type=config['transform_type'], ) return lambda img: model.predict(img)[1]