# Download the dataset from Roboflow
from roboflow import Roboflow
import os
from shutil import copy2

# Roboflow API key and dataset details
rf = Roboflow(api_key="cmb2Ynt8m0CEOg1yl5fU")
project = rf.workspace("nandan-gowda-7qop8").project("prod-final")
version = project.version(1)
dataset = version.download("yolov8")

# Paths
dataset_path = 'Prod-final-1'  # Update the folder name if necessary
filtered_path = 'datasets/imagenet'

# Create directories for filtered data
os.makedirs(f'{filtered_path}/train/images', exist_ok=True)
os.makedirs(f'{filtered_path}/train/labels', exist_ok=True)
os.makedirs(f'{filtered_path}/val/images', exist_ok=True)
os.makedirs(f'{filtered_path}/val/labels', exist_ok=True)

# Filter images with 'book' class
for annotation_file in os.listdir(f'{dataset_path}/train/labels'):
    with open(f'{dataset_path}/train/labels/{annotation_file}') as f:
        lines = f.readlines()
        if any('0 ' in line for line in lines):  # Assuming 'book' class is labeled as '0'
            copy2(f'{dataset_path}/train/labels/{annotation_file}', f'{filtered_path}/train/labels/')
            image_file = annotation_file.replace('.txt', '.jpg')
            copy2(f'{dataset_path}/train/images/{image_file}', f'{filtered_path}/train/images/')

for annotation_file in os.listdir(f'{dataset_path}/valid/labels'):
    with open(f'{dataset_path}/valid/labels/{annotation_file}') as f:
        lines = f.readlines()
        if any('0 ' in line for line in lines):  # Assuming 'book' class is labeled as '0'
            copy2(f'{dataset_path}/valid/labels/{annotation_file}', f'{filtered_path}/val/labels/')
            image_file = annotation_file.replace('.txt', '.jpg')
            copy2(f'{dataset_path}/valid/images/{image_file}', f'{filtered_path}/val/images/')

print("Filtered dataset created at:", filtered_path)