Feroc commited on
Commit
8dce058
1 Parent(s): 6968162

Upload exampleDetection.py

Browse files
Files changed (1) hide show
  1. exampleDetection.py +25 -0
exampleDetection.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import os
3
+ from transformers import AutoProcessor
4
+ from PIL import Image
5
+
6
+ model_path = "D:/nighttest/nightshadeblip_high.pth"
7
+ test_path = 'D:/nighttest/test/'
8
+
9
+ model = torch.load(model_path)
10
+ processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
11
+
12
+ model.eval()
13
+
14
+ for filename in os.listdir(test_path):
15
+ file_path = os.path.join(test_path, filename)
16
+ if os.path.isfile(file_path):
17
+ image = Image.open(file_path)
18
+
19
+ device = "cuda"
20
+ inputs = processor(images=image, return_tensors="pt").to(device)
21
+ pixel_values = inputs.pixel_values
22
+
23
+ generated_ids = model.generate(pixel_values=pixel_values, max_length=50)
24
+ generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
25
+ print(f"[{generated_caption}] {file_path}")