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 tensorflow as tf
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from tensorflow.keras import layers, models

app = Flask(__name__)

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

_pytorch_model = None
_tf_model = None


class CNN_Torch(nn.Module):
    """
    CNN amélioré 4 blocs
    Entrée : (B, 3, 150, 150)
    Sortie : (B, num_classes)
    """

    def __init__(self, num_classes=6):
        super().__init__()

        self.features = nn.Sequential(
            # Block 1: 150 -> 75
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),

            nn.Conv2d(32, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),

            nn.MaxPool2d(2),

            # Block 2: 75 -> 37
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),

            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),

            nn.MaxPool2d(2),
            nn.Dropout2d(0.10),

            # Block 3: 37 -> 18
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),

            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),

            nn.MaxPool2d(2),
            nn.Dropout2d(0.15),

            # Block 4: 18 -> 9
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),

            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),

            nn.MaxPool2d(2),
            nn.Dropout2d(0.20),
        )

        self.gap = nn.AdaptiveAvgPool2d(1)

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

    def forward(self, x):
        x = self.features(x)
        x = self.gap(x)
        x = self.classifier(x)
        return x


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 build_cnn_tf(num_classes: int = 6, input_shape: tuple = (228, 228, 3)):
    return models.Sequential([
        layers.Input(shape=input_shape),

        layers.Conv2D(32, (5, 5), activation="relu"),
        layers.MaxPooling2D(2, 2),

        layers.Conv2D(32, (5, 5), activation="relu"),
        layers.MaxPooling2D(2, 2),

        layers.Conv2D(32, (3, 3), activation="relu"),
        layers.MaxPooling2D(2, 2),

        layers.Conv2D(64, (3, 3), activation="relu"),
        layers.MaxPooling2D(2, 2),

        layers.Conv2D(64, (3, 3), activation="relu"),
        layers.MaxPooling2D(2, 2),

        layers.Flatten(),
        layers.Dense(1024, activation="relu"),
        layers.Dropout(0.20),
        layers.Dense(124, activation="relu"),
        layers.Dropout(0.20),
        layers.Dense(num_classes, activation="softmax"),
    ])


def load_tensorflow():
    global _tf_model
    if _tf_model is not None:
        return _tf_model

    model = build_cnn_tf(num_classes=6, input_shape=(228, 228, 3))
    model.load_weights("parfait_model.keras")

    _tf_model = model
    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.softmax(out, dim=1).cpu().numpy()[0]
                #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]),
            "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)