Add custom nodes, Civitai loras (LFS), and vast.ai setup script
Some checks failed
Python Linting / Run Ruff (push) Has been cancelled
Python Linting / Run Pylint (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.10, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.11, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.12, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-unix-nightly (12.1, , linux, 3.11, [self-hosted Linux], nightly) (push) Has been cancelled
Execution Tests / test (macos-latest) (push) Has been cancelled
Execution Tests / test (ubuntu-latest) (push) Has been cancelled
Execution Tests / test (windows-latest) (push) Has been cancelled
Test server launches without errors / test (push) Has been cancelled
Unit Tests / test (macos-latest) (push) Has been cancelled
Unit Tests / test (ubuntu-latest) (push) Has been cancelled
Unit Tests / test (windows-2022) (push) Has been cancelled
Some checks failed
Python Linting / Run Ruff (push) Has been cancelled
Python Linting / Run Pylint (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.10, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.11, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.12, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-unix-nightly (12.1, , linux, 3.11, [self-hosted Linux], nightly) (push) Has been cancelled
Execution Tests / test (macos-latest) (push) Has been cancelled
Execution Tests / test (ubuntu-latest) (push) Has been cancelled
Execution Tests / test (windows-latest) (push) Has been cancelled
Test server launches without errors / test (push) Has been cancelled
Unit Tests / test (macos-latest) (push) Has been cancelled
Unit Tests / test (ubuntu-latest) (push) Has been cancelled
Unit Tests / test (windows-2022) (push) Has been cancelled
Includes 30 custom nodes committed directly, 7 Civitai-exclusive loras stored via Git LFS, and a setup script that installs all dependencies and downloads HuggingFace-hosted models on vast.ai. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,156 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import cv2
|
||||
import torchvision.transforms as transforms
|
||||
from torch.utils.data import DataLoader
|
||||
from .simple_extractor_dataset import SimpleFolderDataset
|
||||
from .transforms import transform_logits
|
||||
from tqdm import tqdm
|
||||
from PIL import Image
|
||||
|
||||
def get_palette(num_cls):
|
||||
""" Returns the color map for visualizing the segmentation mask.
|
||||
Args:
|
||||
num_cls: Number of classes
|
||||
Returns:
|
||||
The color map
|
||||
"""
|
||||
n = num_cls
|
||||
palette = [0] * (n * 3)
|
||||
for j in range(0, n):
|
||||
lab = j
|
||||
palette[j * 3 + 0] = 0
|
||||
palette[j * 3 + 1] = 0
|
||||
palette[j * 3 + 2] = 0
|
||||
i = 0
|
||||
while lab:
|
||||
palette[j * 3 + 0] |= (((lab >> 0) & 1) << (7 - i))
|
||||
palette[j * 3 + 1] |= (((lab >> 1) & 1) << (7 - i))
|
||||
palette[j * 3 + 2] |= (((lab >> 2) & 1) << (7 - i))
|
||||
i += 1
|
||||
lab >>= 3
|
||||
return palette
|
||||
|
||||
|
||||
def delete_irregular(logits_result):
|
||||
parsing_result = np.argmax(logits_result, axis=2)
|
||||
upper_cloth = np.where(parsing_result == 4, 255, 0)
|
||||
contours, hierarchy = cv2.findContours(upper_cloth.astype(np.uint8),
|
||||
cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_L1)
|
||||
area = []
|
||||
for i in range(len(contours)):
|
||||
a = cv2.contourArea(contours[i], True)
|
||||
area.append(abs(a))
|
||||
if len(area) != 0:
|
||||
top = area.index(max(area))
|
||||
M = cv2.moments(contours[top])
|
||||
cY = int(M["m01"] / M["m00"])
|
||||
|
||||
dresses = np.where(parsing_result == 7, 255, 0)
|
||||
contours_dress, hierarchy_dress = cv2.findContours(dresses.astype(np.uint8),
|
||||
cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_L1)
|
||||
area_dress = []
|
||||
for j in range(len(contours_dress)):
|
||||
a_d = cv2.contourArea(contours_dress[j], True)
|
||||
area_dress.append(abs(a_d))
|
||||
if len(area_dress) != 0:
|
||||
top_dress = area_dress.index(max(area_dress))
|
||||
M_dress = cv2.moments(contours_dress[top_dress])
|
||||
cY_dress = int(M_dress["m01"] / M_dress["m00"])
|
||||
wear_type = "dresses"
|
||||
if len(area) != 0:
|
||||
if len(area_dress) != 0 and cY_dress > cY:
|
||||
irregular_list = np.array([4, 5, 6])
|
||||
logits_result[:, :, irregular_list] = -1
|
||||
else:
|
||||
irregular_list = np.array([5, 6, 7, 8, 9, 10, 12, 13])
|
||||
logits_result[:cY, :, irregular_list] = -1
|
||||
wear_type = "cloth_pant"
|
||||
parsing_result = np.argmax(logits_result, axis=2)
|
||||
# pad border
|
||||
parsing_result = np.pad(parsing_result, pad_width=1, mode='constant', constant_values=0)
|
||||
return parsing_result, wear_type
|
||||
|
||||
|
||||
|
||||
def hole_fill(img):
|
||||
img_copy = img.copy()
|
||||
mask = np.zeros((img.shape[0] + 2, img.shape[1] + 2), dtype=np.uint8)
|
||||
cv2.floodFill(img, mask, (0, 0), 255)
|
||||
img_inverse = cv2.bitwise_not(img)
|
||||
dst = cv2.bitwise_or(img_copy, img_inverse)
|
||||
return dst
|
||||
|
||||
def refine_mask(mask):
|
||||
contours, hierarchy = cv2.findContours(mask.astype(np.uint8),
|
||||
cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_L1)
|
||||
area = []
|
||||
for j in range(len(contours)):
|
||||
a_d = cv2.contourArea(contours[j], True)
|
||||
area.append(abs(a_d))
|
||||
refine_mask = np.zeros_like(mask).astype(np.uint8)
|
||||
if len(area) != 0:
|
||||
i = area.index(max(area))
|
||||
cv2.drawContours(refine_mask, contours, i, color=255, thickness=-1)
|
||||
# keep large area in skin case
|
||||
for j in range(len(area)):
|
||||
if j != i and area[i] > 2000:
|
||||
cv2.drawContours(refine_mask, contours, j, color=255, thickness=-1)
|
||||
return refine_mask
|
||||
|
||||
def refine_hole(parsing_result_filled, parsing_result, arm_mask):
|
||||
filled_hole = cv2.bitwise_and(np.where(parsing_result_filled == 4, 255, 0),
|
||||
np.where(parsing_result != 4, 255, 0)) - arm_mask * 255
|
||||
contours, hierarchy = cv2.findContours(filled_hole, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_L1)
|
||||
refine_hole_mask = np.zeros_like(parsing_result).astype(np.uint8)
|
||||
for i in range(len(contours)):
|
||||
a = cv2.contourArea(contours[i], True)
|
||||
# keep hole > 2000 pixels
|
||||
if abs(a) > 2000:
|
||||
cv2.drawContours(refine_hole_mask, contours, i, color=255, thickness=-1)
|
||||
return refine_hole_mask + arm_mask
|
||||
|
||||
def onnx_inference(lip_session, input_dir, mask_components=[0]):
|
||||
|
||||
transform = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.406, 0.456, 0.485], std=[0.225, 0.224, 0.229])
|
||||
])
|
||||
input_size = [473, 473]
|
||||
|
||||
dataset_lip = SimpleFolderDataset(root=input_dir, input_size=input_size, transform=transform)
|
||||
dataloader_lip = DataLoader(dataset_lip)
|
||||
palette = get_palette(20)
|
||||
with torch.no_grad():
|
||||
for _, batch in enumerate(tqdm(dataloader_lip)):
|
||||
image, meta = batch
|
||||
c = meta['center'].numpy()[0]
|
||||
s = meta['scale'].numpy()[0]
|
||||
w = meta['width'].numpy()[0]
|
||||
h = meta['height'].numpy()[0]
|
||||
|
||||
output = lip_session.run(None, {"input.1": image.numpy().astype(np.float32)})
|
||||
upsample = torch.nn.Upsample(size=input_size, mode='bilinear', align_corners=True)
|
||||
upsample_output = upsample(torch.from_numpy(output[1][0]).unsqueeze(0))
|
||||
upsample_output = upsample_output.squeeze()
|
||||
upsample_output = upsample_output.permute(1, 2, 0) # CHW -> HWC
|
||||
logits_result_lip = transform_logits(upsample_output.data.cpu().numpy(), c, s, w, h,
|
||||
input_size=input_size)
|
||||
parsing_result = np.argmax(logits_result_lip, axis=2)
|
||||
|
||||
output_img = Image.fromarray(np.asarray(parsing_result, dtype=np.uint8))
|
||||
output_img.putpalette(palette)
|
||||
|
||||
mask = np.isin(output_img, mask_components).astype(np.uint8)
|
||||
mask_image = Image.fromarray(mask * 255)
|
||||
mask_image = mask_image.convert("RGB")
|
||||
mask_image = torch.from_numpy(np.array(mask_image).astype(np.float32) / 255.0).unsqueeze(0)
|
||||
|
||||
output_img = output_img.convert('RGB')
|
||||
output_img = torch.from_numpy(np.array(output_img).astype(np.float32) / 255.0).unsqueeze(0)
|
||||
|
||||
return output_img, mask_image
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from .parsing_api import onnx_inference
|
||||
from ...libs.utils import install_package
|
||||
|
||||
class HumanParsing:
|
||||
def __init__(self, model_path):
|
||||
self.model_path = model_path
|
||||
self.session = None
|
||||
|
||||
def __call__(self, input_image, mask_components):
|
||||
if self.session is None:
|
||||
install_package('onnxruntime')
|
||||
import onnxruntime as ort
|
||||
|
||||
session_options = ort.SessionOptions()
|
||||
session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||||
session_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
|
||||
# session_options.add_session_config_entry('gpu_id', str(gpu_id))
|
||||
self.session = ort.InferenceSession(self.model_path, sess_options=session_options,
|
||||
providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
|
||||
|
||||
parsed_image, mask = onnx_inference(self.session, input_image, mask_components)
|
||||
return parsed_image, mask
|
||||
|
||||
|
||||
class HumanParts:
|
||||
|
||||
def __init__(self, model_path):
|
||||
self.model_path = model_path
|
||||
self.session = None
|
||||
# self.classes_dict = {
|
||||
# "background": 0,
|
||||
# "hair": 2,
|
||||
# "glasses": 4,
|
||||
# "top-clothes": 5,
|
||||
# "bottom-clothes": 9,
|
||||
# "torso-skin": 10,
|
||||
# "face": 13,
|
||||
# "left-arm": 14,
|
||||
# "right-arm": 15,
|
||||
# "left-leg": 16,
|
||||
# "right-leg": 17,
|
||||
# "left-foot": 18,
|
||||
# "right-foot": 19,
|
||||
# },
|
||||
self.classes = [0, 13, 2, 4, 5, 9, 10, 14, 15, 16, 17, 18, 19]
|
||||
|
||||
|
||||
def __call__(self, input_image, mask_components):
|
||||
if self.session is None:
|
||||
install_package('onnxruntime')
|
||||
import onnxruntime as ort
|
||||
|
||||
self.session = ort.InferenceSession(self.model_path, providers=['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider'])
|
||||
|
||||
mask, = self.get_mask(self.session, input_image, 0, mask_components)
|
||||
return mask
|
||||
|
||||
def get_mask(self, model, image, rotation, mask_components):
|
||||
image = image.squeeze(0)
|
||||
image_np = image.numpy() * 255
|
||||
|
||||
pil_image = Image.fromarray(image_np.astype(np.uint8))
|
||||
original_size = pil_image.size # to resize the mask later
|
||||
# resize to 512x512 as the model expects
|
||||
pil_image = pil_image.resize((512, 512))
|
||||
center = (256, 256)
|
||||
|
||||
if rotation != 0:
|
||||
pil_image = pil_image.rotate(rotation, center=center)
|
||||
|
||||
# normalize the image
|
||||
image_np = np.array(pil_image).astype(np.float32) / 127.5 - 1
|
||||
image_np = np.expand_dims(image_np, axis=0)
|
||||
|
||||
# use the onnx model to get the mask
|
||||
input_name = model.get_inputs()[0].name
|
||||
output_name = model.get_outputs()[0].name
|
||||
result = model.run([output_name], {input_name: image_np})
|
||||
result = np.array(result[0]).argmax(axis=3).squeeze(0)
|
||||
|
||||
score: int = 0
|
||||
|
||||
mask = np.zeros_like(result)
|
||||
for class_index in mask_components:
|
||||
detected = result == self.classes[class_index]
|
||||
mask[detected] = 255
|
||||
score += mask.sum()
|
||||
|
||||
# back to the original size
|
||||
mask_image = Image.fromarray(mask.astype(np.uint8), mode="L")
|
||||
if rotation != 0:
|
||||
mask_image = mask_image.rotate(-rotation, center=center)
|
||||
|
||||
mask_image = mask_image.resize(original_size)
|
||||
|
||||
# and back to numpy...
|
||||
mask = np.array(mask_image).astype(np.float32) / 255
|
||||
|
||||
# add 2 dimensions to match the expected output
|
||||
mask = np.expand_dims(mask, axis=0)
|
||||
mask = np.expand_dims(mask, axis=0)
|
||||
# ensure to return a "binary mask_image"
|
||||
|
||||
del image_np, result # free up memory, maybe not necessary
|
||||
|
||||
return (torch.from_numpy(mask.astype(np.uint8)),)
|
||||
@@ -0,0 +1,88 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- encoding: utf-8 -*-
|
||||
|
||||
"""
|
||||
@Author : Peike Li
|
||||
@Contact : peike.li@yahoo.com
|
||||
@File : dataset.py
|
||||
@Time : 8/30/19 9:12 PM
|
||||
@Desc : Dataset Definition
|
||||
@License : This source code is licensed under the license found in the
|
||||
LICENSE file in the root directory of this source tree.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from torch.utils import data
|
||||
from .transforms import get_affine_transform
|
||||
|
||||
|
||||
class SimpleFolderDataset(data.Dataset):
|
||||
def __init__(self, root, input_size=[512, 512], transform=None):
|
||||
self.root = root
|
||||
self.input_size = input_size
|
||||
self.transform = transform
|
||||
self.aspect_ratio = input_size[1] * 1.0 / input_size[0]
|
||||
self.input_size = np.asarray(input_size)
|
||||
self.is_pil_image = False
|
||||
if isinstance(root, Image.Image):
|
||||
self.file_list = [root]
|
||||
self.is_pil_image = True
|
||||
elif os.path.isfile(root):
|
||||
self.file_list = [os.path.basename(root)]
|
||||
self.root = os.path.dirname(root)
|
||||
else:
|
||||
self.file_list = os.listdir(self.root)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.file_list)
|
||||
|
||||
def _box2cs(self, box):
|
||||
x, y, w, h = box[:4]
|
||||
return self._xywh2cs(x, y, w, h)
|
||||
|
||||
def _xywh2cs(self, x, y, w, h):
|
||||
center = np.zeros((2), dtype=np.float32)
|
||||
center[0] = x + w * 0.5
|
||||
center[1] = y + h * 0.5
|
||||
if w > self.aspect_ratio * h:
|
||||
h = w * 1.0 / self.aspect_ratio
|
||||
elif w < self.aspect_ratio * h:
|
||||
w = h * self.aspect_ratio
|
||||
scale = np.array([w, h], dtype=np.float32)
|
||||
return center, scale
|
||||
|
||||
def __getitem__(self, index):
|
||||
if self.is_pil_image:
|
||||
img = np.asarray(self.file_list[index])[:, :, [2, 1, 0]]
|
||||
else:
|
||||
img_name = self.file_list[index]
|
||||
img_path = os.path.join(self.root, img_name)
|
||||
img = cv2.imread(img_path, cv2.IMREAD_COLOR)
|
||||
h, w, _ = img.shape
|
||||
|
||||
# Get person center and scale
|
||||
person_center, s = self._box2cs([0, 0, w - 1, h - 1])
|
||||
r = 0
|
||||
trans = get_affine_transform(person_center, s, r, self.input_size)
|
||||
input = cv2.warpAffine(
|
||||
img,
|
||||
trans,
|
||||
(int(self.input_size[1]), int(self.input_size[0])),
|
||||
flags=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=(0, 0, 0))
|
||||
|
||||
input = self.transform(input)
|
||||
meta = {
|
||||
'center': person_center,
|
||||
'height': h,
|
||||
'width': w,
|
||||
'scale': s,
|
||||
'rotation': r
|
||||
}
|
||||
|
||||
return input, meta
|
||||
@@ -0,0 +1,167 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
import cv2
|
||||
import torch
|
||||
|
||||
class BRG2Tensor_transform(object):
|
||||
def __call__(self, pic):
|
||||
img = torch.from_numpy(pic.transpose((2, 0, 1)))
|
||||
if isinstance(img, torch.ByteTensor):
|
||||
return img.float()
|
||||
else:
|
||||
return img
|
||||
|
||||
class BGR2RGB_transform(object):
|
||||
def __call__(self, tensor):
|
||||
return tensor[[2,1,0],:,:]
|
||||
|
||||
def flip_back(output_flipped, matched_parts):
|
||||
'''
|
||||
ouput_flipped: numpy.ndarray(batch_size, num_joints, height, width)
|
||||
'''
|
||||
assert output_flipped.ndim == 4,\
|
||||
'output_flipped should be [batch_size, num_joints, height, width]'
|
||||
|
||||
output_flipped = output_flipped[:, :, :, ::-1]
|
||||
|
||||
for pair in matched_parts:
|
||||
tmp = output_flipped[:, pair[0], :, :].copy()
|
||||
output_flipped[:, pair[0], :, :] = output_flipped[:, pair[1], :, :]
|
||||
output_flipped[:, pair[1], :, :] = tmp
|
||||
|
||||
return output_flipped
|
||||
|
||||
|
||||
def fliplr_joints(joints, joints_vis, width, matched_parts):
|
||||
"""
|
||||
flip coords
|
||||
"""
|
||||
# Flip horizontal
|
||||
joints[:, 0] = width - joints[:, 0] - 1
|
||||
|
||||
# Change left-right parts
|
||||
for pair in matched_parts:
|
||||
joints[pair[0], :], joints[pair[1], :] = \
|
||||
joints[pair[1], :], joints[pair[0], :].copy()
|
||||
joints_vis[pair[0], :], joints_vis[pair[1], :] = \
|
||||
joints_vis[pair[1], :], joints_vis[pair[0], :].copy()
|
||||
|
||||
return joints*joints_vis, joints_vis
|
||||
|
||||
|
||||
def transform_preds(coords, center, scale, input_size):
|
||||
target_coords = np.zeros(coords.shape)
|
||||
trans = get_affine_transform(center, scale, 0, input_size, inv=1)
|
||||
for p in range(coords.shape[0]):
|
||||
target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)
|
||||
return target_coords
|
||||
|
||||
def transform_parsing(pred, center, scale, width, height, input_size):
|
||||
|
||||
trans = get_affine_transform(center, scale, 0, input_size, inv=1)
|
||||
target_pred = cv2.warpAffine(
|
||||
pred,
|
||||
trans,
|
||||
(int(width), int(height)), #(int(width), int(height)),
|
||||
flags=cv2.INTER_NEAREST,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=(0))
|
||||
|
||||
return target_pred
|
||||
|
||||
def transform_logits(logits, center, scale, width, height, input_size):
|
||||
|
||||
trans = get_affine_transform(center, scale, 0, input_size, inv=1)
|
||||
channel = logits.shape[2]
|
||||
target_logits = []
|
||||
for i in range(channel):
|
||||
target_logit = cv2.warpAffine(
|
||||
logits[:,:,i],
|
||||
trans,
|
||||
(int(width), int(height)), #(int(width), int(height)),
|
||||
flags=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=(0))
|
||||
target_logits.append(target_logit)
|
||||
target_logits = np.stack(target_logits,axis=2)
|
||||
|
||||
return target_logits
|
||||
|
||||
|
||||
def get_affine_transform(center,
|
||||
scale,
|
||||
rot,
|
||||
output_size,
|
||||
shift=np.array([0, 0], dtype=np.float32),
|
||||
inv=0):
|
||||
if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
|
||||
print(scale)
|
||||
scale = np.array([scale, scale])
|
||||
|
||||
scale_tmp = scale
|
||||
|
||||
src_w = scale_tmp[0]
|
||||
dst_w = output_size[1]
|
||||
dst_h = output_size[0]
|
||||
|
||||
rot_rad = np.pi * rot / 180
|
||||
src_dir = get_dir([0, src_w * -0.5], rot_rad)
|
||||
dst_dir = np.array([0, (dst_w-1) * -0.5], np.float32)
|
||||
|
||||
src = np.zeros((3, 2), dtype=np.float32)
|
||||
dst = np.zeros((3, 2), dtype=np.float32)
|
||||
src[0, :] = center + scale_tmp * shift
|
||||
src[1, :] = center + src_dir + scale_tmp * shift
|
||||
dst[0, :] = [(dst_w-1) * 0.5, (dst_h-1) * 0.5]
|
||||
dst[1, :] = np.array([(dst_w-1) * 0.5, (dst_h-1) * 0.5]) + dst_dir
|
||||
|
||||
src[2:, :] = get_3rd_point(src[0, :], src[1, :])
|
||||
dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
|
||||
|
||||
if inv:
|
||||
trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
|
||||
else:
|
||||
trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
|
||||
|
||||
return trans
|
||||
|
||||
|
||||
def affine_transform(pt, t):
|
||||
new_pt = np.array([pt[0], pt[1], 1.]).T
|
||||
new_pt = np.dot(t, new_pt)
|
||||
return new_pt[:2]
|
||||
|
||||
|
||||
def get_3rd_point(a, b):
|
||||
direct = a - b
|
||||
return b + np.array([-direct[1], direct[0]], dtype=np.float32)
|
||||
|
||||
|
||||
def get_dir(src_point, rot_rad):
|
||||
sn, cs = np.sin(rot_rad), np.cos(rot_rad)
|
||||
|
||||
src_result = [0, 0]
|
||||
src_result[0] = src_point[0] * cs - src_point[1] * sn
|
||||
src_result[1] = src_point[0] * sn + src_point[1] * cs
|
||||
|
||||
return src_result
|
||||
|
||||
|
||||
def crop(img, center, scale, output_size, rot=0):
|
||||
trans = get_affine_transform(center, scale, rot, output_size)
|
||||
|
||||
dst_img = cv2.warpAffine(img,
|
||||
trans,
|
||||
(int(output_size[1]), int(output_size[0])),
|
||||
flags=cv2.INTER_LINEAR)
|
||||
|
||||
return dst_img
|
||||
Reference in New Issue
Block a user