plantAnalyze / backend /main.py
souleyFall's picture
Upload 3 files
bfc61b9 verified
Raw
History Blame Contribute Delete
5.84 kB
print("[DEBUG] Etape 1 : Imports de base...")
import io
import platform
from pathlib import Path
print("[DEBUG] Etape 2 : Patch WMI & Platform...")
# Contournement : sur certaines machines Windows, les requetes WMI de platform.py
# (via uname(), win32_ver(), etc.) sont extremement lentes ou se bloquent,
# ce qui freeze l'import de torch (qui appelle platform.machine() au chargement).
# On mock platform.machine() et platform.uname() pour retourner instantanement
# des valeurs standard et eviter tout appel systeme bloquant.
platform.machine = lambda: "AMD64"
from collections import namedtuple
_uname_fields = ['system', 'node', 'release', 'version', 'machine', 'processor']
_uname_result = namedtuple('uname_result', _uname_fields)
platform.uname = lambda: _uname_result(
system="Windows",
node="localhost",
release="10",
version="10.0.19045",
machine="AMD64",
processor="Intel64"
)
# On mock aussi les autres raccourcis au cas ou
platform.system = lambda: "Windows"
platform.release = lambda: "10"
platform.version = lambda: "10.0.19045"
print("[DEBUG] Etape 3 : Import torch...")
import torch
print("[DEBUG] Etape 4 : Import torch.nn...")
import torch.nn as nn
print("[DEBUG] Etape 5 : Import torchvision...")
from torchvision import models, transforms
print("[DEBUG] Etape 6 : Import PIL...")
from PIL import Image
print("[DEBUG] Etape 7 : Import fastapi...")
from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
print("[DEBUG] Etape 8 : Initialisation FastAPI...")
app = FastAPI(title="Detecteur de maladies de plantes")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
BASE_DIR = Path(__file__).resolve().parent.parent
FRONTEND_DIR = BASE_DIR / "frontend"
print("[DEBUG] Etape 9 : Configuration DEVICE...")
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"[DEBUG] DEVICE selectionne : {DEVICE}")
def _find_model_path():
for candidate in (BASE_DIR / "model" / "plant_model.pth", BASE_DIR / "model.pth"):
if candidate.exists():
return candidate
return BASE_DIR / "model" / "plant_model.pth"
MODEL_PATH = _find_model_path()
print(f"[DEBUG] Chemin du modele trouve : {MODEL_PATH}")
CLASSES = ["Tomato_Bacterial_spot","Tomato_Early_blight","Tomato_healthy","Potato_Early_blight","Potato_healthy"]
CLASS_LABELS = {
"Tomato_Bacterial_spot": "Tomate - Tache bacterienne",
"Tomato_Early_blight": "Tomate - Mildiou precoce",
"Tomato_healthy": "Tomate - Saine",
"Potato_Early_blight": "Pomme de terre - Mildiou precoce",
"Potato_healthy": "Pomme de terre - Saine",
}
def _build_head(in_features, num_classes, state_dict):
# Detecte l'architecture de la tete a partir des cles du checkpoint :
# tete simple (fc.weight) vs tete Sequential (fc.0.weight / fc.3.weight)
if any(k.startswith("fc.0.") or k.startswith("fc.3.") for k in state_dict):
return nn.Sequential(
nn.Linear(in_features, 256),
nn.ReLU(),
nn.Dropout(0.4),
nn.Linear(256, num_classes),
)
return nn.Linear(in_features, num_classes)
def load_model():
print("[DEBUG] Etape 11 : Chargement de ResNet18...")
classes, labels = CLASSES, CLASS_LABELS
m = models.resnet18(weights=None)
print(f"[DEBUG] ResNet18 instancie. Verification du fichier modele : {MODEL_PATH.exists()}")
if MODEL_PATH.exists():
print(f"[DEBUG] Lecture du checkpoint depuis {MODEL_PATH}...")
ckpt = torch.load(MODEL_PATH, map_location=DEVICE)
print("[DEBUG] Checkpoint lu avec succes. Extraction de l'architecture...")
state_dict = ckpt["model_state_dict"]
classes = ckpt.get("classes") or ckpt.get("class_names") or CLASSES
labels = ckpt.get("class_labels", CLASS_LABELS)
print("[DEBUG] Adaptation de la tete de classification...")
m.fc = _build_head(m.fc.in_features, len(classes), state_dict)
print("[DEBUG] Chargement des poids du modele (load_state_dict)...")
m.load_state_dict(state_dict)
print("[DEBUG] Poids charges.")
else:
print("[DEBUG] Aucun fichier modele trouve. Initialisation de la tete par defaut...")
m.fc = nn.Sequential(
nn.Linear(m.fc.in_features, 256),
nn.ReLU(),
nn.Dropout(0.4),
nn.Linear(256, len(CLASSES)),
)
m.eval()
print("[DEBUG] Modele pret et mis en mode eval().")
return m.to(DEVICE), classes, labels
print("[DEBUG] Etape 10 : Lancement de load_model()...")
model, CLASSES, CLASS_LABELS = load_model()
print("[DEBUG] Etape 12 : Modele charge avec succes !")
TRANSFORM = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]),
])
@app.get("/")
def root():
return FileResponse(FRONTEND_DIR / "index.html")
@app.get("/health")
def health():
return {"status": "ok", "device": str(DEVICE), "classes": CLASSES}
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
img = Image.open(io.BytesIO(await file.read())).convert("RGB")
tensor = TRANSFORM(img).unsqueeze(0).to(DEVICE)
with torch.no_grad():
probs = torch.softmax(model(tensor), dim=1)[0]
idx = probs.argmax().item()
pred = CLASSES[idx]
return {
"prediction": CLASS_LABELS.get(pred, pred),
"confidence": round(probs[idx].item(), 4),
"is_healthy": "healthy" in pred.lower(),
"all_probs": {CLASS_LABELS.get(c,c): round(probs[i].item(),4) for i,c in enumerate(CLASSES)},
}
app.mount("/static", StaticFiles(directory=str(FRONTEND_DIR)), name="static")