File size: 38,889 Bytes
d4e7f2f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3461e1d
 
 
 
d4e7f2f
43d33f6
d4e7f2f
 
 
 
 
 
 
 
 
 
 
 
 
 
db668b3
d4e7f2f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
249fe7e
 
 
 
 
 
 
d4e7f2f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
# -*- coding: utf-8 -*-
# Author: Gaojian Wang@ZJUICSR
# --------------------------------------------------------
# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
# You can find the license in the LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# pip uninstall nvidia_cublas_cu11


import sys
sys.path.append('..')
import os
os.system(f'pip install dlib')
import torch
import numpy as np
from PIL import Image
import models_mae
from torch.nn import functional as F
import dlib

import gradio as gr


# loading model
model = getattr(models_mae, 'mae_vit_base_patch16')()


class ITEM:
    def __init__(self, img, parsing_map):
        self.image = img
        self.parsing_map = parsing_map

face_to_show = ITEM(None, None)


check_region = {'Eyebrows': [2, 3], 
                'Eyes': [4, 5], 
                'Nose': [6], 
                'Mouth': [7, 8, 9], 
                'Face Boundaries': [10, 1, 0],
                'Hair': [10],
                'Skin': [1],
                'Background': [0]}


def get_boundingbox(face, width, height, minsize=None):
    """
    Expects a dlib face to generate a quadratic bounding box.
    :param face: dlib face class
    :param width: frame width
    :param height: frame height
    :param cfg.face_scale: bounding box size multiplier to get a bigger face region
    :param minsize: set minimum bounding box size
    :return: x, y, bounding_box_size in opencv form
    """
    x1 = face.left()
    y1 = face.top()
    x2 = face.right()
    y2 = face.bottom()
    size_bb = int(max(x2 - x1, y2 - y1) * 1.3)
    if minsize:
        if size_bb < minsize:
            size_bb = minsize
    center_x, center_y = (x1 + x2) // 2, (y1 + y2) // 2

    # Check for out of bounds, x-y top left corner
    x1 = max(int(center_x - size_bb // 2), 0)
    y1 = max(int(center_y - size_bb // 2), 0)
    # Check for too big bb size for given x, y
    size_bb = min(width - x1, size_bb)
    size_bb = min(height - y1, size_bb)

    return x1, y1, size_bb


def extract_face(frame):
    face_detector = dlib.get_frontal_face_detector()
    image = np.array(frame.convert('RGB'))
    faces = face_detector(image, 1)
    if len(faces) > 0:
        # For now only take the biggest face
        face = faces[0]
        # Face crop and rescale(follow FF++)
        x, y, size = get_boundingbox(face, image.shape[1], image.shape[0])
        # Get the landmarks/parts for the face in box d only with the five key points
        cropped_face = image[y:y + size, x:x + size]
        # cropped_face = cv2.resize(cropped_face, (224, 224), interpolation=cv2.INTER_CUBIC)
        return Image.fromarray(cropped_face)
    else:
        return None
        

from torchvision.transforms import transforms
def show_one_img_patchify(img, model): 
    x = torch.tensor(img)

    # make it a batch-like
    x = x.unsqueeze(dim=0)
    x = torch.einsum('nhwc->nchw', x)
    x_patches = model.patchify(x)
    
    # visualize the img_patchify
    n = int(np.sqrt(x_patches.shape[1]))
    image_size = int(224/n)  
    padding = 3
    new_img = Image.new('RGB', (n * image_size + padding*(n-1), n * image_size + padding*(n-1)), 'white')
    for i, patch in enumerate(x_patches[0]):
        ax = i % n
        ay = int(i / n)
        patch_img_tensor = torch.reshape(patch, (model.patch_embed.patch_size[0], model.patch_embed.patch_size[1], 3))
        patch_img_tensor = torch.einsum('hwc->chw', patch_img_tensor)
        patch_img = transforms.ToPILImage()(patch_img_tensor)
        new_img.paste(patch_img, (ax * image_size + padding * ax, ay * image_size + padding * ay))

    new_img = new_img.resize((224, 224), Image.BICUBIC)
    return new_img


def show_one_img_parchify_mask(img, parsing_map, mask, model):
    mask = mask.detach()
    mask_patches = mask.unsqueeze(-1).repeat(1, 1, model.patch_embed.patch_size[0]**2 *3)  # (N, H*W, p*p*3)
    mask = model.unpatchify(mask_patches)  # 1 is removing, 0 is keeping
    mask = torch.einsum('nchw->nhwc', mask).detach().cpu()

    # visualize mask
    vis_mask = mask[0].clone()
    vis_mask[vis_mask == 1] = 1  # gray for masked
    vis_mask[vis_mask == 2] = -1  # black for highlight masked facial region
    vis_mask[vis_mask == 0] = 2  # white for visible
    vis_mask = torch.clip(vis_mask * 127, 0, 255).int()
    fasking_mask = vis_mask.numpy().astype(np.uint8)
    fasking_mask = Image.fromarray(fasking_mask)
    
    # visualize the masked image
    im_masked = img
    im_masked[mask[0] == 1] = 127
    im_masked[mask[0] == 2] = 0
    im_masked = Image.fromarray(im_masked)
    
    # visualize the masked image_patchify
    parsing_map_masked = parsing_map
    parsing_map_masked[mask[0] == 1] = 127
    parsing_map_masked[mask[0] == 2] = 0
    
    return [show_one_img_patchify(parsing_map_masked, model), fasking_mask, im_masked]


# Random
class CollateFn_Random:
    def __init__(self, input_size=224, patch_size=16, mask_ratio=0.75):
        self.img_size = input_size
        self.patch_size = patch_size
        self.num_patches_axis = input_size // patch_size
        self.num_patches = (input_size // patch_size) ** 2
        self.mask_ratio = mask_ratio

    def __call__(self, image, parsing_map):
        random_mask = torch.zeros(parsing_map.size(0), self.num_patches, dtype=torch.float32)  # torch.Size([BS, 14, 14])
        random_mask = self.masking(parsing_map, random_mask)

        return {'image': image, 'random_mask': random_mask}

    def masking(self, parsing_map, random_mask):
        """
        :return:
        """
        for i in range(random_mask.size(0)):
            # normalize the masking to strictly target percentage for batch computation.
            num_mask_to_change = int(self.mask_ratio * self.num_patches)
            mask_change_to = 1 if num_mask_to_change >= 0 else 0
            change_indices = torch.randperm(self.num_patches)
            for idx in range(num_mask_to_change):
                random_mask[i, change_indices[idx]] = mask_change_to

        return random_mask


def do_random_masking(image, parsing_map_vis, ratio):
    img = torch.from_numpy(image)
    img = img.unsqueeze(0).permute(0, 3, 1, 2)
    parsing_map = face_to_show.parsing_map
    parsing_map = torch.tensor(parsing_map)
    
    mask_method = CollateFn_Random(input_size=224, patch_size=16, mask_ratio=ratio)
    mask = mask_method(img, parsing_map)['random_mask']
    
    random_patch_on_parsing, random_mask, random_mask_on_image = show_one_img_parchify_mask(image, parsing_map_vis, mask, model)
    
    return random_patch_on_parsing, random_mask, random_mask_on_image

 
# Fasking
class CollateFn_Fasking:
    def __init__(self, input_size=224, patch_size=16, mask_ratio=0.75):
        self.img_size = input_size
        self.patch_size = patch_size
        self.num_patches_axis = input_size // patch_size
        self.num_patches = (input_size // patch_size) ** 2
        self.mask_ratio = mask_ratio
        # --------------------------------------------------------------------------
        self.facial_region_group = [
            [2, 4],  # right eye
            [3, 5],  # left eye
            [6],  # nose
            [7, 8, 9],  # mouth
            [10],  # hair
            [1],  # skin
            [0]  # background
        ]  # ['background', 'face', 'rb', 'lb', 're', 'le', 'nose', 'ulip', 'imouth', 'llip', 'hair']

    def __call__(self, image, parsing_map):
        # image = torch.stack([sample['image'] for sample in samples])  # torch.Size([bs, 3, 224, 224])
        # parsing_map = torch.stack([sample['parsing_map'] for sample in samples])  # torch.Size([bs, 1, 224, 224])
        # parsing_map = parsing_map.squeeze(1)  # torch.Size([BS, 1, 224, 224]) → torch.Size([BS, 224, 224])

        # random select a facial semantic region and get corresponding mask(masking all patches include this region)
        fasking_mask = torch.zeros(parsing_map.size(0), self.num_patches_axis, self.num_patches_axis, dtype=torch.float32)  # torch.Size([BS, 14, 14])
        fasking_mask = self.fasking(parsing_map, fasking_mask)

        return {'image': image, 'fasking_mask': fasking_mask}

    def fasking(self, parsing_map, fasking_mask):
        """
        :return:
        """
        for i in range(parsing_map.size(0)):
            terminate = False
            for seg_group in self.facial_region_group[:-2]:
                if terminate:
                    break
                for comp_value in seg_group:
                    fasking_mask[i] = torch.maximum(
                        fasking_mask[i], F.max_pool2d((parsing_map[i].unsqueeze(0) == comp_value).float(), kernel_size=self.patch_size))
                    if fasking_mask[i].mean() >= ((self.mask_ratio * self.num_patches) / self.num_patches):
                        terminate = True
                        break

        fasking_mask = fasking_mask.view(parsing_map.size(0), -1)
        for i in range(fasking_mask.size(0)):
            # normalize the masking to strictly target percentage for batch computation.
            num_mask_to_change = (self.mask_ratio * self.num_patches - fasking_mask[i].sum(dim=-1)).int()
            mask_change_to = torch.clamp(num_mask_to_change, 0, 1).item()
            select_indices = (fasking_mask[i] == (1 - mask_change_to)).nonzero(as_tuple=False).view(-1)
            change_indices = torch.randperm(len(select_indices))[:torch.abs(num_mask_to_change)]
            fasking_mask[i, select_indices[change_indices]] = mask_change_to

        return fasking_mask


def do_fasking_masking(image, parsing_map_vis, ratio):
    img = torch.from_numpy(image)
    img = img.unsqueeze(0).permute(0, 3, 1, 2)
    parsing_map = face_to_show.parsing_map
    parsing_map = torch.tensor(parsing_map)
    
    mask_method = CollateFn_Fasking(input_size=224, patch_size=16, mask_ratio=ratio)
    mask = mask_method(img, parsing_map)['fasking_mask']
    
    fasking_patch_on_parsing, fasking_mask, fasking_mask_on_image = show_one_img_parchify_mask(image, parsing_map_vis, mask, model)
    
    return fasking_patch_on_parsing, fasking_mask, fasking_mask_on_image


# FRP
class CollateFn_FR_P_Masking:
    def __init__(self, input_size=224, patch_size=16, mask_ratio=0.75):
        self.img_size = input_size
        self.patch_size = patch_size
        self.num_patches_axis = input_size // patch_size
        self.num_patches = (input_size // patch_size) ** 2
        self.mask_ratio = mask_ratio
        self.facial_region_group = [
            [2, 3],  # eyebrows
            [4, 5],  # eyes
            [6],  # nose
            [7, 8, 9],  # mouth
            [10, 1, 0],  # face boundaries
            [10],  # hair
            [1],  # facial skin
            [0]  # background
        ]  # ['background', 'face', 'rb', 'lb', 're', 'le', 'nose', 'ulip', 'imouth', 'llip', 'hair']

    def __call__(self, image, parsing_map):
        # image = torch.stack([sample['image'] for sample in samples])  # torch.Size([bs, 3, 224, 224])
        # parsing_map = torch.stack([sample['parsing_map'] for sample in samples])  # torch.Size([bs, 1, 224, 224])
        # parsing_map = parsing_map.squeeze(1)  # torch.Size([BS, 1, 224, 224]) → torch.Size([BS, 224, 224])

        # random select a facial semantic region and get corresponding mask(masking all patches include this region)
        P_mask = torch.zeros(parsing_map.size(0), self.num_patches_axis, self.num_patches_axis, dtype=torch.float32)  # torch.Size([BS, 14, 14])
        P_mask = self.random_variable_facial_semantics_masking(parsing_map, P_mask)

        return {'image': image, 'P_mask': P_mask}

    def random_variable_facial_semantics_masking(self, parsing_map, P_mask):
        """
        :return:
        """
        P_mask = P_mask.view(P_mask.size(0), -1)
        for i in range(parsing_map.size(0)):

            for seg_group in self.facial_region_group[:-2]:
                mask_in_seg_group = torch.zeros(1, self.num_patches_axis, self.num_patches_axis, dtype=torch.float32)
                if seg_group == [10, 1, 0]:
                    patch_hair_bg = F.max_pool2d(
                        ((parsing_map[i].unsqueeze(0) == 10) + (parsing_map[i].unsqueeze(0) == 0)).float(),
                        kernel_size=self.patch_size)
                    patch_skin = F.max_pool2d((parsing_map[i].unsqueeze(0) == 1).float(), kernel_size=self.patch_size)
                    # skin&hair or skin&bg defined as facial boundaries:
                    mask_in_seg_group = torch.maximum(mask_in_seg_group,
                                                      (patch_hair_bg.bool() & patch_skin.bool()).float())
                else:
                    for comp_value in seg_group:
                        mask_in_seg_group = torch.maximum(mask_in_seg_group,
                                                          F.max_pool2d(
                                                              (parsing_map[i].unsqueeze(0) == comp_value).float(),
                                                              kernel_size=self.patch_size))

                mask_in_seg_group = mask_in_seg_group.view(-1)
                # to_mask_patches_in_seg_group = mask_in_seg_group - (mask_in_seg_group & P_mask[i])
                to_mask_patches_in_seg_group = (mask_in_seg_group - P_mask[i]) > 0
                mask_num = (mask_in_seg_group.sum(dim=-1) * self.mask_ratio -
                            (mask_in_seg_group.sum(dim=-1)-to_mask_patches_in_seg_group.sum(dim=-1))).int()
                if mask_num > 0:
                    select_indices = (to_mask_patches_in_seg_group == 1).nonzero(as_tuple=False).view(-1)
                    change_indices = torch.randperm(len(select_indices))[:mask_num]
                    P_mask[i, select_indices[change_indices]] = 1

            num_mask_to_change = (self.mask_ratio * self.num_patches - P_mask[i].sum(dim=-1)).int()
            mask_change_to = torch.clamp(num_mask_to_change, 0, 1).item()
            select_indices = (P_mask[i] == (1 - mask_change_to)).nonzero(as_tuple=False).view(-1)
            change_indices = torch.randperm(len(select_indices))[:torch.abs(num_mask_to_change)]
            P_mask[i, select_indices[change_indices]] = mask_change_to

        return P_mask
    

def do_FRP_masking(image, parsing_map_vis, ratio):
    img = torch.from_numpy(image)
    img = img.unsqueeze(0).permute(0, 3, 1, 2)
    parsing_map = face_to_show.parsing_map
    parsing_map = torch.tensor(parsing_map)
    
    mask_method = CollateFn_FR_P_Masking(input_size=224, patch_size=16, mask_ratio=ratio)
    masks = mask_method(img, parsing_map)
    mask = masks['P_mask']
    
    FRP_patch_on_parsing, FRP_mask, FRP_mask_on_image = show_one_img_parchify_mask(image, parsing_map_vis, mask, model)
    
    return FRP_patch_on_parsing, FRP_mask, FRP_mask_on_image


# CRFR_R
class CollateFn_CRFR_R_Masking:
    def __init__(self, input_size=224, patch_size=16, mask_ratio=0.75, region='Nose'):
        self.img_size = input_size
        self.patch_size = patch_size
        self.num_patches_axis = input_size // patch_size
        self.num_patches = (input_size // patch_size) ** 2
        self.mask_ratio = mask_ratio
        self.facial_region_group = [
            [2, 3],  # eyebrows
            [4, 5],  # eyes
            [6],  # nose
            [7, 8, 9],  # mouth
            [10, 1, 0],  # face boundaries
            [10],  # hair
            [1],  # facial skin
            [0]  # background
        ]  # ['background', 'face', 'rb', 'lb', 're', 'le', 'nose', 'ulip', 'imouth', 'llip', 'hair']
        self.random_specific_facial_region = check_region[region]

    def __call__(self, image, parsing_map):
        # mage = torch.stack([sample['image'] for sample in samples])  # torch.Size([bs, 3, 224, 224])
        # parsing_map = torch.stack([sample['parsing_map'] for sample in samples]) # torch.Size([bs, 1, 224, 224])
        # parsing_map = parsing_map.squeeze(1)  # torch.Size([BS, 1, 224, 224]) → torch.Size([BS, 224, 224])

        # random select a facial semantic region and get corresponding mask(masking all patches include this region)
        facial_region_mask = torch.zeros(parsing_map.size(0), self.num_patches_axis, self.num_patches_axis, dtype=torch.float32)  # torch.Size([1, H/P, W/P])
        facial_region_mask, random_specific_facial_region = self.masking_all_patches_in_random_specific_facial_region(parsing_map, facial_region_mask)
        # torch.Size([num_patches,]), list

        CRFR_R_mask, facial_region_mask = self.random_variable_masking(facial_region_mask)
        # torch.Size([num_patches,]), torch.Size([num_patches,])

        return {'image': image, 'CRFR_R_mask': CRFR_R_mask, 'fr_mask': facial_region_mask}

    def masking_all_patches_in_random_specific_facial_region(self, parsing_map, facial_region_mask):
        """
        :param parsing_map: [1, img_size, img_size])
        :param facial_region_mask: [1, num_patches ** .5, num_patches ** .5]
        :return: facial_region_mask, random_specific_facial_region
        """
        # random_specific_facial_region = random.choice(self.facial_region_group[:-2])
        # random_specific_facial_region = [6]  # for test: nose
        if self.random_specific_facial_region == [10, 1, 0]:  # facial boundaries, 10-hair 1-skin 0-background
            # True for hair(10) or bg(0) patches:
            patch_hair_bg = F.max_pool2d(((parsing_map == 10) + (parsing_map == 0)).float(),
                                         kernel_size=self.patch_size)
            # True for skin(1) patches:
            patch_skin = F.max_pool2d((parsing_map == 1).float(), kernel_size=self.patch_size)
            # skin&hair or skin&bg is defined as facial boundaries:
            facial_region_mask = (patch_hair_bg.bool() & patch_skin.bool()).float()
        else:
            for facial_region_index in self.random_specific_facial_region:
                facial_region_mask = torch.maximum(facial_region_mask,
                                                   F.max_pool2d((parsing_map == facial_region_index).float(),
                                                                kernel_size=self.patch_size))

        return facial_region_mask.view(parsing_map.size(0), -1), self.random_specific_facial_region

    def random_variable_masking(self, facial_region_mask):
        CRFR_R_mask = facial_region_mask.clone()

        for i in range(facial_region_mask.size(0)):
            num_mask_to_change = (self.mask_ratio * self.num_patches - facial_region_mask[i].sum(dim=-1)).int()
            mask_change_to = 1 if num_mask_to_change >= 0 else 0

            select_indices = (facial_region_mask[i] == (1 - mask_change_to)).nonzero(as_tuple=False).view(-1)
            change_indices = torch.randperm(len(select_indices))[:torch.abs(num_mask_to_change)]
            CRFR_R_mask[i, select_indices[change_indices]] = mask_change_to

            facial_region_mask[i] = CRFR_R_mask[i] if num_mask_to_change < 0 else facial_region_mask[i]

        return CRFR_R_mask, facial_region_mask
    

def do_CRFR_R_masking(image, parsing_map_vis, ratio, region):
    img = torch.from_numpy(image)
    img = img.unsqueeze(0).permute(0, 3, 1, 2)
    parsing_map = face_to_show.parsing_map
    parsing_map = torch.tensor(parsing_map)
    
    mask_method = CollateFn_CRFR_R_Masking(input_size=224, patch_size=16, mask_ratio=ratio, region=region)
    masks = mask_method(img, parsing_map)
    mask = masks['CRFR_R_mask']
    fr_mask = masks['fr_mask']
    
    CRFR_R_patch_on_parsing, CRFR_R_mask, CRFR_R_mask_on_image = show_one_img_parchify_mask(image, parsing_map_vis, mask+fr_mask, model)
    
    return CRFR_R_patch_on_parsing, CRFR_R_mask, CRFR_R_mask_on_image


# CRFR_P
class CollateFn_CRFR_P_Masking:
    def __init__(self, input_size=224, patch_size=16, mask_ratio=0.75, region='Nose'):
        self.img_size = input_size
        self.patch_size = patch_size
        self.num_patches_axis = input_size // patch_size
        self.num_patches = (input_size // patch_size) ** 2
        self.mask_ratio = mask_ratio

        self.facial_region_group = [
            [2, 3],  # eyebrows
            [4, 5],  # eyes
            [6],  # nose
            [7, 8, 9],  # mouth
            [10, 1, 0],  # face boundaries
            [10],  # hair
            [1],  # facial skin
            [0]  # background
        ]  # ['background', 'face', 'rb', 'lb', 're', 'le', 'nose', 'ulip', 'imouth', 'llip', 'hair']
        self.random_specific_facial_region = check_region[region]

    def __call__(self, image, parsing_map):
        # image = torch.stack([sample['image'] for sample in samples])  # torch.Size([bs, 3, 224, 224])
        # parsing_map = torch.stack([sample['parsing_map'] for sample in samples]) # torch.Size([bs, 1, 224, 224])
        # parsing_map = parsing_map.squeeze(1)  # torch.Size([BS, 1, 224, 224]) → torch.Size([BS, 224, 224])

        # random select a facial semantic region and get corresponding mask(masking all patches include this region)
        facial_region_mask = torch.zeros(parsing_map.size(0), self.num_patches_axis, self.num_patches_axis,
                                         dtype=torch.float32)  # torch.Size([1, H/P, W/P])
        facial_region_mask, random_specific_facial_region = self.masking_all_patches_in_random_specific_facial_region(parsing_map, facial_region_mask)
        # torch.Size([num_patches,]), list

        CRFR_P_mask, facial_region_mask = self.random_variable_masking(parsing_map, facial_region_mask, random_specific_facial_region)
        # torch.Size([num_patches,]), torch.Size([num_patches,])

        return {'image': image, 'CRFR_P_mask': CRFR_P_mask, 'fr_mask': facial_region_mask}

    def masking_all_patches_in_random_specific_facial_region(self, parsing_map, facial_region_mask):
        """
        :param parsing_map: [1, img_size, img_size])
        :param facial_region_mask: [1, num_patches ** .5, num_patches ** .5]
        :return: facial_region_mask, random_specific_facial_region
        """
        # random_specific_facial_region = random.choice(self.facial_region_group[:-2])
        # random_specific_facial_region = [4, 5]   # for test: eyes
        if self.random_specific_facial_region == [10, 1, 0]:  # facial boundaries, 10-hair 1-skin 0-background
            # True for hair(10) or bg(0) patches:
            patch_hair_bg = F.max_pool2d(((parsing_map == 10) + (parsing_map == 0)).float(), kernel_size=self.patch_size)
            # True for skin(1) patches:
            patch_skin = F.max_pool2d((parsing_map == 1).float(), kernel_size=self.patch_size)
            # skin&hair or skin&bg is defined as facial boundaries:
            facial_region_mask = (patch_hair_bg.bool() & patch_skin.bool()).float()

            # # True for hair(10) or skin(1) patches:
            # patch_hair_face = F.max_pool2d(((parsing_map == 10) + (parsing_map == 1)).float(),
            #                                kernel_size=self.patch_size)
            # # True for bg(0) patches:
            # patch_bg = F.max_pool2d((parsing_map == 0).float(), kernel_size=self.patch_size)
            # # skin&bg or hair&bg defined as facial boundaries:
            # facial_region_mask = (patch_hair_face.bool() & patch_bg.bool()).float()

        else:
            for facial_region_index in self.random_specific_facial_region:
                facial_region_mask = torch.maximum(facial_region_mask,
                                                   F.max_pool2d((parsing_map == facial_region_index).float(),
                                                                kernel_size=self.patch_size))

        return facial_region_mask.view(parsing_map.size(0), -1), self.random_specific_facial_region

    def random_variable_masking(self, parsing_map, facial_region_mask, random_specific_facial_region):
        CRFR_P_mask = facial_region_mask.clone()
        other_facial_region_group = [region for region in self.facial_region_group if
                                     region != random_specific_facial_region]
        # print(other_facial_region_group)
        for i in range(facial_region_mask.size(0)):  # iterate each map in BS
            num_mask_to_change = (self.mask_ratio * self.num_patches - facial_region_mask[i].sum(dim=-1)).int()
            # mask_change_to = 1 if num_mask_to_change >= 0 else 0
            mask_change_to = torch.clamp(num_mask_to_change, 0, 1).item()

            # masking patches in other facial regions according to the corresponding ratio
            if mask_change_to == 1:
                # mask_ratio_other_fr = remain(unmasked) patches should be masked / remain(unmasked) patches
                mask_ratio_other_fr = (
                        num_mask_to_change / (self.num_patches - facial_region_mask[i].sum(dim=-1)))

                masked_patches = facial_region_mask[i].clone()
                for other_fr in other_facial_region_group:
                    to_mask_patches = torch.zeros(1, self.num_patches_axis, self.num_patches_axis,
                                                  dtype=torch.float32)
                    if other_fr == [10, 1, 0]:
                        patch_hair_bg = F.max_pool2d(
                            ((parsing_map[i].unsqueeze(0) == 10) + (parsing_map[i].unsqueeze(0) == 0)).float(),
                            kernel_size=self.patch_size)
                        patch_skin = F.max_pool2d((parsing_map[i].unsqueeze(0) == 1).float(), kernel_size=self.patch_size)
                        # skin&hair or skin&bg defined as facial boundaries:
                        to_mask_patches = (patch_hair_bg.bool() & patch_skin.bool()).float()
                    else:
                        for facial_region_index in other_fr:
                            to_mask_patches = torch.maximum(to_mask_patches,
                                                            F.max_pool2d((parsing_map[i].unsqueeze(0) == facial_region_index).float(),
                                                                         kernel_size=self.patch_size))

                    # ignore already masked patches:
                    to_mask_patches = (to_mask_patches.view(-1) - masked_patches) > 0
                    # to_mask_patches = to_mask_patches.view(-1) - (to_mask_patches.view(-1) & masked_patches)
                    select_indices = to_mask_patches.nonzero(as_tuple=False).view(-1)
                    change_indices = torch.randperm(len(select_indices))[
                                     :torch.round(to_mask_patches.sum() * mask_ratio_other_fr).int()]
                    CRFR_P_mask[i, select_indices[change_indices]] = mask_change_to
                    # prevent overlap
                    masked_patches = masked_patches + to_mask_patches.float()

                # mask/unmask patch from other facial regions to get CRFR_P_mask with fixed size
                num_mask_to_change = (self.mask_ratio * self.num_patches - CRFR_P_mask[i].sum(dim=-1)).int()
                # mask_change_to = 1 if num_mask_to_change >= 0 else 0
                mask_change_to = torch.clamp(num_mask_to_change, 0, 1).item()
                # prevent unmasking facial_region_mask
                select_indices = ((CRFR_P_mask[i] + facial_region_mask[i]) == (1 - mask_change_to)).nonzero(as_tuple=False).view(-1)
                change_indices = torch.randperm(len(select_indices))[:torch.abs(num_mask_to_change)]
                CRFR_P_mask[i, select_indices[change_indices]] = mask_change_to

            else:
                # if the num of facial_region_mask is over (num_patches*mask_ratio),
                # unmask it to get CRFR_P_mask with fixed size
                select_indices = (facial_region_mask[i] == (1 - mask_change_to)).nonzero(as_tuple=False).view(-1)
                change_indices = torch.randperm(len(select_indices))[:torch.abs(num_mask_to_change)]
                CRFR_P_mask[i, select_indices[change_indices]] = mask_change_to
                facial_region_mask[i] = CRFR_P_mask[i]

        return CRFR_P_mask, facial_region_mask
    

def do_CRFR_P_masking(image, parsing_map_vis, ratio, region):
    img = torch.from_numpy(image)
    img = img.unsqueeze(0).permute(0, 3, 1, 2)
    parsing_map = face_to_show.parsing_map
    parsing_map = torch.tensor(parsing_map)
    
    mask_method = CollateFn_CRFR_P_Masking(input_size=224, patch_size=16, mask_ratio=ratio, region=region)
    masks = mask_method(img, parsing_map)
    mask = masks['CRFR_P_mask']
    fr_mask = masks['fr_mask']
    
    CRFR_P_patch_on_parsing, CRFR_P_mask, CRFR_P_mask_on_image = show_one_img_parchify_mask(image, parsing_map_vis, mask+fr_mask, model)
    
    return CRFR_P_patch_on_parsing, CRFR_P_mask, CRFR_P_mask_on_image


def vis_parsing_maps(parsing_anno):
    part_colors = [[255, 255, 255],
                   [0, 0, 255], [255, 128, 0], [255, 255, 0],
                   [0, 255, 0], [0, 255, 128],
                   [0, 255, 255], [255, 0, 255], [255, 0, 128],
                   [128, 0, 255], [255, 0, 0]]
    vis_parsing_anno = parsing_anno.copy().astype(np.uint8)
    vis_parsing_anno_color = np.zeros((vis_parsing_anno.shape[0], vis_parsing_anno.shape[1], 3)) + 255

    num_of_class = np.max(vis_parsing_anno)

    for pi in range(1, num_of_class + 1):
        index = np.where(vis_parsing_anno == pi)
        vis_parsing_anno_color[index[0], index[1], :] = part_colors[pi]

    vis_parsing_anno_color = vis_parsing_anno_color.astype(np.uint8)
    return vis_parsing_anno_color


#from facer import facer
import facer
def do_face_parsing(img):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    face_detector = facer.face_detector('retinaface/mobilenet', device=device, threshold=0.3)  # 0.3 for FF++
    face_parser = facer.face_parser('farl/lapa/448', device=device)  # celebm parser
    
    img = extract_face(img)
    with torch.inference_mode():
        img = img.resize((224, 224), Image.BICUBIC)
        image = torch.from_numpy(np.array(img.convert('RGB')))
        image = image.unsqueeze(0).permute(0, 3, 1, 2).to(device=device)
        try:
            faces = face_detector(image)
            faces = face_parser(image, faces)

            seg_logits = faces['seg']['logits']
            seg_probs = seg_logits.softmax(dim=1)  # nfaces x nclasses x h x w
            seg_probs = seg_probs.data  # torch.Size([1, 11, 224, 224])
            parsing = seg_probs.argmax(1)  # [1, 224, 224]

            parsing_map = parsing.data.cpu().numpy()  # [1, 224, 224] int64
            parsing_map = parsing_map.astype(np.int8)  # smaller space
            parsing_map_vis = vis_parsing_maps(parsing_map.squeeze(0))
            
        except KeyError:
            return gr.update()
    
    face_to_show.image = img
    face_to_show.parsing_map = parsing_map
    return img, parsing_map_vis, show_one_img_patchify(parsing_map_vis, model)


# WebUI
with gr.Blocks() as demo:
    # gr.Markdown("<h1 style='text-align: center;'>🧑‍ Visualization Demo of Facial Masking Strategies</h1>")

    gr.HTML("<h1 style='text-align: center;'>🧑‍ Visualization Demo of Facial Masking Strategies</h1>")

    gr.Markdown(
        "This is a demo of visualizing different facial masking strategies that are introduced in [FSFM-3C](https://fsfm-3c.github.io/) for facial masked image modeling (MIM)."
    )
    gr.Markdown(
        "- <b>Random Masking</b>: Random masking all patches."
    )
    gr.Markdown(
        "- <b>Fasking-I</b>: Use a face parser to divide facial regions and priority masking non-skin and non-background regions."
    )
    gr.Markdown(
        "- <b>FRP</b>: Facial Region Proportional masking, which masks an equal portion of patches in each facial region to the overall masking ratio."
    )
    gr.Markdown(
        "- <b>CRFR-R</b>: (1) Covering a Random Facial Region followed by (2) Random masking other patche."
    )
    gr.Markdown(
        "- <b>CRFR-P _(suggested in FSFM-3C)_</b>: (1) Covering a Random Facial Region followed by (2) Proportional masking masking other regions."
    )

    with gr.Column():
        image = gr.Image(label="Upload/Capture/Paste a facial image", type="pil")

        image_submit_btn = gr.Button("🖱️ Face Parsing")
        with gr.Row():  
            ori_image = gr.Image(interactive=False, label="Detected Face")
            parsing_map_vis = gr.Image(interactive=False, label="Face Parsing")
            patch_parsing_map = gr.Image(interactive=False, label="Patchify")
        gr.HTML('<div class="spacer-20"></div>')
    
    with gr.Column():  # Random
        random_submit_btn = gr.Button("🖱️ Random Masking")
        ratio_random = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.5, label="Masking Ratio for Random Masking")
        with gr.Row():
            random_patch_on_parsing = gr.Image(interactive=False, label="Mask/Parsing")
            random_mask = gr.Image(interactive=False, label="Mask")
            random_mask_on_image = gr.Image(interactive=False, label="Masked Face")
        gr.HTML('<div class="spacer-20"></div>')

    with gr.Column():  # Fasking-I
        fasking_submit_btn = gr.Button("🖱️ Fasking-I")
        ratio_fasking = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.5, label="Masking Ratio for Fasking")
        with gr.Row():
            fasking_patch_on_parsing = gr.Image(interactive=False, label="Mask/Parsing")
            fasking_mask = gr.Image(interactive=False, label="Mask")
            fasking_mask_on_image = gr.Image(interactive=False, label="Masked Face")
        gr.HTML('<div class="spacer-20"></div>')
        
    with gr.Column():  # FRP
        FRP_submit_btn = gr.Button("🖱️ FRP")
        ratio_FRP = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.5, label="Masking Ratio for FRP")
        with gr.Row():
            FRP_patch_on_parsing = gr.Image(interactive=False, label="Mask/Parsing")
            FRP_mask = gr.Image(interactive=False, label="Mask")
            FRP_mask_on_image = gr.Image(interactive=False, label="Masked Face")
        gr.HTML('<div class="spacer-20"></div>')
               
    with gr.Column():  # CRFR-R
        CRFR_R_submit_btn = gr.Button("🖱️ CRFR-R")
        ratio_CRFR_R = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.5, label="Masking Ratio for CRFR-R")
        mask_region_CRFR_R = gr.Radio(choices=['Eyebrows', 'Eyes', 'Nose', 'Mouth', 'Face Boundaries', 'Hair','Skin','Background'],
                                      value='Eyes',
                                      label="Facial Region (for CRFR, highlighted by black)")
        with gr.Row():
            CRFR_R_patch_on_parsing = gr.Image(interactive=False, label="Mask/Parsing")
            CRFR_R_mask = gr.Image(interactive=False, label="Mask")
            CRFR_R_mask_on_image = gr.Image(interactive=False, label="Masked Face")
        gr.HTML('<div class="spacer-20"></div>')
        
    with gr.Column():  # CRFR-P
        CRFR_P_submit_btn = gr.Button("🖱️ CRFR-P (suggested in FSFM-3C)")
        ratio_CRFR_P = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.5, label="Masking Ratio for CRFR-P")
        mask_region_CRFR_P = gr.Radio(choices=['Eyebrows', 'Eyes', 'Nose', 'Mouth', 'Face Boundaries', 'Hair', 'Skin', 'Background'],
                                      value='Eyes',
                                      label="Facial Region (for CRFR, highlighted by black)")
        with gr.Row():
            CRFR_P_patch_on_parsing = gr.Image(interactive=False, label="Mask/Parsing")
            CRFR_P_mask = gr.Image(interactive=False, label="Mask")
            CRFR_P_mask_on_image = gr.Image(interactive=False, label="Masked Face")

    gr.HTML(
        '<div style="display: flex; justify-content: center; gap: 20px; margin-bottom: 20px;">'
        '<a href="https://mapmyvisitors.com/web/1bxvj" title="Visit tracker">'
        '<img src="https://mapmyvisitors.com/map.png?d=jXP1NOOyT1tgXoUPspzdoiPssfitXz38c5uReUt9G9M&cl=ffffff">'
        '</a>'
        '</div>'
    )
    
    parseing_map = []
    image_submit_btn.click(
        fn = do_face_parsing,
        inputs=image,
        outputs=[ori_image, parsing_map_vis, patch_parsing_map]
    )
    random_submit_btn.click(
        fn = do_random_masking,
        inputs=[ori_image, parsing_map_vis, ratio_random],
        outputs=[random_patch_on_parsing, random_mask, random_mask_on_image],
    )
    ratio_random.change(
        fn = do_random_masking,
        inputs=[ori_image, parsing_map_vis, ratio_random],
        outputs=[random_patch_on_parsing, random_mask, random_mask_on_image],
    )
    fasking_submit_btn.click(
        fn = do_fasking_masking,
        inputs=[ori_image, parsing_map_vis, ratio_fasking],
        outputs=[fasking_patch_on_parsing, fasking_mask, fasking_mask_on_image],
    )
    ratio_fasking.change(
        fn = do_fasking_masking,
        inputs=[ori_image, parsing_map_vis, ratio_fasking],
        outputs=[fasking_patch_on_parsing, fasking_mask, fasking_mask_on_image],
    )
    FRP_submit_btn.click(
        fn = do_FRP_masking,
        inputs=[ori_image, parsing_map_vis, ratio_FRP],
        outputs=[FRP_patch_on_parsing, FRP_mask, FRP_mask_on_image],
    )
    ratio_FRP.change(
        fn = do_FRP_masking,
        inputs=[ori_image, parsing_map_vis, ratio_FRP],
        outputs=[FRP_patch_on_parsing, FRP_mask, FRP_mask_on_image],
    )
    CRFR_R_submit_btn.click(
        fn = do_CRFR_R_masking,
        inputs=[ori_image, parsing_map_vis, ratio_CRFR_R, mask_region_CRFR_R],
        outputs=[CRFR_R_patch_on_parsing, CRFR_R_mask, CRFR_R_mask_on_image],
    )
    ratio_CRFR_R.change(
        fn = do_CRFR_R_masking,
        inputs=[ori_image, parsing_map_vis, ratio_CRFR_R, mask_region_CRFR_R],
        outputs=[CRFR_R_patch_on_parsing, CRFR_R_mask, CRFR_R_mask_on_image],
    )
    mask_region_CRFR_R.change(
        fn = do_CRFR_R_masking,
        inputs=[ori_image, parsing_map_vis, ratio_CRFR_R, mask_region_CRFR_R],
        outputs=[CRFR_R_patch_on_parsing, CRFR_R_mask, CRFR_R_mask_on_image],
    )
    CRFR_P_submit_btn.click(
        fn = do_CRFR_P_masking,
        inputs=[ori_image, parsing_map_vis, ratio_CRFR_P, mask_region_CRFR_P],
        outputs=[CRFR_P_patch_on_parsing, CRFR_P_mask, CRFR_P_mask_on_image],
    )
    ratio_CRFR_P.change(
        fn=do_CRFR_P_masking,
        inputs=[ori_image, parsing_map_vis, ratio_CRFR_P, mask_region_CRFR_P],
        outputs=[CRFR_P_patch_on_parsing, CRFR_P_mask, CRFR_P_mask_on_image],
    )
    mask_region_CRFR_P.change(
        fn=do_CRFR_P_masking,
        inputs=[ori_image, parsing_map_vis, ratio_CRFR_P, mask_region_CRFR_P],
        outputs=[CRFR_P_patch_on_parsing, CRFR_P_mask, CRFR_P_mask_on_image],
    )

if __name__ == "__main__":
    gr.close_all()
    demo.queue()
    demo.launch()