Upload run_test.py with huggingface_hub

run_test.py

@@ -0,0 +1,271 @@
+import cv2
+import os
+import einops
+import numpy as np
+import torch
+import argparse
+from cldm.model import create_model, load_state_dict
+from cldm.ddim_hacked import DDIMSampler
+from cldm.hack import disable_verbosity, enable_sliced_attention
+from datasets.data_utils import * 
+from omegaconf import OmegaConf
+from tqdm import tqdm
+import albumentations as A
+
+save_memory = False
+disable_verbosity()
+if save_memory:
+    enable_sliced_attention()
+
+config = OmegaConf.load('./configs/inference.yaml')
+model_ckpt =  config.pretrained_model
+model_config = config.config_file
+
+model = create_model(model_config).cpu()
+model.load_state_dict(load_state_dict(model_ckpt, location='cuda'))
+model = model.cuda()
+ddim_sampler = DDIMSampler(model)
+
+
+def get_input(batch, k):
+    x = batch[k]
+    if len(x.shape) == 3:
+        x = x[None, ...]
+
+    x = torch.tensor(x)
+    x = einops.rearrange(x, 'b h w c -> b c h w')
+    x = x.to(memory_format=torch.contiguous_format).float()
+    return x
+
+def get_unconditional_conditioning(N, obj_thr):
+    x = [torch.zeros((1, 3, 224, 224)).to(model.device)] * N
+    single_uc = model.get_learned_conditioning(x)
+    uc = single_uc.unsqueeze(-1).repeat(1, 1, 1, obj_thr)
+    return {"pch_code": uc}
+
+def inference(item, back_image):
+    obj_thr = 2
+    num_samples = 1
+    H, W = 512, 512
+    guidance_scale = 5.0
+    
+    # 1. Condition & Mask Extraction
+    xc = []
+    xc_mask = []
+    for i in range(obj_thr):
+        xc.append(get_input(item, f"view{i}").cuda())
+        xc_mask.append(get_input(item, f"mask{i}"))
+
+    # 2. Cross-Attention Condition (pch_code)
+    c_list = [model.get_learned_conditioning(xc_i) for xc_i in xc]
+    c_tensor = torch.stack(c_list).permute(1, 2, 3, 0) # [B, Tokens, Dim, Obj]
+    cond_cross = {"pch_code": c_tensor}
+
+    # 3. Mask Condition
+    c_mask = torch.stack(xc_mask).permute(1, 2, 3, 4, 0) # Align with BasicTransformerBlock
+
+    # 4. ControlNet / Concat Condition
+    hint = item['hint']
+    control = torch.from_numpy(hint.copy()).float().cuda()
+    control = torch.stack([control] * num_samples, dim=0)
+    control = einops.rearrange(control, 'b h w c -> b c h w').clone()
+
+    # 5. Build Final Condition Dictionaries
+    cond = {
+        "c_concat": [control], 
+        "c_crossattn": [cond_cross], 
+        "c_mask": [c_mask]
+    }
+    
+    # Correctly unwrap the UC dictionary
+    uc_pch = get_unconditional_conditioning(num_samples, obj_thr)
+    un_cond = {
+        "c_concat": [control], 
+        "c_crossattn": [uc_pch], 
+        "c_mask": [c_mask]
+    }
+
+    # 6. Sampling
+    if save_memory:
+        model.low_vram_shift(is_diffusing=True)
+
+    shape = (4, H // 8, W // 8)
+    model.control_scales = [1.0] * 13
+    
+    samples, _ = ddim_sampler.sample(
+        50, num_samples, shape, cond, 
+        verbose=False, eta=0.0,
+        unconditional_guidance_scale=guidance_scale,
+        unconditional_conditioning=un_cond
+    )
+
+    if save_memory:
+        model.low_vram_shift(is_diffusing=False)
+
+    # 7. Post-processing
+    x_samples = model.decode_first_stage(samples)
+    x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy()
+    
+    pred = np.clip(x_samples[0], 0, 255).astype(np.uint8)
+    
+    # Resize and crop
+    side = max(back_image.shape[0], back_image.shape[1])
+    pred = cv2.resize(pred, (side, side))
+    pred = crop_back(pred, back_image, item['extra_sizes'], item['hint_sizes0'], item['hint_sizes1'], is_masked=True) 
+    
+    return pred
+
+
+def process_pairs_multiple(mask, tar_image, patch_dir, counter=0, max_ratio=0.8):
+    # 1. Process Reference Object (View)
+    view = cv2.imread(patch_dir)
+    view = cv2.cvtColor(view, cv2.COLOR_BGR2RGB)
+    view = pad_to_square(view, pad_value=255, random=False)
+    view = cv2.resize(view.astype(np.uint8), (224, 224))
+    view = view.astype(np.float32) / 255.0
+
+    # 2. BBox and Mask Logic
+    box_yyxx = get_bbox_from_mask(mask)
+    
+    # Define crop area (using full image here)
+    H1, W1 = tar_image.shape[0], tar_image.shape[1]
+    box_yyxx_crop = [0, H1, 0, W1]
+    
+    # Handle box within crop
+    y1, y2, x1, x2 = box_in_box(box_yyxx, box_yyxx_crop)
+
+    # 3. Create Collage (Input Hint)
+    # Background with hole (zeroed out at object position)
+    collage = tar_image.copy()
+    source_collage = collage.copy()
+    collage[y1:y2, x1:x2, :] = 0
+
+    # Binary mask for the current object hole
+    collage_mask = np.zeros_like(tar_image, dtype=np.float32)
+    collage_mask[y1:y2, x1:x2, :] = 1.0
+
+    # 4. Square Padding & Resizing
+    # Pad all to square (pad_value 2 for mask indicates padding area)
+    tar_square = pad_to_square(tar_image, pad_value=0, random=False)
+    collage_square = pad_to_square(collage, pad_value=0, random=False)
+    mask_square = pad_to_square(collage_mask, pad_value=2, random=False)
+    
+    H2, W2 = collage_square.shape[0], collage_square.shape[1]
+
+    # Resize to model input size
+    tar_res = cv2.resize(tar_square, (512, 512)).astype(np.float32)
+    col_res = cv2.resize(collage_square, (512, 512)).astype(np.float32)
+    mask_res = cv2.resize(mask_square, (512, 512), interpolation=cv2.INTER_NEAREST).astype(np.float32)
+
+    # 5. Mask Value Normalization
+    # Original logic: mask=1 for object, 0 for background, -1 for padding
+    mask_res[mask_res == 2] = -1
+    
+    # For conditioning: keep a 0/1 version for cross-attn mask
+    c_mask = np.where(mask_res[..., 0:1] == 1, 1.0, 0.0).astype(np.float32)
+
+    # 6. Final Item Assembly
+    # Normalize images to [-1, 1]
+    tar_res = tar_res / 127.5 - 1.0
+    col_res = col_res / 127.5 - 1.0
+    
+    # Hint: Concatenate background with the (-1, 0, 1) mask
+    hint_final = np.concatenate([col_res, mask_res[..., :1]], axis=-1)
+
+    item = {
+        f'view{counter}': view,
+        f'hint{counter}': hint_final,
+        f'mask{counter}': c_mask,
+        f'hint_sizes{counter}': np.array([y1, x1, y2, x2]),
+        'jpg': tar_res, # Targets are same for all counters in a pair
+        'collage': source_collage,
+        'extra_sizes': np.array([H1, W1, H2, W2])
+    }
+
+    return item
+
+
+def process_composition(item, obj_thr):
+    collage = item['collage'].copy()
+    collage_mask = np.zeros((collage.shape[0], collage.shape[1], 1), dtype=np.float32)
+
+    for i in reversed(range(obj_thr)):
+        y1, x1, y2, x2 = item['hint_sizes'+str(i)]
+        collage[y1:y2, x1:x2, :] = 0
+        collage_mask[y1:y2,x1:x2,:] = 1.0
+
+    collage = pad_to_square(collage, pad_value = 0, random = False).astype(np.uint8)
+    
+    collage_mask = pad_to_square(collage_mask, pad_value = 2, random = False).astype(np.float32)
+
+    collage = cv2.resize(collage.astype(np.uint8), (512, 512)).astype(np.float32) / 127.5 - 1.0 
+    collage_mask = cv2.resize(collage_mask, (512, 512), interpolation=cv2.INTER_NEAREST).astype(np.float32)
+    
+    if len(collage_mask.shape) == 2: 
+        collage_mask = collage_mask[..., None]
+    
+    collage_mask[collage_mask == 2] = -1.0
+
+    collage_final = np.concatenate([collage, collage_mask[:,:,:1]] , -1)
+    
+    item.update({'hint': collage_final.copy()})
+    return item
+
+def run_inference(input_dir, output_dir, sample_num=31, obj_thr=2):
+    """
+    Core inference loop for multi-object composition.
+    """
+    os.makedirs(output_dir, exist_ok=True)
+    comp_image_dir = os.path.join(output_dir, 'composed')
+    os.makedirs(comp_image_dir, exist_ok=True)
+
+    img_ids = sorted([d for d in os.listdir(input_dir) if os.path.isdir(os.path.join(input_dir, d))])
+
+    for img_id in tqdm(img_ids, desc="Processing images"):
+        img_folder = os.path.join(input_dir, img_id)
+        img_path = os.path.join(img_folder, 'image.jpg')
+        
+        if not os.path.exists(img_path):
+            continue
+
+        # 1. Load background image
+        back_image = cv2.imread(img_path)
+        back_image = cv2.cvtColor(back_image, cv2.COLOR_BGR2RGB)
+
+        # 2. Iteratively process multiple objects
+        item_with_collage = {}
+        for j in range(obj_thr):
+        # for j in reversed(range(obj_thr)):
+            patch_path = os.path.join(img_folder, f"object_{j}.png")
+            mask_path = os.path.join(img_folder, f"object_{j}_mask.png")
+            
+            if not (os.path.exists(patch_path) and os.path.exists(mask_path)):
+                print(f"Warning: Object {j} missing in {img_id}")
+                continue
+
+            tar_mask = (cv2.imread(mask_path)[:, :, 0] > 128).astype(np.uint8)
+            
+            # Pass counter=j to ensure keys like 'view0', 'view1' are unique
+            item = process_pairs_multiple(tar_mask, back_image, patch_path, counter=j)
+            item_with_collage.update(item)
+
+        # 3. Composition & Model Prediction
+        # Ensure process_composition merges 'hint0', 'hint1' into a single 'hint'
+        item_with_collage = process_composition(item_with_collage, obj_thr)
+        
+        # Using inference_single_image_multi as defined previously
+        gen_image = inference(item_with_collage, back_image)
+        
+        # 4. Save result
+        save_name = f'composed_{img_id}.png'
+        cv2.imwrite(os.path.join(comp_image_dir, save_name), gen_image[:, :, ::-1])
+    
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input', type=str, help='Input data directory')
+    parser.add_argument('--output', type=str, help='Output save directory')
+    parser.add_argument('--obj_thr', type=int, default=2, help='Number of objects to compose')
+    args = parser.parse_args()
+    
+    run_inference(args.input, args.output, obj_thr=args.obj_thr)