File size: 1,549 Bytes
1c7d0a2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
---
license: apache-2.0
---
アニメイラストのセリフや擬音を検出するモデルです
ベースモデルにInternViT-6B-448px-V1-5を使用しています。
https://huggingface.co/OpenGVLab/InternViT-6B-448px-V1-5を使用しています。


ベースモデルのpooler_output層にこんな感じに繋げば使えると思います。
```python
import torch
from PIL import Image
from transformers import AutoModel, CLIPImageProcessor

class CustomModel(nn.Module):
    def __init__(self, base_model, num_classes=2):
        super(CustomModel, self).__init__()
        self.base_model = base_model
        self.classifier = nn.Linear(base_model.config.hidden_size, num_classes).to(torch.bfloat16)
    
    def forward(self, x):
        outputs = self.base_model(x)
        pooled_output = outputs.pooler_output
        logits = self.classifier(pooled_output)
        return logits

base_model = AutoModel.from_pretrained(
    'OpenGVLab/InternViT-6B-448px-V1-5',
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
    trust_remote_code=True).cuda().eval()

model = CustomModel(base_model, num_classes=2).to(device).eval()
model.classifier.load_state_dict(torch.load("checkpoints/classifier_weights.pth"))

image = Image.open('./examples/image1.jpg').convert('RGB')

image_processor = CLIPImageProcessor.from_pretrained('OpenGVLab/InternViT-6B-448px-V1-5')

pixel_values = image_processor(images=image, return_tensors='pt').pixel_values.to(torch.bfloat16).cuda()

with torch.no_grad():
  outputs = model(pixel_values)

```