cocoa-diseases-detection/api.py

# from fastapi import FastAPI, File, UploadFile, Form, HTTPException
# from pydantic import BaseModel
# from typing import Optional
# import uuid
# import io
# import cv2
# import numpy as np
# from helper import load_model, show_model_not_loaded_warning, model_path
# import logging
# import os


# logger = logging.getLogger(__name__)

# app = FastAPI()

# class_names = {0: "Fitoftora", 1: "Monilia", 2: "Sana", 3: "Healthy"}

# class PredictionInfo(BaseModel):
#     nb_classe: int
#     nb_box: int
#     confidence_min: float
#     pred_classes: set
#     bbx_coordinates: list[tuple[float, float, float, float, float]]
#     pred_img_id: uuid.UUID
#     box_id: uuid.UUID
#     diseases_id: uuid.UUID
#     det_image_path: str

# def predict_image(confidence: float, image_bytes: bytes, model):
#     try:
#         nparr = np.frombuffer(image_bytes, np.uint8)
#         image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

#         if image is None:
#             raise HTTPException(
#                 status_code=400,
#                 detail="L'image est vide. Assurez-vous de télécharger une image valide."
#             )

        
#         # Convert image to RGB
#         image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)


#         res = model.predict(image_rgb, conf=confidence)
#         nb_classe = len(set(res[0].boxes.cls.tolist()))
#         nb_box = len(res[0].boxes)
#         confidence_min = min(res[0].boxes.conf.tolist())
#         pred_classes = set([class_names[int(box.cls.item())] for box in res[0].boxes])
#         pred_img_id = uuid.uuid4()
#         box_id = uuid.uuid4()
#         diseases_id = uuid.uuid4()


#         # Normalisation des coordonnées des boîtes englobantes
#         bbx_coordinates = [
#             (box.xyxy[0][0], box.xyxy[0][1],
#             box.xyxy[0][2], box.xyxy[0][3], box.conf)
#             for box in res[0].boxes]



#         # Define image saving path for debugging (optional)
#         current_directory = os.getcwd()
#         IMAGE_SAVE_PATH = os.path.join(current_directory, "images_bbx")
#         os.makedirs(IMAGE_SAVE_PATH, exist_ok=True)
#         image_with_boxes_path = os.path.join(IMAGE_SAVE_PATH, f"imagebbx_{pred_img_id}.jpg")
#         res_plotted = res[0].plot()[:, :, ::-1]
#         cv2.imwrite(image_with_boxes_path, res_plotted)

#         return PredictionInfo(
#             pred_img_id=pred_img_id,
#             diseases_id=diseases_id,
#             box_id=box_id,
#             nb_classe=nb_classe,
#             nb_box=nb_box,
#             confidence_min=confidence_min,
#             pred_classes=pred_classes,
#             bbx_coordinates=bbx_coordinates,
#             det_image_path=image_with_boxes_path
#         )
#     except Exception as e:
#         logger.exception("Une erreur est survenue pendant la prédiction : %s", e)
#         raise HTTPException(status_code=500, detail="Une erreur interne est survenue.")

# @app.post("/predict_image/")
# async def predict_image_route(confidence: float = Form(...), image: UploadFile = File(...)):
#     try:
#         model = load_model(model_path)
#         image_bytes = await image.read()
#         return predict_image(confidence, image_bytes, model)
#     except HTTPException as e:
#         raise e
#     except Exception as e:
#         logger.exception("Une erreur interne est survenue lors de la prédiction d'image.")
#         raise HTTPException(status_code=500, detail="Une erreur interne est survenue.")




from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from pydantic import BaseModel
from typing import Optional
import uuid
import io
import cv2
import numpy as np
from helper import load_model, show_model_not_loaded_warning, model_path
import logging
import os
import urllib.parse
from starlette.responses import FileResponse

logger = logging.getLogger(__name__)

app = FastAPI()

class_names = {0: "Fitoftora", 1: "Monilia", 2: "Sana", 3: "Healthy"}

class PredictionInfo(BaseModel):
    nb_classe: int
    nb_box: int
    confidence_min: float
    pred_classes: set
    bbx_coordinates: list[tuple[float, float, float, float, float]]
    pred_img_id: uuid.UUID
    box_id: uuid.UUID
    diseases_id: uuid.UUID
    det_image_path: str

def predict_image(confidence: float, image_bytes: bytes, model):
    try:
        nparr = np.frombuffer(image_bytes, np.uint8)
        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

        if image is None:
            raise HTTPException(
                status_code=400,
                detail="L'image est vide. Assurez-vous de télécharger une image valide."
            )

        # Convertir l'image en RGB
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        res = model.predict(image_rgb, conf=confidence)
        nb_classe = len(set(res[0].boxes.cls.tolist()))
        nb_box = len(res[0].boxes)
        confidence_min = min(res[0].boxes.conf.tolist())
        pred_classes = set([class_names[int(box.cls.item())] for box in res[0].boxes])
        pred_img_id = uuid.uuid4()
        box_id = uuid.uuid4()
        diseases_id = uuid.uuid4()

        # Normaliser les coordonnées des boîtes englobantes
        bbx_coordinates = [
            (box.xyxy[0][0], box.xyxy[0][1],
            box.xyxy[0][2], box.xyxy[0][3], box.conf)
            for box in res[0].boxes]

        # Définir le chemin d'enregistrement de l'image pour le débogage (optionnel)
        current_directory = os.getcwd()
        IMAGE_SAVE_PATH = os.path.join(current_directory, "images_bbx")
        os.makedirs(IMAGE_SAVE_PATH, exist_ok=True)
        image_with_boxes_path = os.path.join(IMAGE_SAVE_PATH, f"imagebbx_{pred_img_id}.jpg")
        res_plotted = res[0].plot()[:, :, ::-1]
        cv2.imwrite(image_with_boxes_path, res_plotted)

        return PredictionInfo(
            pred_img_id=pred_img_id,
            diseases_id=diseases_id,
            box_id=box_id,
            nb_classe=nb_classe,
            nb_box=nb_box,
            confidence_min=confidence_min,
            pred_classes=pred_classes,
            bbx_coordinates=bbx_coordinates,
            det_image_path=image_with_boxes_path
        )
    except Exception as e:
        logger.exception("Une erreur est survenue pendant la prédiction : %s", e)
        raise HTTPException(status_code=500, detail="Une erreur interne est survenue.")

@app.post("/predict_image/")
async def predict_image_route(confidence: float = Form(...), image: UploadFile = File(...)):
    try:
        model = load_model(model_path)
        image_bytes = await image.read()
        prediction_info = predict_image(confidence, image_bytes, model)

        # Créer le lien de téléchargement pour l'image avec les boîtes englobantes
        image_download_link = f"http://0.0.0.0:8001/download_image?image_path={urllib.parse.quote(prediction_info.det_image_path)}"

        return {
            "nb_classe": prediction_info.nb_classe,
            "nb_box": prediction_info.nb_box,
            "confidence_min": prediction_info.confidence_min,
            "pred_classes": prediction_info.pred_classes,
            "bbx_coordinates": prediction_info.bbx_coordinates,
            "pred_img_id": prediction_info.pred_img_id,
            "box_id": prediction_info.box_id,
            "diseases_id": prediction_info.diseases_id,
            "image_download_link": image_download_link
        }
    except HTTPException as e:
        raise e
    except Exception as e:
        logger.exception("Une erreur interne est survenue lors de la prédiction d'image.")
        raise HTTPException(status_code=500, detail="Une erreur interne est survenue.")

@app.get("/download_image/")
async def download_image(image_path: str):
    try:
        return FileResponse(image_path)
    except Exception as e:
        logger.exception("Une erreur est survenue lors du téléchargement de l'image : %s", e)
        raise HTTPException(status_code=500, detail="Une erreur interne est survenue lors du téléchargement de l'image.")