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)