Initial Docker ML app deployment

.gitignore

@@ -0,0 +1,25 @@
+# See https://help.github.com/articles/ignoring-files/ for more about ignoring files.
+
+# dependencies
+/node_modules
+/.pnp
+.pnp.js
+/venv
+/__pycache__
+/myenv
+# testing
+/coverage
+
+# production
+/build
+
+# misc
+.DS_Store
+.env.local
+.env.development.local
+.env.test.local
+.env.production.local
+
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*

DockerFIle

@@ -0,0 +1,22 @@
+FROM python:3.10-slim
+
+# Create a non-root user like Hugging Face recommends
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+# Set working directory
+WORKDIR /app
+
+# Install Python dependencies
+COPY --chown=user requirements.txt requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the rest of the app
+COPY --chown=user . /app
+
+# Expose port
+EXPOSE 7860
+
+# Start Flask app
+CMD ["python", "app.py"]

app.py

@@ -0,0 +1,28 @@
+from flask import Flask, request, jsonify
+from flask_cors import CORS
+from model import load_model, predict_species, get_label_names
+
+app = Flask(__name__)
+
+CORS(app)
+
+# Load model once
+model = load_model()
+label_names = get_label_names()
+
+@app.route('/predict', methods=['GET'])
+def predict():
+    image_url = request.args.get('url')
+    if not image_url:
+        return jsonify({'error': 'URL parameter is missing'}), 400
+
+    try:
+        predicted_species = predict_species(model, image_url, label_names)
+        return jsonify({'species': predicted_species})
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+if __name__ == '__main__':
+    import os
+    app.run(host='0.0.0.0', port=int(os.environ.get('PORT', 7860)))
+

model.py

@@ -0,0 +1,58 @@
+import timm
+import torch
+from PIL import Image
+from torchvision import transforms
+import requests
+from io import BytesIO
+
+
+def load_model():
+    """Load the pre-trained model."""
+    model = timm.create_model("hf_hub:timm/vit_large_patch14_clip_336.laion2b_ft_in12k_in1k_inat21", pretrained=True)
+    model.eval()
+    return model
+
+
+def get_label_names():
+    """Fetch the class labels from the Hugging Face Hub."""
+    config_url = "https://huggingface.co/timm/vit_large_patch14_clip_336.laion2b_ft_in12k_in1k_inat21/resolve/main/config.json"
+    response = requests.get(config_url)
+    response.raise_for_status()
+    config = response.json()
+    return config["label_names"]
+
+
+def preprocess_image(image_url):
+    """Fetch and preprocess the image."""
+    preprocess = transforms.Compose([
+        transforms.Resize(336),
+        transforms.CenterCrop(336),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+
+    response = requests.get(image_url)
+    response.raise_for_status()
+    image = Image.open(BytesIO(response.content))
+    input_tensor = preprocess(image).unsqueeze(0)  # Add a batch dimension
+    return input_tensor
+
+
+def predict_species(model, image_url, label_names):
+    """Make a prediction using the model."""
+    input_tensor = preprocess_image(image_url)
+
+    # Move to GPU if available
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = model.to(device)
+    input_tensor = input_tensor.to(device)
+
+    # Make prediction
+    with torch.no_grad():
+        output = model(input_tensor)
+        _, predicted_class = torch.max(output, 1)
+
+    # Map prediction to species
+    predicted_species = label_names[predicted_class.item()]
+    return predicted_species
+#finish

requirements.txt

@@ -0,0 +1,7 @@
+flask
+flask-cors
+timm
+torch
+torchvision
+pillow
+requests