Views
No views yet
Senko (The Helpful Fox Senko-san)ShiroSoraOther (The rest, not Senko)Senko (Сенко-сан)Shiro (Широ)Sora (Сора)Other (Остальные, не Сенко).pth file. You need to define the wrapper class to load it.
Модель упакована в автономный файл .pth. Вам нужно определить класс-обертку для её загрузки.1import torch
2import torch.nn as nn
3from transformers import AutoModel, AutoConfig, AutoImageProcessor
4from PIL import Image
5import numpy as np
6
7# 1. Define the wrapper class / Определяем класс
8class StandaloneDino(nn.Module):
9 def __init__(self, config_dict, num_classes):
10 super().__init__()
11 config = AutoConfig.for_model(**config_dict)
12 self.backbone = AutoModel.from_config(config)
13 self.classifier = nn.Linear(config.hidden_size, num_classes)
14
15 def forward(self, x):
16 outputs = self.backbone(pixel_values=x)
17 return self.classifier(outputs.pooler_output)
18
19# 2. Load the model / Загружаем модель
20# Download 'Senko_Detector_DinoV3_v1.pth' from Files tab
21MODEL_FILE = "Senko_Detector_DinoV3_v1.pth"
22DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
23
24checkpoint = torch.load(MODEL_FILE, map_location=DEVICE)
25model = StandaloneDino(checkpoint['architecture_config'], len(checkpoint['classes']))
26model.load_state_dict(checkpoint['state_dict'])
27model.to(DEVICE).eval()
28
29# 3. Preprocessing & Inference / Инференс
30def prepare_native(image, patch_size=16):
31 w, h = image.size
32 # Resize only if image is too huge / Ресайз только если картинка гигантская
33 if max(w, h) > 2500:
34 scale = 2500 / max(w, h)
35 image = image.resize((int(w * scale), int(h * scale)), Image.Resampling.BICUBIC)
36 w, h = image.size
37
38 # Crop to patch size / Подгонка под размер патча 16
39 new_w = w - (w % patch_size)
40 new_h = h - (h % patch_size)
41 if new_w != w or new_h != h:
42 image = image.crop((0, 0, new_w, new_h))
43 return image
44
45# Load Image
46image_path = "test.jpg"
47image = Image.open(image_path).convert("RGB")
48image = prepare_native(image)
49
50# Normalize
51p_conf = checkpoint['processor_config']
52mean = np.array(p_conf.get('image_mean', [0.485, 0.456, 0.406])).reshape(1, 1, 3)
53std = np.array(p_conf.get('image_std', [0.229, 0.224, 0.225])).reshape(1, 1, 3)
54
55img_arr = np.array(image).astype(np.float32) / 255.0
56input_tensor = (img_arr - mean) / std
57input_tensor = torch.from_numpy(input_tensor.transpose(2, 0, 1)).unsqueeze(0).to(DEVICE)
58
59# Predict
60with torch.no_grad():
61 logits = model(input_tensor)
62 probs = torch.softmax(logits, dim=1)
63 conf, pred_idx = torch.max(probs, 1)
64
65 print(f"Class: {checkpoint['classes'][pred_idx]} | Confidence: {conf.item():.4f}")
