chore: change text output
This commit is contained in:
@@ -4,15 +4,12 @@ from PIL import Image
|
||||
import timm
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
|
||||
# ТЕПЕРЬ BLIP-2
|
||||
from transformers import Blip2Processor, Blip2ForConditionalGeneration
|
||||
|
||||
class ImageProcessor:
|
||||
def __init__(self, model_path: Path | str):
|
||||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
# --- ПОТОК 1: ТВОЙ АВТОРСКИЙ ЭМБЕДДИНГ (Core) ---
|
||||
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))
|
||||
@@ -26,15 +23,13 @@ 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
|
||||
print("Загрузка BLIP-2...")
|
||||
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
|
||||
torch_dtype=torch.float16
|
||||
).to(self.device)
|
||||
print("✅ BLIP-2 и ResNet-50 готовы.")
|
||||
print("BLIP-2 и ResNet-50 готовы.")
|
||||
|
||||
@torch.no_grad()
|
||||
def extract_embedding(self, image: Image.Image) -> np.ndarray:
|
||||
@@ -47,10 +42,8 @@ class ImageProcessor:
|
||||
"""Генерирует описание через BLIP-2."""
|
||||
img_rgb = image.convert('RGB')
|
||||
|
||||
# Инференс BLIP-2 требует float16 для V100
|
||||
inputs = self.blip_processor(images=img_rgb, return_tensors="pt").to(self.device, torch.float16)
|
||||
|
||||
# Генерируем описание
|
||||
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()
|
||||
return caption
|
||||
@@ -56,5 +56,5 @@ Example: {{"energy": 0.5, "flux": 0.2, "centroid": 0.4, "pitch": 0.3, "hnr": 0.8
|
||||
return profile
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"⚠️ Ошибка связи с локальной LLM: {e}")
|
||||
print(f"Ошибка связи с локальной LLM: {e}")
|
||||
return None
|
||||
Reference in New Issue
Block a user