libcom 0.1.8__py3-none-any.whl

libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/utils.py

@@ -0,0 +1,72 @@
+# -*- coding: utf-8 -*-
+import cv2
+import numpy as np
+
+skeleton = [[15, 13], [13, 11], [16, 14], [14, 12], [11, 12], [5, 11], [6, 12], [5, 6], [5, 7], [6, 8], [7, 9], [8, 10],
+            [1, 2], [0, 1], [0, 2], [1, 3], [2, 4], [3, 5], [4, 6]]
+
+pose_kpt_color = [[51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255], [0, 255, 0],
+                  [255, 128, 0], [0, 255, 0], [255, 128, 0], [0, 255, 0], [255, 128, 0], [0, 255, 0], [255, 128, 0],
+                  [0, 255, 0], [255, 128, 0], [0, 255, 0], [255, 128, 0]]
+
+pose_link_color = [[0, 255, 0], [0, 255, 0], [255, 128, 0], [255, 128, 0],
+                   [51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255], [0, 255, 0], [255, 128, 0],
+                   [0, 255, 0], [255, 128, 0], [51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255],
+                   [51, 153, 255], [51, 153, 255], [51, 153, 255]]
+
+
+def imshow_keypoints(img,
+                     pose_result,
+                     kpt_score_thr=0.1,
+                     radius=2,
+                     thickness=2):
+    """Draw keypoints and links on an image.
+
+    Args:
+            img (ndarry): The image to draw poses on.
+            pose_result (list[kpts]): The poses to draw. Each element kpts is
+                a set of K keypoints as an Kx3 numpy.ndarray, where each
+                keypoint is represented as x, y, score.
+            kpt_score_thr (float, optional): Minimum score of keypoints
+                to be shown. Default: 0.3.
+            thickness (int): Thickness of lines.
+    """
+
+    img_h, img_w, _ = img.shape
+    img = np.zeros(img.shape)
+
+    for idx, kpts in enumerate(pose_result):
+        if idx > 1:
+            continue
+        kpts = kpts['keypoints']
+        # print(kpts)
+        kpts = np.array(kpts, copy=False)
+
+        # draw each point on image
+        assert len(pose_kpt_color) == len(kpts)
+
+        for kid, kpt in enumerate(kpts):
+            x_coord, y_coord, kpt_score = int(kpt[0]), int(kpt[1]), kpt[2]
+
+            if kpt_score < kpt_score_thr or pose_kpt_color[kid] is None:
+                # skip the point that should not be drawn
+                continue
+
+            color = tuple(int(c) for c in pose_kpt_color[kid])
+            cv2.circle(img, (int(x_coord), int(y_coord)), radius, color, -1)
+
+        # draw links
+
+        for sk_id, sk in enumerate(skeleton):
+            pos1 = (int(kpts[sk[0], 0]), int(kpts[sk[0], 1]))
+            pos2 = (int(kpts[sk[1], 0]), int(kpts[sk[1], 1]))
+
+            if (pos1[0] <= 0 or pos1[0] >= img_w or pos1[1] <= 0 or pos1[1] >= img_h or pos2[0] <= 0
+                    or pos2[0] >= img_w or pos2[1] <= 0 or pos2[1] >= img_h or kpts[sk[0], 2] < kpt_score_thr
+                    or kpts[sk[1], 2] < kpt_score_thr or pose_link_color[sk_id] is None):
+                # skip the link that should not be drawn
+                continue
+            color = tuple(int(c) for c in pose_link_color[sk_id])
+            cv2.line(img, pos1, pos2, color, thickness=thickness)
+
+    return img

libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/fusion/dual_encoder_fusion.py

@@ -0,0 +1,93 @@
+import torch
+from torch import nn, einsum
+from libcom.painterly_image_harmonization.source.PHDiffusion.ldm.modules.attention import CrossAttention
+from inspect import isfunction
+import math
+import torch
+import torch.nn.functional as F
+from torch import nn, einsum
+from einops import rearrange, repeat
+from typing import Optional, Any
+from libcom.painterly_image_harmonization.source.PHDiffusion.ldm.modules.diffusionmodules.util import checkpoint
+import numpy
+
+def Normalize(in_channels):
+    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+# our dual encoder fusion
+class CrossAttentionInteraction(nn.Module):
+    def __init__(self, in_channels, n_heads=8, d_head=64,
+                 dropout=0.):
+        super().__init__()
+
+        inner_dim = n_heads * d_head
+
+        self.norm = Normalize(inner_dim)
+
+        self.crossAtt_1 = CrossAttention(query_dim=inner_dim, heads=n_heads, dim_head=d_head, dropout=dropout,
+                                         context_dim=inner_dim)
+
+        self.crossAtt_2 = CrossAttention(query_dim=inner_dim, heads=n_heads, dim_head=d_head, dropout=dropout,
+                                         context_dim=inner_dim)
+
+        self.fc = nn.Conv1d(in_channels=inner_dim * 2, out_channels=inner_dim, kernel_size=1)
+
+        self.conv1 = nn.Conv2d(in_channels=inner_dim, out_channels=inner_dim, kernel_size=1)
+
+    def downsample(self, image_tensor, width, height):
+        image_upsample_tensor = torch.nn.functional.interpolate(image_tensor, size=[width, height])
+        image_upsample_tensor = image_upsample_tensor.clamp(0, 1)
+        return image_upsample_tensor
+
+    def forward(self, adapter_feature, unet_feature, fg_mask):
+        ori_adapter_feature = adapter_feature
+
+        b, c, h, w = adapter_feature.shape
+        fg_mask = self.downsample(fg_mask, h, w)
+
+        original_h = adapter_feature.shape[-1]
+
+        adapter_feature = self.norm(adapter_feature)
+
+        fg = adapter_feature * fg_mask
+
+        fg = rearrange(fg, 'b c h w -> b c (h w)').contiguous()
+
+        fg = rearrange(fg, 'b c h -> b h c').contiguous()
+
+        bg = adapter_feature * (1 - fg_mask)
+
+        bg = rearrange(bg, 'b c h w -> b (h w) c').contiguous()
+
+        adapter_feature = self.crossAtt_1(fg, bg, mask=fg_mask, is_foreground='ada')
+
+        adapter_feature = adapter_feature.permute(0, 2, 1)
+
+        unet_feature = self.norm(unet_feature)
+
+        unet_fg = unet_feature * fg_mask
+        unet_fg = rearrange(unet_fg, 'b c h w -> b c (h w)').contiguous()
+
+        unet_fg = rearrange(unet_fg, 'b c h -> b h c').contiguous()
+
+        unet_bg = unet_feature * (1 - fg_mask)
+        unet_bg = rearrange(unet_bg, 'b c h w -> b (h w) c').contiguous()
+        unet_feature = self.crossAtt_2(unet_fg, unet_bg, mask=fg_mask, is_foreground='unet')
+
+        unet_feature = unet_feature.permute(0, 2, 1)
+
+        interact_feature = self.fc(torch.cat([adapter_feature, unet_feature], dim=1))  # 1 640 mm -> 1 320 mm
+
+        interact_feature = interact_feature.repeat(1, 1, int(original_h * original_h / interact_feature.shape[-1]))
+
+        b, c, h = interact_feature.shape
+
+        interact_feature = interact_feature.reshape(b, c, original_h, original_h)
+
+        interact_feature = self.conv1(interact_feature)
+
+        new_adapter_feature = interact_feature * fg_mask + ori_adapter_feature * (1 - fg_mask)
+
+        return new_adapter_feature
+

libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/image_degradation/__init__.py

@@ -0,0 +1,2 @@
+from libcom.painterly_image_harmonization.source.ldm.modules.image_degradation.bsrgan import degradation_bsrgan_variant as degradation_fn_bsr
+from libcom.painterly_image_harmonization.source.ldm.modules.image_degradation.bsrgan_light import degradation_bsrgan_variant as degradation_fn_bsr_light