app.py

"""
Intel Scene Classifier — Flask App
"""

import io
import os
import urllib.request
import numpy as np
from flask import Flask, jsonify, render_template, request
from PIL import Image

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms

app = Flask(__name__)

CLASSES = ["buildings", "forest", "glacier", "mountain", "sea", "street"]
IMG_SIZE = 150

_pytorch_model = None
_tf_model = None


class CNN_Torch(nn.Module):
    def __init__(self, num_classes=6):
        super().__init__()

        self.features = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),

            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Dropout2d(0.1),

            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Dropout2d(0.2),
        )

        self.gap = nn.AdaptiveAvgPool2d(1)

        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(256, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.gap(x)
        x = self.classifier(x)
        return F.log_softmax(x, dim=1)


def load_pytorch():
    global _pytorch_model
    if _pytorch_model is not None:
        return _pytorch_model

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = CNN_Torch(num_classes=6).to(device)

    state_dict = torch.load("parfait_model.pth", map_location=device)
    model.load_state_dict(state_dict)
    model.eval()

    tf_transform = transforms.Compose([
        transforms.Resize((IMG_SIZE, IMG_SIZE)),
        transforms.ToTensor(),
        transforms.Normalize(
            [0.485, 0.456, 0.406],
            [0.229, 0.224, 0.225]
        ),
    ])

    _pytorch_model = (model, device, tf_transform)
    return _pytorch_model


def load_tensorflow():
    global _tf_model
    if _tf_model is None:
        import tensorflow as tf
        _tf_model = tf.keras.models.load_model("parfait_model.keras")
    return _tf_model


def read_input_image():
    if "image" in request.files and request.files["image"].filename:
        return Image.open(io.BytesIO(request.files["image"].read())).convert("RGB")

    image_url = request.form.get("image_url", "").strip()
    if image_url:
        with urllib.request.urlopen(image_url) as response:
            return Image.open(io.BytesIO(response.read())).convert("RGB")

    raise ValueError("No image provided")


@app.route("/")
def index():
    return render_template("index.html")


@app.route("/predict", methods=["POST"])
def predict():
    framework = request.form.get("model", "pytorch")

    try:
        pil_img = read_input_image()
    except Exception:
        return jsonify({"error": "Fichier image invalide"}), 400

    try:
        if framework == "pytorch":
            model, device, tf_transform = load_pytorch()
            tensor = tf_transform(pil_img).unsqueeze(0).to(device)

            with torch.no_grad():
                out = model(tensor)
                probs = torch.exp(out).cpu().numpy()[0]

        elif framework == "tensorflow":
            model = load_tensorflow()
            arr = np.array(
                pil_img.resize((IMG_SIZE, IMG_SIZE)),
                dtype=np.float32
            )
            arr = np.expand_dims(arr, axis=0)
            probs = model.predict(arr, verbose=0)[0]

        else:
            return jsonify({"error": "Framework non supporté"}), 400

        pred_idx = int(np.argmax(probs))
        return jsonify({
            "class": CLASSES[pred_idx],
            "confidence": float(probs[pred_idx]),
            "probabilities": {
                c: float(p) for c, p in zip(CLASSES, probs)
            },
        })

    except FileNotFoundError as e:
        return jsonify({
            "error": f"Modèle introuvable : {e}. Placez les fichiers .pth et .keras à la racine."
        }), 500
    except Exception as e:
        return jsonify({"error": str(e)}), 500


if __name__ == "__main__":
    port = int(os.environ.get("PORT", 5000))
    app.run(host="0.0.0.0", port=port, debug=False)