cyberai-1
Update tensorflow
ac05c13
raw
history blame
6.5 kB
"""
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)