Upload 23 files

d72a22b verified 8 months ago

2.97 kB

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()

	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()