Spaces:
Sleeping
Sleeping
File size: 1,197 Bytes
86e9c13 |
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 |
# app/models/plant_vision.py
import torch
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForImageClassification
from app.models.llm import explain_species
from app.utils.config import DEVICE, PLANT_MODEL_ID
processor = AutoImageProcessor.from_pretrained(PLANT_MODEL_ID)
model = AutoModelForImageClassification.from_pretrained(
PLANT_MODEL_ID
).to(DEVICE)
model.eval()
@torch.no_grad()
def predict_plant(image: Image.Image):
"""
Returns:
{
"species": str,
"common_name": str | None,
"confidence": float,
"description": str
}
"""
inputs = processor(
images=image,
return_tensors="pt"
).to(DEVICE)
outputs = model(**inputs)
probs = torch.softmax(outputs.logits, dim=-1)
confidence, idx = probs.max(dim=-1)
label = model.config.id2label[idx.item()]
llm_result = explain_species(
species=label,
confidence=confidence.item(),
domain="plant"
)
return {
"species": label,
"common_name": llm_result["common_name"],
"confidence": confidence.item(),
"description": llm_result["description"],
}
|