feat: Add output image resizing for DWPose

This commit is contained in:
blessedcoolant 2024-02-10 03:35:10 +05:30 committed by Kent Keirsey
parent f7998b4be0
commit 67cbfeb33d
6 changed files with 67 additions and 19 deletions

View File

@ -292,7 +292,7 @@ class OpenposeImageProcessorInvocation(ImageProcessorInvocation):
image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res) image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res)
def run_processor(self, image): def run_processor(self, image):
openpose_processor = OpenposeDetector.from_pretrained("lllyasviel/Annotators") openpose_processor = OpenposeDetector.from_pretrained(pretrained_model_or_path="lllyasviel/Annotators")
processed_image = openpose_processor( processed_image = openpose_processor(
image, image,
detect_resolution=self.detect_resolution, detect_resolution=self.detect_resolution,
@ -649,8 +649,15 @@ class DWPoseImageProcessorInvocation(ImageProcessorInvocation):
draw_body: bool = InputField(default=True) draw_body: bool = InputField(default=True)
draw_face: bool = InputField(default=False) draw_face: bool = InputField(default=False)
draw_hands: bool = InputField(default=False) draw_hands: bool = InputField(default=False)
image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res)
def run_processor(self, image): def run_processor(self, image):
dwpose = DWPoseDetector() dwpose = DWPoseDetector()
processed_image = dwpose(image, draw_face=self.draw_face, draw_hands=self.draw_hands, draw_body=self.draw_body) processed_image = dwpose(
image,
draw_face=self.draw_face,
draw_hands=self.draw_hands,
draw_body=self.draw_body,
resolution=self.image_resolution,
)
return processed_image return processed_image

View File

@ -1,12 +1,13 @@
import numpy as np import numpy as np
import torch import torch
from controlnet_aux.util import resize_image
from PIL import Image from PIL import Image
from invokeai.backend.image_util.dwpose.utils import draw_bodypose, draw_facepose, draw_handpose from invokeai.backend.image_util.dwpose.utils import draw_bodypose, draw_facepose, draw_handpose
from invokeai.backend.image_util.dwpose.wholebody import Wholebody from invokeai.backend.image_util.dwpose.wholebody import Wholebody
def draw_pose(pose, H, W, draw_face=True, draw_body=True, draw_hands=True): def draw_pose(pose, H, W, draw_face=True, draw_body=True, draw_hands=True, resolution=512):
bodies = pose["bodies"] bodies = pose["bodies"]
faces = pose["faces"] faces = pose["faces"]
hands = pose["hands"] hands = pose["hands"]
@ -23,7 +24,11 @@ def draw_pose(pose, H, W, draw_face=True, draw_body=True, draw_hands=True):
if draw_face: if draw_face:
canvas = draw_facepose(canvas, faces) canvas = draw_facepose(canvas, faces)
dwpose_image = Image.fromarray(canvas) dwpose_image = resize_image(
canvas,
resolution,
)
dwpose_image = Image.fromarray(dwpose_image)
return dwpose_image return dwpose_image
@ -32,7 +37,9 @@ class DWPoseDetector:
def __init__(self) -> None: def __init__(self) -> None:
self.pose_estimation = Wholebody() self.pose_estimation = Wholebody()
def __call__(self, image: Image.Image, draw_face=False, draw_body=True, draw_hands=False) -> Image.Image: def __call__(
self, image: Image.Image, draw_face=False, draw_body=True, draw_hands=False, resolution=512
) -> Image.Image:
np_image = np.array(image) np_image = np.array(image)
H, W, C = np_image.shape H, W, C = np_image.shape
@ -64,4 +71,6 @@ class DWPoseDetector:
bodies = {"candidate": body, "subset": score} bodies = {"candidate": body, "subset": score}
pose = {"bodies": bodies, "hands": hands, "faces": faces} pose = {"bodies": bodies, "hands": hands, "faces": faces}
return draw_pose(pose, H, W, draw_face=draw_face, draw_hands=draw_hands, draw_body=draw_body) return draw_pose(
pose, H, W, draw_face=draw_face, draw_hands=draw_hands, draw_body=draw_body, resolution=resolution
)

View File

@ -1,4 +1,4 @@
import { Flex, FormControl, FormLabel, Switch } from '@invoke-ai/ui-library'; import { CompositeNumberInput, CompositeSlider, Flex, FormControl, FormLabel, Switch } from '@invoke-ai/ui-library';
import { useProcessorNodeChanged } from 'features/controlAdapters/components/hooks/useProcessorNodeChanged'; import { useProcessorNodeChanged } from 'features/controlAdapters/components/hooks/useProcessorNodeChanged';
import { CONTROLNET_PROCESSORS } from 'features/controlAdapters/store/constants'; import { CONTROLNET_PROCESSORS } from 'features/controlAdapters/store/constants';
import type { RequiredDWPoseImageProcessorInvocation } from 'features/controlAdapters/store/types'; import type { RequiredDWPoseImageProcessorInvocation } from 'features/controlAdapters/store/types';
@ -18,7 +18,7 @@ type Props = {
const DWPoseProcessor = (props: Props) => { const DWPoseProcessor = (props: Props) => {
const { controlNetId, processorNode, isEnabled } = props; const { controlNetId, processorNode, isEnabled } = props;
const { draw_body, draw_face, draw_hands } = processorNode; const { image_resolution, draw_body, draw_face, draw_hands } = processorNode;
const processorChanged = useProcessorNodeChanged(); const processorChanged = useProcessorNodeChanged();
const { t } = useTranslation(); const { t } = useTranslation();
@ -43,6 +43,13 @@ const DWPoseProcessor = (props: Props) => {
[controlNetId, processorChanged] [controlNetId, processorChanged]
); );
const handleImageResolutionChanged = useCallback(
(v: number) => {
processorChanged(controlNetId, { image_resolution: v });
},
[controlNetId, processorChanged]
);
return ( return (
<ProcessorWrapper> <ProcessorWrapper>
<Flex sx={{ flexDir: 'row', gap: 6 }}> <Flex sx={{ flexDir: 'row', gap: 6 }}>
@ -59,6 +66,24 @@ const DWPoseProcessor = (props: Props) => {
<Switch defaultChecked={DEFAULTS.draw_hands} isChecked={draw_hands} onChange={handleDrawHandsChanged} /> <Switch defaultChecked={DEFAULTS.draw_hands} isChecked={draw_hands} onChange={handleDrawHandsChanged} />
</FormControl> </FormControl>
</Flex> </Flex>
<FormControl isDisabled={!isEnabled}>
<FormLabel>{t('controlnet.imageResolution')}</FormLabel>
<CompositeSlider
value={image_resolution}
onChange={handleImageResolutionChanged}
defaultValue={DEFAULTS.image_resolution}
min={0}
max={4096}
marks
/>
<CompositeNumberInput
value={image_resolution}
onChange={handleImageResolutionChanged}
defaultValue={DEFAULTS.image_resolution}
min={0}
max={4096}
/>
</FormControl>
</ProcessorWrapper> </ProcessorWrapper>
); );
}; };

View File

@ -232,6 +232,7 @@ export const CONTROLNET_PROCESSORS: ControlNetProcessorsDict = {
default: { default: {
id: 'dwpose_image_processor', id: 'dwpose_image_processor',
type: 'dwpose_image_processor', type: 'dwpose_image_processor',
image_resolution: 512,
draw_body: true, draw_body: true,
draw_face: false, draw_face: false,
draw_hands: false, draw_hands: false,

View File

@ -157,7 +157,7 @@ export type RequiredOpenposeImageProcessorInvocation = O.Required<
*/ */
export type RequiredDWPoseImageProcessorInvocation = O.Required< export type RequiredDWPoseImageProcessorInvocation = O.Required<
DWPoseImageProcessorInvocation, DWPoseImageProcessorInvocation,
'type' | 'draw_body' | 'draw_face' | 'draw_hands' 'type' | 'image_resolution' | 'draw_body' | 'draw_face' | 'draw_hands'
>; >;
/** /**

File diff suppressed because one or more lines are too long