Views
No views yet
1graph TD
2 A[Input Image] --> B[YOLOv8n Detector]
3 B --> C[Detected Strawberries]
4 C --> D[Crop & Resize]
5 D --> E[EfficientNet-B0 Classifier]
6 E --> F[Ripeness Prediction]
7 F --> G[Decision: Pick Only Ripe]
8
9 style A fill:#f9f9f9
10 style B fill:#e3f2fd
11 style E fill:#fff3e0
12 style G fill:#c8e6c9| Component | Model | Architecture | Performance | Size | Purpose |
|---|---|---|---|---|---|
| Detection | YOLOv8n | Object Detection | mAP@50: 83.07% | 6.2MB | Locate strawberries |
| Classification | EfficientNet-B0 | Image Classification | Accuracy: 91.94% | 56MB | Classify ripeness |
| Model | Architecture | Performance | Size | Best For |
|---|---|---|---|---|
| YOLOv8n | YOLOv8 Nano | mAP@50: 98.9% | 5.7MB | Edge deployment, real-time |
| YOLOv8s | YOLOv8 Small | mAP@50: 93.7% | 21MB | Higher accuracy applications |
| YOLOv11n | YOLOv11 Nano | Testing | 10.4MB | Latest architecture testing |
1# Clone repository
2git clone https://github.com/theonegareth/strawberryPicker.git
3cd strawberryPicker
4
5# Install dependencies
6pip install -r requirements.txt1from huggingface_hub import hf_hub_download
2
3# Download detection model
4detector_path = hf_hub_download(
5 repo_id="theonegareth/strawberryPicker",
6 filename="detection/best.pt"
7)
8
9# Download classification model
10classifier_path = hf_hub_download(
11 repo_id="theonegareth/strawberryPicker",
12 filename="classification/best_ripeness_classifier.pth"
13)
14
15print(f"Models downloaded to:\n- {detector_path}\n- {classifier_path}")1import torch
2import cv2
3from PIL import Image
4from torchvision import transforms
5import numpy as np
6
7# Load detection model
8detector = torch.hub.load('ultralytics/yolov8', 'custom', path=detector_path)
9
10# Load classification model
11classifier = torch.load(classifier_path, map_location='cpu')
12classifier.eval()
13
14# Preprocessing for classifier
15transform = transforms.Compose([
16 transforms.Resize((128, 128)),
17 transforms.ToTensor(),
18 transforms.Normalize(mean=[0.485, 0.456, 0.406],
19 std=[0.229, 0.224, 0.225])
20])
21
22# Process image
23def detect_and_classify(image_path):
24 """
25 Detect strawberries and classify their ripeness
26
27 Args:
28 image_path: Path to input image
29
30 Returns:
31 results: List of dicts with bbox, ripeness, confidence
32 """
33 # Load image
34 image = cv2.imread(image_path)
35 image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
36
37 # Detect strawberries
38 detection_results = detector(image_rgb)
39
40 results = []
41 for result in detection_results:
42 boxes = result.boxes.xyxy.cpu().numpy()
43 confidences = result.boxes.conf.cpu().numpy()
44 class_ids = result.boxes.cls.cpu().numpy()
45
46 for box, conf, cls_id in zip(boxes, confidences, class_ids):
47 if conf < 0.5: # Filter low confidence detections
48 continue
49
50 x1, y1, x2, y2 = map(int, box)
51
52 # Crop strawberry
53 crop = image_rgb[y1:y2, x1:x2]
54 if crop.size == 0:
55 continue
56
57 # Classify ripeness
58 crop_pil = Image.fromarray(crop)
59 input_tensor = transform(crop_pil).unsqueeze(0)
60
61 with torch.no_grad():
62 output = classifier(input_tensor)
63 probabilities = torch.softmax(output, dim=1)
64 predicted_class = torch.argmax(probabilities, dim=1).item()
65 confidence = probabilities[0][predicted_class].item()
66
67 # Ripeness classes
68 classes = ['unripe', 'partially-ripe', 'ripe', 'overripe']
69
70 results.append({
71 'bbox': (x1, y1, x2, y2),
72 'ripeness': classes[predicted_class],
73 'confidence': confidence,
74 'detection_confidence': float(conf),
75 'detection_class': int(cls_id)
76 })
77
78 return results
79
80# Example usage
81if __name__ == "__main__":
82 image_path = "strawberries.jpg"
83 results = detect_and_classify(image_path)
84
85 print(f"Detected {len(results)} strawberries:")
86 for i, result in enumerate(results, 1):
87 print(f" {i}. Ripeness: {result['ripeness']} "
88 f"(conf: {result['confidence']:.2f})")strawberryPicker/
├── detection/ # YOLOv8n detection model (Two-stage system)
│ ├── best.pt # PyTorch weights
│ └── README.md # Model documentation
├── classification/ # EfficientNet-B0 classification model (Two-stage system)
│ ├── best_ripeness_classifier.pth # PyTorch weights
│ ├── training_summary.md
│ └── README.md # Model documentation
├── yolov8n/ # YOLOv8 Nano model (98.9% mAP@50)
│ ├── best.pt # PyTorch weights
│ ├── best.onnx # ONNX format
│ ├── best_fp16.onnx # FP16 ONNX for edge deployment
│ └── README.md # Model documentation
├── yolov8s/ # YOLOv8 Small model (93.7% mAP@50)
│ ├── best.pt # PyTorch weights
│ ├── strawberry_yolov8s_enhanced.pt # Enhanced version
│ └── README.md # Model documentation
├── yolov11n/ # YOLOv11 Nano model (Testing)
│ ├── strawberry_yolov11n.pt # PyTorch weights
│ ├── strawberry_yolov11n.onnx # ONNX format
│ └── README.md # Model documentation
├── scripts/ # Optimization scripts
├── benchmark_results/ # Performance benchmarks
├── results/ # Training results/plots
├── LICENSE # MIT license
├── CITATION.cff # Academic citation
├── sync_to_huggingface.py # Automation script
├── requirements.txt # Python dependencies
├── inference_example.py # Basic inference script
├── webcam_inference.py # Real-time webcam demo
└── README.md # This file1# Pseudo-code for robotics integration
2for strawberry in detected_strawberries:
3 if strawberry.ripeness == 'ripe':
4 robot_arm.move_to(strawberry.position)
5 robot_arm.pick()1# Conveyor belt sorting
2if ripeness == 'ripe':
3 conveyor.route_to('premium_package')
4elif ripeness == 'partially-ripe':
5 conveyor.route_to('delayed_shipping')
6else:
7 conveyor.route_to('rejection_bin')1# Track ripeness distribution over time
2daily_ripeness_counts = analyze_temporal_ripeness(images_over_time)
3optimal_harvest_day = find_peak_ripe_day(daily_ripeness_counts)1@misc{strawberryPicker2024,
2 title={Strawberry Picker AI System: A Two-Stage Approach for Automated Harvesting},
3 author={The One Gareth},
4 year={2024},
5 publisher={HuggingFace},
6 url={https://huggingface.co/theonegareth/strawberryPicker}
7}