from transformers import AutoImageProcessor, AutoModelForObjectDetection
from PIL import Image
import torch

# Load model and processor
processor = AutoImageProcessor.from_pretrained("MohamedKhayat/fruit-detector-yolos-base")
model = AutoModelForObjectDetection.from_pretrained("MohamedKhayat/fruit-detector-yolos-base")

# Load and process image
image = Image.open("fruit_image.jpg")
inputs = processor(images=image, return_tensors="pt")

# Run inference
with torch.no_grad():
    outputs = model(**inputs)

# Post-process results
target_sizes = torch.tensor([[image.height, image.width]])
results = processor.post_process_object_detection(
    outputs, 
    threshold=0.5, 
    target_sizes=target_sizes
)[0]

for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
    box = box.tolist()
    print(f"Detected {model.config.id2label[label.item()]} with confidence {score:.2f} at {box}")

Training

Base Model: hustvl/yolos-base
Architecture: yolos_base
Task: Object Detection
mAP@50:95 Score: 0.5585
Input Size: 640x640

This model was trained on a custom fruit detection dataset.

Training Repository: transformers-for-fruit-object-detection-internship

License

Apache 2.0