Create inference.py

inference.py

@@ -0,0 +1,28 @@
+import torch
+from transformers import ViTForImageClassification, ViTFeatureExtractor
+from PIL import Image
+import json
+
+# Load metadata from the JSON file
+with open('metadata.json') as f:
+    metadata = json.load(f)
+
+def predict(image_path: str):
+    # Load the fine-tuned model and feature extractor from Hugging Face Hub
+    model = ViTForImageClassification.from_pretrained("yigagilbert/image-quality-model")
+    feature_extractor = ViTFeatureExtractor.from_pretrained("yigagilbert/image-quality-model")
+    
+    # Open and preprocess the image
+    image = Image.open(image_path)
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    
+    # Perform inference
+    with torch.no_grad():
+        outputs = model(**inputs)
+        predicted_value = outputs.logits.squeeze().item()
+    
+    # Scale the predicted value to match the dataset's max value
+    max_value = metadata.get('max_value', 1.0)  # Default to 1.0 if not found
+    predicted_value_scaled = predicted_value * max_value
+    
+    return predicted_value_scaled