| from transformers import AutoFeatureExtractor, EfficientNetForImageClassification | |
| import torch | |
| from PIL import Image | |
| import io | |
| import base64 | |
| def pipeline(image_bytes): | |
| image = Image.open(io.BytesIO(base64.b64decode(image_bytes))).convert('RGB') | |
| feature_extractor = AutoFeatureExtractor.from_pretrained(".") | |
| model = EfficientNetForImageClassification.from_pretrained(".") | |
| # Replace the classification head with a regression head | |
| model.classifier = torch.nn.Linear(model.classifier.in_features, 1) | |
| # Load the custom weights | |
| model.load_state_dict(torch.load("model.pt", map_location=torch.device('cpu'))) | |
| model.eval() | |
| inputs = feature_extractor(images=image, return_tensors="pt") | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| prediction = outputs.logits.item() # For regression, we directly use the output | |
| return {"prediction": float(prediction)} | |
| def run(raw_image_bytes): | |
| return pipeline(raw_image_bytes) | |