Transfer Learning

def read_data_yaml(yaml_path):
    with open(yaml_path, 'r') as stream:
        data_config = yaml.safe_load(stream)
    return data_config.get("names", [])

data_yaml_path = "waste_detection/data.yaml"
class_names = read_data_yaml(data_yaml_path)
print("Class names from data.yaml:", class_names)

def count_images_in_folder(base_path):
    splits = ['train', 'valid', 'test']
    image_counts = {}

    for split in splits:
        image_folder = os.path.join(base_path, split, 'images')
        if os.path.exists(image_folder):
            image_counts[split] = len([
                f for f in os.listdir(image_folder)
                if f.lower().endswith(('.jpg', '.png', '.jpeg'))
            ])
    return image_counts

base_path = 'waste_detection'
results = count_images_in_folder(base_path)
for split, count in results.items():
    print(f"Number of images in {split}: {count}")

def display_random_images_from_folder(folder_path, image_extensions=None, num_images=4):
    if image_extensions is None:
        image_extensions = ['.jpg', '.jpeg', '.png']

    image_files = [f for f in os.listdir(folder_path)
                   if os.path.splitext(f)[1].lower() in image_extensions]

    selected_files = random.sample(image_files, min(num_images, len(image_files)))

    plt.figure(figsize=(10, 8))
    for i, filename in enumerate(selected_files):
        image_path = os.path.join(folder_path, filename)
        image = cv2.imread(image_path)
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        plt.subplot(2, 2, i + 1)
        plt.imshow(image_rgb)
        plt.title(filename, fontsize=10)
        plt.axis('off')

    plt.tight_layout()
    plt.show()

display_random_images_from_folder("/content/waste_detection/train/images")

def polygon_to_bbox(polygon):
    xs = polygon[0::2]
    ys = polygon[1::2]
    min_x = min(xs)
    max_x = max(xs)
    min_y = min(ys)
    max_y = max(ys)

    x_center = (min_x + max_x) / 2.0
    y_center = (min_y + max_y) / 2.0
    width = max_x - min_x
    height = max_y - min_y
    return x_center, y_center, width, height

def convert_polygon_to_yolo_lines(input_file):
    yolo_lines = []
    with open(input_file, 'r') as f_in:
        for line in f_in:
            line = line.strip()
            if not line:
                continue
            tokens = line.split()
            cls = tokens[0]
            coords = list(map(float, tokens[1:]))
            x_center, y_center, width, height = polygon_to_bbox(coords)
            yolo_line = f"{cls} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}"
            yolo_lines.append(yolo_line)
    return yolo_lines

def visualize_polygon_annotations(image_path, label_path, class_list):
    image = cv2.imread(image_path)
    height, width = image.shape[:2]
    yolo_boxes = convert_polygon_to_yolo_lines(label_path)

    for yolo_line in yolo_boxes:
        tokens = yolo_line.strip().split()
        cls_idx = int(tokens[0])
        cls_name = class_list[cls_idx]
        x_center, y_center, w, h = map(float, tokens[1:])

        x_center *= width
        y_center *= height
        w *= width
        h *= height

        x1 = int(x_center - w / 2)
        y1 = int(y_center - h / 2)
        x2 = int(x_center + w / 2)
        y2 = int(y_center + h / 2)

        cv2.rectangle(image, (x1, y1), (x2, y2), color=(0, 255, 0), thickness=2)
        cv2.putText(image, cls_name, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 0), 2)

    cv2_imshow(image)

visualize_polygon_annotations(image_path, label_path, class_names)

from ultralytics import YOLO

model = YOLO('yolov8n.pt')
run = model.train(
    data=str(DATA_YAML),
    epochs=20,
    imgsz=640,
    batch=16,
    project='runs',
    name='exp_epochs_20'
)

val_results = best_model.val(data=str(DATA_YAML), split='val')
print(f"VAL mAP@0.5    = {val_results.box.map50:.3f}")
print(f"VAL mAP@0.5:0.95 = {val_results.box.map:.3f}")

test_results = best_model.val(data=str(DATA_YAML), split='test')
print(f"TEST mAP@0.5    = {test_results.box.map50:.3f}")

results = best_model.predict(img_path, conf=0.25, verbose=False)
boxes = results[0].boxes.xyxy.cpu().numpy()
classes = results[0].boxes.cls.cpu().numpy()

for box, cls in zip(boxes, classes):
    x1, y1, x2, y2 = map(int, box)
    cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), 2)
    cv2.putText(img, class_names[int(cls)], (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)