File size: 2,385 Bytes
4b1a7b1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
import os
import shutil
from pycocotools.coco import COCO
from tqdm import tqdm
def prepare_coco_three_classes(coco_json_path, images_dir, output_dir, split):
# Cibles :bagages
target_classes = ['suitcase', "backpack", "handbag"]
# Initialisation COCO
coco = COCO(coco_json_path)
# Récupérer les IDs de catégories pour les classes ciblées
class_id_map = {}
for idx, class_name in enumerate(target_classes):
cat_ids = coco.getCatIds(catNms=[class_name])
if cat_ids:
class_id_map[cat_ids[0]] = idx
print(f"Classes sélectionnées (COCO id → YOLO id) : {class_id_map}")
# Récupérer tous les IDs d’images contenant au moins une des classes cibles
img_ids = set()
for cat_id in class_id_map.keys():
img_ids.update(coco.getImgIds(catIds=[cat_id]))
imgs = coco.loadImgs(list(img_ids))
# Créer les dossiers de sortie
output_images = os.path.join(output_dir, "images", split)
output_labels = os.path.join(output_dir, "labels", split)
os.makedirs(output_images, exist_ok=True)
os.makedirs(output_labels, exist_ok=True)
for img in tqdm(imgs, desc=f"Processing {split}"):
img_path = os.path.join(images_dir, img["file_name"])
if not os.path.exists(img_path):
continue # Éviter les erreurs si l’image n’est pas là
shutil.copy(img_path, os.path.join(output_images, img["file_name"]))
# Récupérer les annotations pour les catégories cibles
ann_ids = coco.getAnnIds(imgIds=img['id'], catIds=list(class_id_map.keys()))
anns = coco.loadAnns(ann_ids)
label_path = os.path.join(output_labels, img["file_name"].replace(".jpg", ".txt"))
with open(label_path, "w") as f:
for ann in anns:
if ann.get("iscrowd", 0):
continue # On évite les annotations crowd
cat_id = ann["category_id"]
yolo_class_id = class_id_map[cat_id]
x, y, w, h = ann["bbox"]
cx = x + w / 2
cy = y + h / 2
cx /= img["width"]
cy /= img["height"]
w /= img["width"]
h /= img["height"]
f.write(f"{yolo_class_id} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}\n")
|