Upload app.py with huggingface_hub

app.py

@@ -1,27 +1,20 @@
-import torch
 from transformers import AutoFeatureExtractor, EfficientNetForImageClassification
+import torch
 from PIL import Image
-import base64
 import io
+import base64
 
-def preprocess_image(image_data):
-    if image_data.startswith('data:image'):
-        image_data = image_data.split(',')[1]
-    image_bytes = base64.b64decode(image_data)
-    image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
-    return image
-
-def predict(image_data):
-    image = preprocess_image(image_data)
+def pipeline(image_bytes):
+    image = Image.open(io.BytesIO(base64.b64decode(image_bytes))).convert('RGB')
     
-    feature_extractor = AutoFeatureExtractor.from_pretrained("./huggingface_model")
-    model = EfficientNetForImageClassification.from_pretrained("./huggingface_model")
+    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.load_state_dict(torch.load("model.pt", map_location=torch.device('cpu')))
     model.eval()
     
     inputs = feature_extractor(images=image, return_tensors="pt")
@@ -33,5 +26,5 @@ def predict(image_data):
     
     return {"prediction": float(prediction)}
 
-def run(raw_image_data):
-    return predict(raw_image_data)
+def run(raw_image_bytes):
+    return pipeline(raw_image_bytes)