File size: 13,584 Bytes
8c9048a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
from diffusers import StableDiffusionBrushNetPipeline, BrushNetModel, UniPCMultistepScheduler
import os
import torch
import cv2
import numpy as np
import pandas as pd
from tqdm import tqdm
from PIL import Image



imagelist = pd.read_csv("examples/brushnet/paper_imagelist_for_inpainting.csv").values

def image_resize(image, width = None, height = None, inter = cv2.INTER_AREA):
    # initialize the dimensions of the image to be resized and
    # grab the image size
    dim = None
    (h, w) = image.shape[:2]

    # if both the width and height are None, then return the
    # original image
    if width is None and height is None:
        return image

    # check to see if the width is None
    if width is None:
        # calculate the ratio of the height and construct the
        # dimensions
        r = height / float(h)
        dim = (int(w * r), height)

    # otherwise, the height is None
    else:
        # calculate the ratio of the width and construct the
        # dimensions
        r = width / float(w)
        dim = (width, int(h * r))

    # resize the image
    resized = cv2.resize(image, dim, interpolation = inter)

    # return the resized image
    return resized



# choose the base model here
base_model_path = "data/ckpt/realisticVisionV60B1_v51VAE"
# base_model_path = "runwayml/stable-diffusion-v1-5"

# input brushnet ckpt path
brushnet_path = "data/ckpt/segmentation_mask_brushnet_ckpt"

# choose whether using blended operation
blended = False

occupations = ['backpacker', 'ballplayer', 'bartender', 'basketball_player', 'boatman', 'carpenter', 'cheerleader', 'climber', 'computer_user', 'craftsman', 'dancer', 'disk_jockey', 'doctor', 'drummer', 'electrician', 'farmer', 'fireman', 'flutist', 'gardener', 'guard', 'guitarist', 'gymnast', 'hairdresser', 'horseman', 'judge', 'laborer', 'lawman', 'lifeguard', 'machinist', 'motorcyclist', 'nurse', 'painter', 'patient', 'prayer', 'referee', 'repairman', 'reporter', 'retailer', 'runner', 'sculptor', 'seller', 'singer', 'skateboarder', 'soccer_player', 'soldier', 'speaker', 'student', 'teacher', 'tennis_player', 'trumpeter', 'waiter']
facet = pd.read_csv("../../datasets/facet/annotations/annotations.csv", header=0).rename(columns={'Unnamed: 0': 'sample_idx'}) # Bounding boxes

root = "../../datasets/facet/images_bb"
# mask_root = "../Color-Invariant-Skin-Segmentation/FCN/paper_output/final_skin_mask"
mask_root = "../Color-Invariant-Skin-Segmentation/FCN/output/person_masks"
# mask_root = "../Color-Invariant-Skin-Segmentation/FCN/paper_output/clothes_mask"


# output_dir = "facet_paper_skin_ours"
# output_dir = "facet_paper_clothes_only"
output_dir = "/home/kis/datasets/facet_paper_whole_body_occupation_prompt_filelist/"
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

steps = 10
# conditioning scale
brushnet_conditioning_scale = 1.0
seed = 12345

brushnet = BrushNetModel.from_pretrained(brushnet_path, torch_dtype=torch.float16)
pipe = StableDiffusionBrushNetPipeline.from_pretrained(
    base_model_path, brushnet=brushnet, torch_dtype=torch.float16, low_cpu_mem_usage=False, safety_checker=None, requires_safety_checker=False
)

# speed up diffusion process with faster scheduler and memory optimization
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
# remove following line if xformers is not installed or when using Torch 2.0.
# pipe.enable_xformers_memory_efficient_attention()
# memory optimization.
pipe.enable_model_cpu_offload()


for category in occupations:

    n_imgs = facet[facet['class1'] == category]['person_id'].shape[0]
    
    for id_ in tqdm(range(n_imgs)):
        img = facet[facet['class1'] == category].iloc[id_]
        
        if int(img['visible_face']) != 1:
            continue
        if int(img['gender_presentation_masc']) == 1:
            gender = 'male'
        elif int(img['gender_presentation_fem']) == 1:
            gender = 'female'
        else:
            continue

        if gender == 'male':
            bb = eval(img["bounding_box"])

            input_box = np.array([int(bb['x']), int(bb['y']), int(bb['x'])+int(bb['width']), int(bb['y'])+int(bb['height'])])

            img_id = str(img['filename']).replace(".jpg", "")
            bb_id = str(img['person_id'])


            if not int(bb_id) in imagelist:
                print("file not in the in-painting list, skipping")
                continue

            if os.path.exists(f"{output_dir}/{bb_id}_original.png"):
                print(f"Skipping image {bb_id}: already processed")
                continue

            image_path = f"{root}/{bb_id}.jpg"
            mask_path = f"{mask_root}/{bb_id}.jpg"

            if not os.path.exists(mask_path):
                print(f"No mask found for image id {bb_id}")
                continue

            if not os.path.exists(image_path):
                print(f"No image found for image id {bb_id}")
                continue

            init_image = cv2.imread(image_path)
            (h, w) = init_image.shape[:2]
            if h < 224 or w < 224:
                print(f"Skipping image as it is too small: {h}x{w}")
                continue

            init_image = image_resize(init_image, width=min(512, init_image.shape[1]))
            
            init_image = init_image[:,:,::-1]
            mask_image = 1.*(cv2.resize(cv2.imread(mask_path), (init_image.shape[1], init_image.shape[0])).sum(-1)>255)[:,:,np.newaxis]

            mask = Image.fromarray(mask_image.astype(np.uint8).repeat(3,-1)*255).convert("RGB")
            size = np.asarray(mask).shape
            ret, output = cv2.threshold(np.asarray(mask),127,255,cv2.THRESH_BINARY)
            is_empty = True if np.sum(output > 127) / (size[0]*size[1]) < 0.30  else False
            if is_empty:
                print(f"Skipping image {bb_id}: mask percentage: {np.sum(output > 127) / (size[0]*size[1])}")
                continue

            cv2.imwrite(f"{output_dir}/{bb_id}_original.png", cv2.imread(image_path))

            
            init_image = init_image * (1-mask_image)
            init_image = Image.fromarray(init_image.astype(np.uint8)).convert("RGB")
            mask_image = Image.fromarray(mask_image.astype(np.uint8).repeat(3,-1)*255).convert("RGB")
            cv2.imwrite(f"{output_dir}/{bb_id}_mask.png", np.asarray(mask_image))
        
            generator = torch.Generator("cuda").manual_seed(seed)
            # caption = f"A photo of a woman"# who is a {category}"
            caption = f"A photo of a woman who is a {category}"
            print(f"Image {bb_id}: {caption}")
            image = pipe(
                caption, 
                init_image, 
                mask_image, 
                num_inference_steps=steps, 
                generator=generator,
                brushnet_conditioning_scale=brushnet_conditioning_scale
            ).images[0]

            if blended:
                image_np=np.array(image)
                init_image_np=cv2.imread(image_path)[:,:,::-1]
                mask_np = 1.*(cv2.imread(mask_path).sum(-1)>255)[:,:,np.newaxis]

                # blur, you can adjust the parameters for better performance
                mask_blurred = cv2.GaussianBlur(mask_np*255, (21, 21), 0)/255
                mask_blurred = mask_blurred[:,:,np.newaxis]
                mask_np = 1-(1-mask_np) * (1-mask_blurred)

                image_pasted=init_image_np * (1-mask_np) + image_np*mask_np
                image_pasted=image_pasted.astype(image_np.dtype)
                image=Image.fromarray(image_pasted)

            # image.save(f"examples/brushnet/{output_dir}/{bb_id}_male_to_female.png")
            image.save(f"{output_dir}/{bb_id}_male_to_female.png")


            # caption = f"A photo of a man"# who is a {category}"
            caption = f"A photo of a man who is a {category}"
            print(f"Image {bb_id}: {caption}")
            image = pipe(
                caption, 
                init_image, 
                mask_image, 
                num_inference_steps=steps, 
                generator=generator,
                brushnet_conditioning_scale=brushnet_conditioning_scale
            ).images[0]

            if blended:
                image_np=np.array(image)
                init_image_np=cv2.imread(image_path)[:,:,::-1]
                mask_np = 1.*(cv2.imread(mask_path).sum(-1)>255)[:,:,np.newaxis]

                # blur, you can adjust the parameters for better performance
                mask_blurred = cv2.GaussianBlur(mask_np*255, (21, 21), 0)/255
                mask_blurred = mask_blurred[:,:,np.newaxis]
                mask_np = 1-(1-mask_np) * (1-mask_blurred)

                image_pasted=init_image_np * (1-mask_np) + image_np*mask_np
                image_pasted=image_pasted.astype(image_np.dtype)
                image=Image.fromarray(image_pasted)

            image.save(f"{output_dir}/{bb_id}_male_to_male.png")
        elif gender == "female":
            bb = eval(img["bounding_box"])

            input_box = np.array([int(bb['x']), int(bb['y']), int(bb['x'])+int(bb['width']), int(bb['y'])+int(bb['height'])])

            img_id = str(img['filename']).replace(".jpg", "")
            bb_id = str(img['person_id'])

            if not int(bb_id) in imagelist:
                print("file not in the in-painting list, skipping")
                continue
            if os.path.exists(f"{output_dir}/{bb_id}_original.png"):
                continue

            image_path = f"{root}/{bb_id}.jpg"
            mask_path = f"{mask_root}/{bb_id}.jpg"

            if not os.path.exists(mask_path):
                print(f"No mask found for image id {bb_id}")
                continue

            if not os.path.exists(image_path):
                print(f"Image not found for id {bb_id}")
                continue





            init_image = cv2.imread(image_path)
            (h, w) = init_image.shape[:2]
            if h < 224 or w < 224:
                continue
            init_image = image_resize(init_image, width=min(512, init_image.shape[1]))
            
            init_image = init_image[:,:,::-1]
            mask_image = 1.*(cv2.resize(cv2.imread(mask_path), (init_image.shape[1], init_image.shape[0])).sum(-1)>255)[:,:,np.newaxis]


            mask = Image.fromarray(mask_image.astype(np.uint8).repeat(3,-1)*255).convert("RGB")
            size = np.asarray(mask).shape
            ret, output = cv2.threshold(np.asarray(mask),127,255,cv2.THRESH_BINARY)
            is_empty = True if np.sum(output > 127) / (size[0]*size[1]) < 0.10  else False
            if is_empty:
                print(f"Image {bb_id}: mask percentage: {np.sum(output > 127) / (size[0]*size[1])}")
                continue
        
            cv2.imwrite(f"{output_dir}/{bb_id}_original.png", cv2.imread(image_path))
            init_image = init_image * (1-mask_image)
            init_image = Image.fromarray(init_image.astype(np.uint8)).convert("RGB")
            mask_image = Image.fromarray(mask_image.astype(np.uint8).repeat(3,-1)*255).convert("RGB")
            # cv2.imwrite(f"examples/brushnet/{output_dir}/{bb_id}_mask.png", np.asarray(mask_image))
            cv2.imwrite(f"{output_dir}/{bb_id}_mask.png", np.asarray(mask_image))


            generator = torch.Generator("cuda").manual_seed(seed)
            # caption = f"A photo of a woman"# who is a {category}"
            caption = f"A photo of a woman who is a {category}"
            image = pipe(
                caption, 
                init_image, 
                mask_image, 
                num_inference_steps=steps, 
                generator=generator,
                brushnet_conditioning_scale=brushnet_conditioning_scale
            ).images[0]

            if blended:
                image_np=np.array(image)
                init_image_np=cv2.imread(image_path)[:,:,::-1]
                mask_np = 1.*(cv2.imread(mask_path).sum(-1)>255)[:,:,np.newaxis]

                # blur, you can adjust the parameters for better performance
                mask_blurred = cv2.GaussianBlur(mask_np*255, (21, 21), 0)/255
                mask_blurred = mask_blurred[:,:,np.newaxis]
                mask_np = 1-(1-mask_np) * (1-mask_blurred)

                image_pasted=init_image_np * (1-mask_np) + image_np*mask_np
                image_pasted=image_pasted.astype(image_np.dtype)
                image=Image.fromarray(image_pasted)

            image.save(f"{output_dir}/{bb_id}_female_to_female.png")


            caption = f"A photo of a man who is a {category}"
            image = pipe(
                caption, 
                init_image, 
                mask_image, 
                num_inference_steps=steps, 
                generator=generator,
                brushnet_conditioning_scale=brushnet_conditioning_scale
            ).images[0]

            if blended:
                image_np=np.array(image)
                init_image_np=cv2.imread(image_path)[:,:,::-1]
                mask_np = 1.*(cv2.imread(mask_path).sum(-1)>255)[:,:,np.newaxis]

                # blur, you can adjust the parameters for better performance
                mask_blurred = cv2.GaussianBlur(mask_np*255, (21, 21), 0)/255
                mask_blurred = mask_blurred[:,:,np.newaxis]
                mask_np = 1-(1-mask_np) * (1-mask_blurred)

                image_pasted=init_image_np * (1-mask_np) + image_np*mask_np
                image_pasted=image_pasted.astype(image_np.dtype)
                image=Image.fromarray(image_pasted)

            image.save(f"{output_dir}/{bb_id}_female_to_male.png")