TACO dataset: Exploratie

import glob
import os
import random
from pprint import pprint

import kagglehub
import matplotlib.pyplot as plt
import pandas as pd
import plotly.express as px
import yaml
from PIL import Image, ImageDraw

# Download the dataset
path = kagglehub.dataset_download("vencerlanz09/taco-dataset-yolo-format")

print(path)

with open(path + "/README.dataset.txt") as f:
    print(f.read())

with open(path + "/README.roboflow.txt") as f:
    print(f.read())

with open(path + "/data.yaml") as f:
    meta = yaml.safe_load(f)

# yaml is loaded as a dictionary; using p(retty)print for a nicer output
pprint(meta)

# the path to the images needs to be updated
# os.listdir(meta["train"]) # this generates a FileNotFoundError

meta["train"] = path + "/train/images"
meta["val"] = path + "/valid/images"
meta["test"] = path + "/test/images"

# show the first five image files
os.listdir(meta["train"])[:5]

img_paths = glob.glob(path + "/**/*.jpg", recursive=True)
img_paths[:5]

# create an empty dictionary
# each image contains one or more labeled objects
# each object is labeled with class_id, center_x, center_y, width and height
labels = {
    "img_path": [],
    "subset": [],
    "class_id": [],
    "center_x": [],
    "center_y": [],
    "width": [],
    "height": [],
}

# loop through all the images
for i, img_path in enumerate(img_paths):
    # print some progress info
    if i % 1000 == 0:
        print(f"{i} / {len(img_paths)}")

    # determine whether the image is part of the training, validation or test set
    if "train" in img_path:
        subset = "train"
    elif "valid" in img_path:
        subset = "val"
    elif "test" in img_path:
        subset = "test"
    else:
        raise ValueError(f"Unknown set for image path {img_path}")

    # get the corresponding label file
    label_path = img_path.replace("images", "labels").replace(".jpg", ".txt")
    with open(label_path, "r") as f:
        # each object in the image corresponds to a different line in the label file
        lines = f.readlines()
        for line in lines:
            class_id, center_x, center_y, width, height = line.strip().split()
            labels["img_path"].append(img_path)
            labels["subset"].append(subset)
            labels["class_id"].append(class_id)
            labels["center_x"].append(center_x)
            labels["center_y"].append(center_y)
            labels["width"].append(width)
            labels["height"].append(height)

df = pd.DataFrame(labels)
df.head()

# check the data types
df.info()

# adjust the dtypes
df.class_id = df.class_id.astype(int)
df.center_x = df.center_x.astype(float)
df.center_y = df.center_y.astype(float)
df.width = df.width.astype(float)
df.height = df.height.astype(float)

df.info()

px.pie(df, names="subset", title="Subset distribution")

df["class_name"] = df.class_id.apply(lambda x: meta["names"][x])

px.pie(df, names="class_name", title="Class distribution")

# img_path = random.choice(img_paths)
img_path = path + "/train/images/000087_JPG_jpg.rf.8b6a2fc865187b4fe55927ac0af3e044.jpg"

df.loc[df.img_path == img_path]

def draw_bounding_box(img_path):
    # Load the image using PIL
    img = Image.open(img_path)
    img_width, img_height = img.size

    # Create a drawing object
    draw = ImageDraw.Draw(img)

    for _, row in df.loc[df.img_path == img_path].iterrows():
        # Convert YOLO format (center_x, center_y, width, height) to pixel coordinates (x1, y1, x2, y2)
        x1 = int((row.center_x - row.width / 2) * img_width)
        y1 = int((row.center_y - row.height / 2) * img_height)
        x2 = int((row.center_x + row.width / 2) * img_width)
        y2 = int((row.center_y + row.height / 2) * img_height)

        # Draw the bounding box
        draw.rectangle([x1, y1, x2, y2], outline="green", width=2)

        # Add the class name as text
        text = meta["names"][row.class_id]
        draw.text((x1, y1 - 10), text, fill="white")

    # Display the image
    plt.imshow(img)
    plt.axis("off")
    plt.show()


draw_bounding_box(img_path)

# The bounding box area could be an interesting feature to analyze
df["area"] = df.width * df.height

# check the distribution for the classes with a boxplot
px.box(df, x="class_name", y="area", title="Bounding Box Area Distribution by Class")

# using a log scale for the y-axis
px.box(
    df,
    x="class_name",
    y="area",
    title="Bounding Box Area Distribution by Class (log scale)",
    log_y=True,
)

# display the image with the largest bounding box for a cigarette
cigarette = df[df.class_name == "Cigarette"].sort_values(by="area", ascending=False).iloc[0]
img_path = cigarette.img_path

draw_bounding_box(img_path)

# Select all the images with a Bottle and a Bottle cap label
bottle_img_paths = df[df.class_name == "Bottle"].img_path.unique()
bottle_cap_img_paths = df[df.class_name == "Bottle cap"].img_path.unique()
common_img_paths = set(bottle_img_paths).intersection(set(bottle_cap_img_paths))
print(f"Number of images with both Bottle and Bottle cap: {len(common_img_paths)}")

# Select from those instances where any of the Bottle areas is smaller that the largest Bottle cap area
suspect_bottles_img_paths = []
for img_path in common_img_paths:
    # largest Bottle cap:
    max_cap_area = df[(df.img_path == img_path) & (df.class_name == "Bottle cap")].area.max()
    # any Bottle smaller than that?
    suspect_bottle = df[
        (df.img_path == img_path) & (df.class_name == "Bottle") & (df.area < max_cap_area)
    ]
    if len(suspect_bottle) > 0:
        suspect_bottles_img_paths.append(img_path)

print(
    f"Number of images with a Bottle smaller than the largest Bottle cap: {len(suspect_bottles_img_paths)}"
)

draw_bounding_box(random.choice(suspect_bottles_img_paths))

# Select an images with Bottle caps and no Bottles
bottle_cap_no_bottle_img_paths = set(bottle_cap_img_paths).difference(set(bottle_img_paths))

draw_bounding_box(random.choice(list(bottle_cap_no_bottle_img_paths)))

# Select an images with Bottles and no Bottle caps
bottle_no_bottle_cap_img_paths = set(bottle_img_paths).difference(set(bottle_cap_img_paths))

draw_bounding_box(random.choice(list(bottle_no_bottle_cap_img_paths)))

# Swap Bottle cap and Bottle labels
bottle_cap_idx = meta["names"].index("Bottle cap")
bottle_idx = meta["names"].index("Bottle")
meta["names"][bottle_cap_idx], meta["names"][bottle_idx] = (
    meta["names"][bottle_idx],
    meta["names"][bottle_cap_idx],
)

# Verify the swap
meta["names"]