File size: 2,965 Bytes
d72a22b | 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 77 78 79 80 81 82 83 84 85 | import os
from pathlib import Path
from ultralytics import YOLO
# ==========================================
# USER CONFIGURATION
# ==========================================
# Path to your YOLO model file (.pt)
MODEL_PATH = r"G:\.shortcut-targets-by-id\1ig9YbfY4r3ujzEoaN6_wIKnZPNNel2ma\Mary_Paz\Projeto mestrado cafe (1)\Papers\Article_2_models\03_trained_models\clr_YOLOV8.pt"
# Path to the image file OR the folder containing images you want to process
SOURCE_PATH = r"G:\My Drive\Organic_coffee_sev - Copy\total\IMG_318.jpg"
# Folder where the results will be saved
OUTPUT_DIR = r"G:\.shortcut-targets-by-id\1ig9YbfY4r3ujzEoaN6_wIKnZPNNel2ma\Mary_Paz\Projeto mestrado cafe (1)\Papers\Article_2_models\03_trained_models\figure_example1"
# Confidence threshold for detection (0.0 to 1.0)
CONF_THRESHOLD = 0.5
# ==========================================
def process_images(model_path, source_path, output_dir, conf_threshold):
"""
Run YOLO inference on images.
"""
# 1. Load the model
try:
model = YOLO(model_path)
except Exception as e:
print(f"Error loading model from {model_path}: {e}")
return
# 2. Resolve paths
source = Path(source_path)
output = Path(output_dir)
output.mkdir(parents=True, exist_ok=True)
# 3. Collect images
images_to_process = []
if source.is_file():
images_to_process.append(source)
elif source.is_dir():
# Add common image extensions
extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.webp']
for ext in extensions:
images_to_process.extend(source.glob(ext))
# Also check uppercase extensions
images_to_process.extend(source.glob(ext.upper()))
else:
print(f"Error: Source path '{source}' does not exist.")
return
if not images_to_process:
print(f"No images found at {source}")
return
print(f"Processing {len(images_to_process)} images...")
# 4. Run inference
for img_path in images_to_process:
print(f"Predicting: {img_path.name}")
try:
model.predict(
source=str(img_path),
save=True,
project=str(output),
name="", # Don't create valid/train subfolders, just save to project dir usually
exist_ok=True, # Allow overwriting/appending to existing project dir
conf=conf_threshold
)
except Exception as e:
print(f"Failed to process {img_path.name}: {e}")
print(f"Done! Results saved to {output.resolve()}")
def main():
print("Starting Leaf Extraction...")
print(f"Model: {MODEL_PATH}")
print(f"Source: {SOURCE_PATH}")
print(f"Output: {OUTPUT_DIR}")
process_images(MODEL_PATH, SOURCE_PATH, OUTPUT_DIR, CONF_THRESHOLD)
if __name__ == "__main__":
main()
|