File size: 1,890 Bytes
2d30574 |
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 46 47 48 49 50 51 52 53 54 |
import requests
import torch
import torch.nn.functional as F
from transformers import ImageClassificationPipeline, AutoConfig, AutoTokenizer
from PIL import Image
from io import BytesIO
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_name_or_path = "Akazi/resnet_c_s_redwood_finetuned"
config = AutoConfig.from_pretrained(model_name_or_path)
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
class CustomImageClassificationPipeline(ImageClassificationPipeline):
def _sanitize_parameters(self, **pipeline_parameters):
return {}, {}, {}
def preprocess(self, inputs):
if isinstance(inputs, str):
if inputs.startswith("http://") or inputs.startswith("https://"):
response = requests.get(inputs)
image = Image.open(BytesIO(response.content))
else:
image = Image.open(inputs)
elif isinstance(inputs, bytes):
image = Image.open(BytesIO(inputs))
else:
image = inputs
inputs = tokenizer(image, return_tensors="pt")
inputs.to(device)
return inputs
def _forward(self, model_inputs):
return self.model(**model_inputs)
def postprocess(self, model_outputs):
logits = model_outputs.logits
probabilities = F.softmax(logits, dim=1).detach().cpu().numpy()[0]
labels = [config.id2label[i] for i in range(len(probabilities))]
outputs = [{"label": label, "score": round(score, 5)} for label, score in zip(labels, probabilities)]
return outputs
# Create the pipeline
pipeline = CustomImageClassificationPipeline(model=model_name_or_path, device=device)
# Example usage
image_path = "path/to/your/image.jpg"
results = pipeline(image_path)
# Print the results
for result in results:
print(f"Label: {result['label']}, Score: {result['score']}")
|