@@ -5,18 +5,18 @@ import timm
from pathlib import Path
import numpy as np
# ТЕПЕРЬ BLIP-2
from transformers import Blip2Processor , Blip2ForConditionalGeneration
# НОВЫЙ ИМПОРТ ДЛЯ VLM
from transformers import BlipProcessor , BlipForConditionalGeneration
class ImageProcessor :
def __init__ ( self , model_path : Path | str ) :
self . device = torch . device ( ' cuda ' if torch . cuda . is_available ( ) else ' cpu ' )
# --- ПОТОК 1: ТВОЙ АВТОРСКИЙ ЭМБЕДДИНГ (Core ) ---
# --- ПОТОК 1: ЭМОЦИИ (ResNet-50 ) ---
print ( " ⏳ Загрузка эмоционального модуля (ResNet-50)... " )
self . emo_model = timm . create_model ( ' resnet50 ' , pretrained = False , num_classes = 8 )
if Path ( model_path ) . exists ( ) :
self . emo_model . load_state_dict ( torch . load ( model_path , map_location = self . device ) )
self . emo_model . fc = torch . nn . Identity ( )
self . emo_model . to ( self . device ) . eval ( )
@@ -26,31 +26,32 @@ class ImageProcessor:
T . Normalize ( mean = [ 0.485 , 0.456 , 0.406 ] , std = [ 0.229 , 0.224 , 0.225 ] )
] )
# --- ПОТОК 2: СЕМАНТИЧЕСКИЙ ЭКСПЕР Т BLIP-2 ---
print ( " ⏳ Загрузка тяжелой артиллерии: BLIP-2 ... " )
# Используем версию opt-2.7b — она идеально сбалансирована для V100
self . blip_processor = Blip2Processor . from_pretrained ( " Salesforce/blip2-opt-2.7b " )
self . blip_model = Blip2ForConditionalGeneration . from_pretrained (
" Salesforce/blip2-opt-2.7b " ,
torch_dtype = torch . float16 # Обязательно для скорости на V100
) . to ( self . device )
print ( " ✅ BLIP-2 и ResNet-50 готовы. " )
# --- ПОТОК 2: СЕМАНТИКА И КОНТЕКС Т ( BLIP Large) ---
print ( " ⏳ Загрузка мощной VLM модели (BLIP) для описания сцен ... " )
# Используем версию Large, так как позволяет железо V100
self . blip_processor = BlipProcessor . from_pretrained ( " Salesforce/blip-image-captioning-large " )
self . blip_model = BlipForConditionalGeneration . from_pretrained ( " Salesforce/blip-image-captioning-large " ) . to ( self . device )
print ( " ✅ Обе нейросети визуального анализа успешно загружены на V100! " )
@torch.no_grad ( )
def extract_embedding ( self , image : Image . Image ) - > np . ndarray :
""" Извлекает 2048-мерный вектор эмоций. """
img_rgb = image . convert ( ' RGB ' )
img_tensor = self . emo_transform ( img_rgb ) . unsqueeze ( 0 ) . to ( self . device )
return self . emo_model ( img_tensor ) . cpu ( ) . numpy ( ) . flatten ( )
@torch.no_grad ( )
def describe_scene ( self , image : Image . Image ) - > str :
""" Генерирует описание через BLIP-2 . """
""" Генерирует текстовое описание картинки (Captioning) для LLM . """
img_rgb = image . convert ( ' RGB ' )
# Инференс BLIP-2 требует float16 для V100
inputs = self . blip_processor ( images = img_rgb , return_tensors = " pt " ) . to ( self . device , torch . float16 )
# Готовим картинку для BLIP
inputs = self . blip_processor ( img_rgb , return_tensors = " pt " ) . to ( self . device )
# Генерируем описание
generated_ids = self . blip_model . generate ( * * inputs , max_new_tokens = 40 )
caption = self . blip_processor . batch_decode ( generated_ids , skip_special_tokens = True ) [ 0 ] . strip ( )
# Генерируем описание (max_new_tokens ограничим, чтобы было лаконично)
out = self . blip_model . generate ( * * inputs , max_new_tokens = 30 )
# Декодируем тензор в строку
caption = self . blip_processor . decode ( out [ 0 ] , skip_special_tokens = True )
return caption