File size: 1,003 Bytes
cc72945 027b8fc cc72945 027b8fc cc72945 027b8fc cc72945 027b8fc cc72945 027b8fc cc72945 027b8fc |
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 |
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)
|