ArtiAgent - DefectDiffu/src/pipeline/README.md

Image Artifacts Pipeline

A modular pipeline for generating synthetic image artifacts using GSAM (Grounded Segment Anything Model) for part detection and FLUX diffusion model for artifact generation.

Overview

This pipeline provides clean, modular Python components for the two-stage artifact generation workflow:

📦 pipeline/
├── 📄 __init__.py              # Package initialization
├── 📄 data_loader.py           # Dataset handling (COCO, ImageNet, Custom)
├── 📄 gsam_detector.py         # GSAM model integration (GroundingDINO + SAM)
├── 📄 instance_processor.py    # Instance filtering and bbox operations
├── 📄 flux_generator.py        # FLUX model operations
├── 📄 visualization.py         # Image visualization utilities
├── 📄 prompts.py               # OpenAI API utilities for vocabulary generation
└── 📄 README.md               # This file

Architecture

┌─────────────────┐
│  COCO/ImageNet  │
│  Custom Dataset │
└────────┬────────┘
         │
         ▼
┌─────────────────┐
│  Data Loader    │
│  (data_loader)  │
└────────┬────────┘
         │
         ▼
┌─────────────────┐     ┌──────────────┐
│  GSAM Detector  │◄────┤  OpenAI API  │
│  (GroundingDINO │     │  (Vocabulary)│
│   + SAM/SAM-HQ) │     └──────────────┘
└────────┬────────┘
         │
         ▼
┌─────────────────┐
│Instance Processor│
│ (Filter, Sample,│
│  Create Patches)│
└────────┬────────┘
         │
         ▼
┌─────────────────┐
│ FLUX Generator  │
│ (Diffusion with │
│ Patch Guidance) │
└────────┬────────┘
         │
         ▼
┌─────────────────┐
│  Visualizer     │
│ (Results/Masks) │
└─────────────────┘

Components

1. Data Loaders

Handles dataset loading and image sampling for COCO, ImageNet, and custom directories.

from pipeline import COCODataLoader, ImageNetDataLoader, CustomDirectoryDataLoader

# COCO Dataset
coco_loader = COCODataLoader(
    dataset_path="/path/to/coco/annotations",
    image_path="/path/to/coco/images"
)
cat_ids = coco_loader.get_category_ids(['person'])
img_info, img_array, caption = coco_loader.sample_image_by_category(cat_ids)

# ImageNet Dataset
imagenet_loader = ImageNetDataLoader(
    dataset_path="/path/to/imagenet",
    split="train"
)

# Custom Directory
custom_loader = CustomDirectoryDataLoader(
    directory_path="/path/to/images"
)

2. GSAMDetector

Integrates GroundingDINO and SAM/SAM-HQ for part detection with OpenAI-generated vocabulary.

from pipeline import GSAMDetector

detector = GSAMDetector(
    grounding_config="GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py",
    grounding_checkpoint="weight/groundingdino_swint_ogc.pth",
    sam_checkpoint="weight/sam_vit_h_4b8939.pth",
    sam_version="vit_h",
    use_sam_hq=False,
    box_threshold=0.3,
    text_threshold=0.25,
    device='cuda'
)

# Detect parts in image
entity_predictions, subentity_predictions, vis_output = detector.detect_parts(
    image=img_array,
    entities=['person'],
    subentities=['person head', 'person arm', 'person leg'],
    entity_subentity_mapping={'person': ['person head', 'person arm', 'person leg']},
    min_area_ratio=0.005,
    max_area_ratio=0.5
)

3. InstanceProcessor

Handles instance filtering, sampling, bounding box operations, and patch annotation creation.

from pipeline import InstanceProcessor

# Sample instance by confidence score
sampled_instance = InstanceProcessor.sample_instance_by_score(
    predictions, 
    min_area_ratio=0.01, 
    max_area_ratio=0.5
)

# Generate bbox suggestions for addition artifacts
suggested_bbox = InstanceProcessor.generate_bbox_suggestion(
    predictions=predictions,
    reference_bbox=reference_bbox,
    class_name=class_name,
    vocab=vocab,
    max_ref_overlap=0.3,
    min_entity_overlap=0.1
)

# Create annotation with patch indices
annotation_data = InstanceProcessor.create_annotation_dict(
    instance=sampled_instance,
    img_shape=image.shape,
    artifact_type='distortion',
    patch_size=16
)

4. FluxGenerator

Manages FLUX diffusion model operations for artifact generation with patch-based guidance.

from pipeline import FluxGenerator, FluxConfig

# Configure FLUX model
config = FluxConfig(
    name='flux-dev',
    guidance=5.0,
    num_steps=25,
    pe_step=0.5,  # Position encoding step
    seed=42
)

# Artifact-type-specific PE steps
config_advanced = FluxConfig(
    name='flux-dev',
    guidance=5.0,
    num_steps=25,
    pe_step={
        'addition': 0.3,
        'removal': 0.3,
        'distortion': 0.5
    },
    seed=42
)

generator = FluxGenerator(device='cuda', config=config_advanced)

# Generate artifact image
generated_image = generator.generate_with_artifacts(
    source_prompt="a photo of a person",
    target_prompt="a photo of a person",
    bbox=target_bbox,
    bbox_ref=reference_bbox,
    artifact_type='distortion',
    source_img=image
)

5. ImageVisualizer

Provides visualization utilities for debugging and quality assurance.

from pipeline import ImageVisualizer

visualizer = ImageVisualizer()

# Show single image with caption
visualizer.show_image(image, caption, title="Original", base_dir="output/")

# Show comparison
visualizer.show_comparison(
    original_image, 
    generated_image, 
    artifact_data,
    caption="Distortion Artifact",
    base_dir="output/"
)

# Show bounding box overlay
visualizer.show_bbox_overlay(
    image, 
    target_bbox,
    base_dir="output/",
    filename="bbox_overlay.png"
)

# Show patch masks
visualizer.show_patch_masks(
    image,
    reference_patches,
    target_patches,
    base_dir="output/"
)

Artifact Types

The pipeline supports three types of image artifacts:

1. Distortion

Modifies the appearance of existing parts while keeping them in place.

• Uses reference patches to guide where distortion is applied
• Configurable distortion kernels (jitter, swirl, voronoi)

2. Removal

Removes detected parts from images naturally.

• Uses reference patches to identify removal areas
• FLUX inpainting fills removed areas contextually

3. Addition

Adds new instances of detected parts in suitable locations.

• Uses reference patches as source templates
• Generates target patches using IoU-based intelligent placement
• Maintains visual consistency with surrounding context

Configuration

Detection Parameters

Parameter	Default	Description
box_threshold	0.3	Detection confidence threshold
text_threshold	0.25	Text-image matching threshold
min_area_ratio	0.005	Minimum part size (0.5% of image)
max_area_ratio	0.5	Maximum part size (50% of image)
nms_threshold	0.5	Non-maximum suppression threshold

Generation Parameters

Parameter	Default	Description
guidance	5.0	Guidance scale for FLUX
num_steps	25	Number of diffusion steps
pe_step	0.5	Position encoding step size
inject	25	Injection step in diffusion
seed	42	Random seed for reproducibility

Dependencies

• torch - PyTorch for deep learning
• openai - For vocabulary generation
• pycocotools - For COCO dataset handling
• supervision - For detection utilities
• groundingdino - For grounded detection
• segment_anything - For segmentation
• matplotlib - For visualization
• PIL - For image processing
• numpy - For numerical operations

Installation

1. Install GroundingDINO and SAM following their respective installation guides
2. Download model weights and place them in src/weight/ directory
3. Set up OpenAI API key: export OPENAI_API_KEY='your-key'
4. Install required dependencies

See the main README for detailed installation instructions.

Usage

See batch_gsam_segmentation.py and batch_flux_generation.py for complete batch processing examples.

Performance Tips

1. GPU Memory: Use SAM vit_b for limited GPU memory, vit_h for best quality
2. Filtering: Adjust area ratios to balance quality vs. quantity
3. Speed: Lower num_steps (15-20) for faster generation
4. Quality: Higher num_steps (25-35) for better results

Troubleshooting

Common Issues

1. GSAM setup issues: Ensure GroundingDINO and SAM weights are downloaded
2. OpenAI API errors: Check API key and rate limits
3. COCO dataset errors: Verify dataset paths and structure
4. GPU memory issues: Use smaller SAM model or reduce batch size

License

This pipeline integrates multiple open-source components, each with their own licenses. See the main repository LICENSE and model_licenses/ directory for details.