import sys
import os
sys.path.insert(0, os.path.dirname(__file__))

from flask import Flask, request, jsonify
from transformers import AutoModelForImageClassification, AutoImageProcessor
from PIL import Image
import io, base64, torch
import torch.nn.functional as F

app = Flask(__name__)

processor = AutoImageProcessor.from_pretrained("microsoft/resnet-50")
model = AutoModelForImageClassification.from_pretrained("anismizi/skin-type-classifier")
model.eval()
labels = ["dry", "oily"]

@app.route('/', methods=['GET'])
def home():
    return {"status": "🧴 Skin Type API!", "predict": "/predict POST base64 image"}

@app.route('/predict', methods=['POST'])
def predict():
    try:
        data = request.json
        img_data = data['image'].split(',')[1]
        img_bytes = base64.b64decode(img_data)
        image = Image.open(io.BytesIO(img_bytes)).convert('RGB')
        inputs = processor(images=image, return_tensors="pt")
        with torch.no_grad():
            outputs = model(**inputs)
            probs = F.softmax(outputs.logits, dim=-1)
            pred_id = probs.argmax().item()
        return jsonify({
            "skin_type": labels[pred_id],
            "confidence": f"{probs[0][pred_id].item():.1%}",
            "dry": f"{probs[0][0].item():.1%}",
            "oily": f"{probs[0][1].item():.1%}"
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500

application = app