File size: 2,559 Bytes
0dc7e53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | # ================================
# 0. PATCH pour huggingface_hub (contourne l'absence de HfFolder)
# ================================
import huggingface_hub
if not hasattr(huggingface_hub, 'HfFolder'):
class HfFolder:
_token = None
@staticmethod
def get_token():
return HfFolder._token
@staticmethod
def save_token(token):
HfFolder._token = token
huggingface_hub.HfFolder = HfFolder
# ================================
# 1. PATCH pour contourner le bug de Gradio 4.44.0
# (TypeError: argument of type 'bool' is not iterable)
# ================================
import gradio_client.utils
original_get_type = gradio_client.utils.get_type
def patched_get_type(schema):
if isinstance(schema, bool):
return "boolean"
return original_get_type(schema)
gradio_client.utils.get_type = patched_get_type
# ================================
# 2. IMPORTS STANDARDS
# ================================
import gradio as gr
import numpy as np
from PIL import Image
from ultralytics import YOLO
# ================================
# 3. CHARGEMENT DU MODÈLE (YOLOv8 classification)
# ================================
MODEL_PATH = "best.pt" # le modèle doit être présent à la racine
model = YOLO(MODEL_PATH)
# ================================
# 4. FONCTION DE PRÉDICTION
# ================================
def predict(image):
"""
Prend une image PIL, retourne un dictionnaire {classe: probabilité}
"""
# Redimensionner à 224x224 (taille utilisée lors de l'entraînement)
image = image.resize((224, 224))
# Convertir en numpy array
img_np = np.array(image)
# Faire la prédiction avec YOLO
results = model.predict(img_np, verbose=False)
probs = results[0].probs # objet Probs
# Les probabilités sont dans probs.data (tensor)
prob_weed = float(probs.data[1].cpu().numpy()) # classe 1 = weed
prob_crop = float(probs.data[0].cpu().numpy()) # classe 0 = crop
return {"crop": prob_crop, "weed": prob_weed}
# ================================
# 5. INTERFACE GRADIO
# ================================
iface = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil", label="Chargez une image de champ"),
outputs=gr.Label(num_top_classes=2, label="Prédiction"),
title="🌾 Classification Culture / Mauvaise Herbe (YOLOv8)",
description="Modèle YOLOv8-cls entraîné sur le dataset DeepWeeds pour distinguer les cultures des mauvaises herbes."
)
if __name__ == "__main__":
iface.launch(share=True) |