Update README.md

README.md

@@ -1,3 +1,45 @@
----
-license: apache-2.0
----
+---
+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)
+
+```