|
import json |
|
import os |
|
import torch |
|
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline |
|
|
|
|
|
model_path = "your_model_path" |
|
tokenizer = AutoTokenizer.from_pretrained("OttoYu/Tree-Dbh") |
|
model = AutoModelForSequenceClassification.from_pretrained("OttoYu/Tree-Dbh") |
|
|
|
|
|
text_classification = pipeline( |
|
"text-classification", |
|
model=model, |
|
tokenizer=tokenizer, |
|
device=0 if torch.cuda.is_available() else -1, |
|
return_all_scores=True, |
|
) |
|
|
|
|
|
def predict_tree_properties(dbh): |
|
|
|
input_text = f"dbh: {dbh}" |
|
|
|
|
|
results = text_classification(input_text) |
|
probs = results[0]["scores"] |
|
|
|
|
|
tree_height = probs[0] * 100 |
|
crown_spread = probs[1] * 10 |
|
|
|
|
|
return {"tree_height": tree_height, "crown_spread": crown_spread} |
|
|
|
|
|
def run_inference(): |
|
|
|
dbh = input("Enter the dbh value (in cm): ") |
|
|
|
|
|
tree_properties = predict_tree_properties(dbh) |
|
print(f"Predicted Tree Height: {tree_properties['tree_height']:.2f} m") |
|
print(f"Predicted Crown Spread: {tree_properties['crown_spread']:.2f} m") |
|
|
|
|
|
run_inference() |
|
|