import os

# Force Hugging Face cache directory to a writable path
os.environ["HF_HOME"] = "/tmp/huggingface"
os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface/transformers"
os.environ["HF_DATASETS_CACHE"] = "/tmp/huggingface/datasets"


from flask import Flask, request, jsonify
from PIL import Image
import torch
import torchvision.transforms as transforms
from transformers import AutoModelForImageClassification

# Load model
MODEL_NAME = "anismizi/skin-type-classifier"
model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
model.eval()

# Define preprocessing
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

app = Flask(__name__)

@app.route("/")
def home():
    return jsonify({"message": "Skin Type Classifier API is running!"})

@app.route("/predict", methods=["POST"])
def predict():
    if 'file' not in request.files:
        return jsonify({"error": "No file provided"}), 400
    file = request.files['file']
    try:
        image = Image.open(file.stream).convert("RGB")
        input_tensor = transform(image).unsqueeze(0)
        with torch.no_grad():
            outputs = model(input_tensor)
            probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
            predicted_class = probabilities.argmax().item()
            confidence = probabilities[0][predicted_class].item()

        labels = ["dry", "oily"]
        result = {
            "predicted_class": labels[predicted_class],
            "confidence": round(confidence * 100, 2)
        }
        return jsonify(result)
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)