Real-time Object Detection and Classification for Recycling

import os
HOME = os.getcwd()
print(HOME)
!pip install ultralytics==8.0.20

from IPython import display
display.clear_output()

import ultralytics
ultralytics.checks()
from ultralytics import YOLO
from IPython.display import display, Image


!mkdir {HOME}/datasets
%cd {HOME}/datasets

!pip install roboflow --quiet

from roboflow import Roboflow
rf = Roboflow(api_key="fR5rAAgUtxpvu1aEzWsQ")
project = rf.workspace("roboflow-universe-projects").project("taco-object-detection-kcxyn")
dataset = project.version(2).download("yolov8")

%cat {dataset.location}/data.yaml

%cd {HOME}

!yolo task=detect mode=train model=yolov8m.pt data={dataset.location}/data.yaml epochs=30 imgsz=640 plots=True

!ls {HOME}/runs/detect/train/

%cd {HOME}
Image(filename=f'{HOME}/runs/detect/train/confusion_matrix.png', width=600)

%cd {HOME}
Image(filename=f'{HOME}/runs/detect/train/results.png', width=600)

%cd {HOME}
Image(filename=f'{HOME}/runs/detect/train/val_batch0_pred.jpg', width=600)

%cd {HOME}

!yolo task=detect mode=val model={HOME}/runs/detect/train/weights/best.pt data={dataset.location}/data.yaml


%cd {HOME}
!yolo task=detect mode=predict model={HOME}/runs/detect/train/weights/best.pt conf=0.25 source={dataset.location}/test/images save=True


import glob
from IPython.display import Image, display

for image_path in glob.glob(f'{HOME}/runs/detect/predict/*.jpg')[:3]:
      display(Image(filename=image_path, width=600))
      print("\n")
%pwd
%ls

from ultralytics import YOLO

model = YOLO('runs/detect/train/weights/best.pt')

model.export(format='onnx')

import argparse
import torch
import numpy as np
import onnx
import onnxruntime as ort
import vai_q_onnx
from onnxruntime.quantization import QuantFormat, QuantType, CalibrationDataReader
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
import cv2
from pathlib import Path
import os

# Define the paths
script_dir = os.path.dirname(os.path.abspath(__file__))
models_dir = os.path.join(script_dir, 'models')
onnx_model_path = os.path.join(models_dir, 'best.onnx')

class RecyclingDataset(Dataset):
    def __init__(self, data_dir, transform=None):
        self.data_dir = Path(data_dir)
        self.image_dir = self.data_dir / 'images'
        self.label_dir = self.data_dir / 'labels'
        self.transform = transform

        self.image_files = list(self.image_dir.glob('*.jpg'))  # Adjust file extension if needed
        print(f"Found {len(self.image_files)} images in {self.image_dir}")

    def __len__(self):
        return len(self.image_files)

    def __getitem__(self, idx):
        img_path = str(self.image_files[idx])
        label_path = str(self.label_dir / (self.image_files[idx].stem + '.txt'))

        image = cv2.imread(img_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        if self.transform:
            image = self.transform(image)
        
        # Read labels
        labels = self.read_labels(label_path, img_path)
        
        return image, labels
    
    def read_labels(self, label_path, img_path):
        if os.path.exists(label_path):
            try:
                with open(label_path, 'r') as f:
                    lines = f.readlines()
                    labels = []
                    for line in lines:
                        parts = list(map(float, line.strip().split()))
                        if len(parts) >= 6:
                            class_id = int(parts[0])
                            polygon = np.array(parts[1:]).reshape(-1, 2)
                            x_min, y_min = polygon.min(axis=0)
                            x_max, y_max = polygon.max(axis=0)
                            x_center = (x_min + x_max) / 2
                            y_center = (y_min + y_max) / 2
                            width = x_max - x_min
                            height = y_max - y_min
                            labels.append([class_id, x_center, y_center, width, height])
                        else:
                            print(f"Warning: Invalid label format in file {label_path}")
                    labels = np.array(labels)
            except Exception as e:
                print(f"Error reading label file {label_path}: {str(e)}")
                labels = np.zeros((1, 5))  # Placeholder label
        else:
            print(f"Warning: Label file not found for {img_path}")
            labels = np.zeros((1, 5))  # Placeholder label
        
        return labels

def prepare_dataset(data_dir, batch_size=1, quantization=False):
    transform = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize((640, 640)),  # Adjust size as needed
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    actual_batch_size = 1 if quantization else batch_size
    
    # Use a smaller subset of the validation data for calibration
    subset = 'valid' if quantization else 'train'
    
    dataset_path = os.path.join(data_dir, subset)
    print(f"Loading dataset from: {dataset_path}")
    
    dataset = RecyclingDataset(dataset_path, transform=transform)
    
    if len(dataset) == 0:
        raise ValueError(f"No images found in {os.path.join(dataset_path, 'images')}")
    
    dataloader = DataLoader(dataset, batch_size=actual_batch_size, shuffle=not quantization)
    
    print(f"Created dataloader with {len(dataset)} images")
    print(f"Image directory: {dataset.image_dir}")
    print(f"Label directory: {dataset.label_dir}")
    
    # Print first few file names
    print("First few image files:")
    for img_file in dataset.image_files[:5]:
        print(f"  {img_file}")
    
    return dataloader

class YOLOCalibrationDataReader(CalibrationDataReader):
    def __init__(self, data_loader):
        super().__init__()
        self.data_loader = data_loader
        self.iterator = iter(self.data_loader)

    def get_next(self) -> dict:
        try:
            images, _ = next(self.iterator)
            return {"images": images[0].unsqueeze(0).numpy()}
        except StopIteration:
            return None

def yolo_calibration_reader(data_loader):
    return YOLOCalibrationDataReader(data_loader)

def quantize(quantize_loader, model_name):
    print(f"Quantizing {model_name}...")
    onnx_model_path = os.path.join(models_dir, 'best.onnx')
    onnx_model = onnx.load(onnx_model_path)
    onnx.checker.check_model(onnx_model)
    
    input_model_path = onnx_model_path
    output_model_path = os.path.join(models_dir, f"{model_name}_recycling_detection.qdq.U8S8.onnx")
    
    data_reader = yolo_calibration_reader(quantize_loader)
    
    try:
        vai_q_onnx.quantize_static(
            input_model_path,
            output_model_path,
            data_reader,
            quant_format=QuantFormat.QDQ,
            calibrate_method=vai_q_onnx.PowerOfTwoMethod.MinMSE,
            activation_type=QuantType.QUInt8,
            weight_type=QuantType.QInt8,
            enable_ipu_cnn=True,
            extra_options={'ActivationSymmetric': True}
        )
        print(f"Quantized Model Saved at {output_model_path}")
    except Exception as e:
        print(f"Error during quantization: {str(e)}")
        # Add more debug information
        print("Model input shape:", onnx_model.graph.input[0].type.tensor_type.shape)
        print("First batch from data_reader:")
        first_batch = data_reader.get_next()
        if first_batch:
            print("Shape:", first_batch['images'].shape)
            print("Type:", first_batch['images'].dtype)
        else:
            print("No data from data_reader")

def load_quantized_model(model_path):
    session = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])
    return session

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-model", type=str, default='yolo')
    parser.add_argument("--data_dir", type=str, required=True, help="Path to the dataset directory")
    args = parser.parse_args()
    return args

def main():
    args = get_args()
    
    print(f"Data directory: {args.data_dir}")
    print(f"Model name: {args.model}")
    
    try:
        print("Preparing quantization dataset...")
        quantize_loader = prepare_dataset(args.data_dir, quantization=True)
        print("Quantizing model...")
        quantize(quantize_loader, args.model)
    except Exception as e:
        print(f"An error occurred: {str(e)}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()

import cv2
import numpy as np
import onnx
import onnxruntime as ort
from pathlib import Path

# Load the quantized YOLO model
quantized_model_path = r'C:\Users\merte\OneDrive\Desktop\Resnet-v2\models\quantized_model.onnx'
model = onnx.load(quantized_model_path)

# Set up ONNX Runtime session
providers = ['CPUExecutionProvider']
provider_options = [{}]
session = ort.InferenceSession(model.SerializeToString(), providers=providers,
                               provider_options=provider_options)

# Class names for TACO dataset (59 classes)
class_names = ['Aerosol', 'Aluminium blister pack', 'Aluminium foil', 'Battery', 'Broken glass', 
               'Carded blister pack', 'Cigarette', 'Clear plastic bottle', 'Corrugated carton', 
               'Crisp packet', 'Disposable food container', 'Disposable plastic cup', 'Drink can', 
               'Drink carton', 'Egg carton', 'Foam cup', 'Foam food container', 'Food Can', 
               'Food waste', 'Garbage bag', 'Glass bottle', 'Glass cup', 'Glass jar', 
               'Magazine paper', 'Meal carton', 'Metal bottle cap', 'Metal lid', 'Normal paper', 
               'Other carton', 'Other plastic bottle', 'Other plastic container', 'Other plastic cup', 
               'Other plastic wrapper', 'Other plastic', 'Paper bag', 'Paper cup', 'Paper straw', 
               'Pizza box', 'Plastic bottle cap', 'Plastic film', 'Plastic glooves', 'Plastic lid', 
               'Plastic straw', 'Plastic utensils', 'Polypropylene bag', 'Pop tab', 'Rope & strings', 
               'Scrap metal', 'Shoe', 'Single-use carrier bag', 'Six pack rings', 'Spread tub', 
               'Squeezable tube', 'Styrofoam piece', 'Tissues', 'Toilet tube', 'Tupperware', 
               'Unlabeled litter', 'Wrapping paper']

def preprocess_image(image, input_size=(640, 640)):
    original_height, original_width = image.shape[:2]
    
    # Resize and pad image
    ratio = min(input_size[0] / original_width, input_size[1] / original_height)
    new_size = (int(original_width * ratio), int(original_height * ratio))
    resized = cv2.resize(image, new_size, interpolation=cv2.INTER_LINEAR)
    
    padded = np.full((input_size[0], input_size[1], 3), 114, dtype=np.uint8)
    padded[:new_size[1], :new_size[0]] = resized
    
    # Normalize and change to CHW format
    padded = padded.astype(np.float32) / 255.0
    padded = padded.transpose(2, 0, 1)
    
    return np.expand_dims(padded, axis=0), (original_height, original_width), new_size
def postprocess(output, orig_shape, new_size, conf_threshold=0.5, iou_threshold=0.45):
    predictions = np.squeeze(output[0])
    
    # Assuming the output is in the format [x, y, w, h, conf, class_scores]
    num_classes = min(predictions.shape[1] - 5, len(class_names))
    
    # Apply sigmoid to confidence scores and class scores
    scores = 1 / (1 + np.exp(-predictions[:, 4]))
    class_scores = 1 / (1 + np.exp(-predictions[:, 5:5+num_classes]))
    
    # Get predicted classes
    class_ids = np.argmax(class_scores, axis=1)
    
    # Filter based on confidence threshold and valid class range
    mask = (scores > conf_threshold) & (class_ids < num_classes)
    boxes = predictions[mask, :4]
    scores = scores[mask]
    class_ids = class_ids[mask]
    
    # Convert boxes from [x, y, w, h] to [x1, y1, x2, y2]
    boxes[:, 2] = boxes[:, 0] + boxes[:, 2]
    boxes[:, 3] = boxes[:, 1] + boxes[:, 3]
    
    # Apply NMS
    indices = cv2.dnn.NMSBoxes(boxes, scores, conf_threshold, iou_threshold)
    
    if len(indices) > 0:
        indices = indices.flatten()
        boxes = boxes[indices]
        scores = scores[indices]
        class_ids = class_ids[indices]
        
        # Scale boxes back to original image
        input_height, input_width = new_size
        orig_height, orig_width = orig_shape
        
        scale = min(input_width / orig_width, input_height / orig_height)
        offset_x = (input_width - orig_width * scale) / 2
        offset_y = (input_height - orig_height * scale) / 2
        
        boxes[:, [0, 2]] = (boxes[:, [0, 2]] - offset_x) / scale
        boxes[:, [1, 3]] = (boxes[:, [1, 3]] - offset_y) / scale
        
        # Additional filtering based on box size (optional)
        valid_detections = (boxes[:, 2] - boxes[:, 0]) > 20 & (boxes[:, 3] - boxes[:, 1]) > 20
        boxes = boxes[valid_detections]
        scores = scores[valid_detections]
        class_ids = class_ids[valid_detections]
        
        return boxes, scores, class_ids
    else:
        return [], [], []

# Update the draw_detections function to include more information:
def draw_detections(frame, boxes, scores, class_ids):
    for box, score, class_id in zip(boxes, scores, class_ids):
        x1, y1, x2, y2 = box.astype(int)
        
        # Debug print
        print(f"class_id: {class_id}, score: {score:.4f}, box: {box}")
        
        # Check if class_id is valid
        if 0 <= class_id < len(class_names):
            label = f'{class_names[class_id]}: {score:.2f}'
        else:
            label = f'Unknown ({class_id}): {score:.2f}'
        
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        
        # Draw the label background
        label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        cv2.rectangle(frame, (x1, y1 - label_size[1] - 10), (x1 + label_size[0], y1), (0, 255, 0), -1)
        
        # Draw the label text
        cv2.putText(frame, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)

def main():
    cap = cv2.VideoCapture(0)  # Open the default camera 
    
    while True:
        ret, frame = cap.read()
        if not ret:
            print("Failed to capture frame")
            break
    while True:    
        # Preprocess the frame
        input_data, orig_shape, new_size = preprocess_image(frame)
        
        # Run inference
        try:
            outputs = session.run(None, {'images': input_data})
            
            print(f"Model output shape: {outputs[0].shape}")
            print(f"Model output type: {outputs[0].dtype}")
            print(f"Model output min: {outputs[0].min()}, max: {outputs[0].max()}")
            
            boxes, scores, class_ids = postprocess(outputs[0], orig_shape, new_size)
            
            print(f"Number of detections: {len(boxes)}")
            print(f"Class IDs: {class_ids}")
            print(f"Scores: {scores}")
            print(f"Boxes: {boxes}")
            
            draw_detections(frame, boxes, scores, class_ids)
        
        except Exception as e:
            print(f"Error during inference or postprocessing: {e}")
              
        # Display the frame
        cv2.imshow('Recycling Classification', frame)
        
        # Break the loop if 'q' is pressed
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    # Release the capture and close windows
    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()