File size: 3,696 Bytes
c8c00f0 | 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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | # scripts/convert_roboflow_to_triplets.py
import json
import shutil
from pathlib import Path
from PIL import Image
import numpy as np
def convert_roboflow_to_triplets(roboflow_dir: str, output_dir: str):
"""
Convert Roboflow COCO format to patch triplet format.
Roboflow COCO structure:
train/
_annotations.coco.json
image1.jpg
image2.jpg
"""
roboflow_path = Path(roboflow_dir)
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
# Load COCO annotations
anno_file = roboflow_path / "train" / "_annotations.coco.json"
with open(anno_file) as f:
coco = json.load(f)
# Build image ID → filename map
images = {img['id']: img['file_name'] for img in coco['images']}
# Group annotations by image
from collections import defaultdict
image_annotations = defaultdict(list)
for ann in coco['annotations']:
image_annotations[ann['image_id']].append(ann)
triplet_id = 0
for img_id, filename in images.items():
img_path = roboflow_path / "train" / filename
if not img_path.exists():
continue
img = Image.open(img_path).convert('RGB')
img_array = np.array(img)
H, W = img_array.shape[:2]
anns = image_annotations.get(img_id, [])
if not anns:
continue
for ann in anns:
# 1. Skip non-defect category annotations (category_id == 2)
if ann.get('category_id') == 2:
continue
# 2. Skip full-image bounding boxes
x, y, w, h = ann['bbox']
if w >= W and h >= H:
continue
triplet_id += 1
triplet_dir = output_path / f"defect_{triplet_id:04d}"
triplet_dir.mkdir(exist_ok=True)
# Get bbox
x, y, w, h = ann['bbox']
x1, y1, x2, y2 = int(x), int(y), int(x+w), int(y+h)
# Create mask from segmentation if available, else bbox
if 'segmentation' in ann and ann['segmentation']:
# COCO polygon segmentation
from pycocotools import mask as maskUtils
rles = maskUtils.frPyObjects(ann['segmentation'], H, W)
mask = maskUtils.decode(rles)
if len(mask.shape) == 3:
mask = np.any(mask, axis=2).astype(np.uint8) * 255
else:
# Fallback: bbox mask
mask = np.zeros((H, W), dtype=np.uint8)
mask[y1:y2, x1:x2] = 255
# Save clean image (original_target)
img.save(triplet_dir / "original_target.png")
# Save mask (original_masked)
Image.fromarray(mask).save(triplet_dir / "original_masked.png")
# For artifact_target, we need the defective version.
# Since this is a real defect dataset, the original image IS the defect.
# For synthetic training, you may want to inpaint the defect out to create "clean",
# but for RAG retrieval, we can use the same image as artifact_target.
img.save(triplet_dir / "artifact_target.png")
print(f"Created triplet {triplet_id}: {filename} → {triplet_dir}")
print(f"\nTotal triplets created: {triplet_id}")
print(f"Output: {output_path}")
if __name__ == "__main__":
convert_roboflow_to_triplets(
roboflow_dir="./data/external/Manufacturing_Defect_Detection",
output_dir="data/external/roboflow_manufacturing"
) |