Added a way to crop image in standard size for all images in batch and final output images will be divisible by 8

This commit is contained in:
Ram Deshmukh
2024-04-25 15:11:04 +05:30
parent 7595071a3c
commit 5fd4064b1e
9 changed files with 914 additions and 22 deletions
+42 -22
View File
@@ -400,7 +400,7 @@ class TRI3DExtractPartsBatch:
# If no contours were found, just return None
if not contours:
return None
return None, None
# Combine all contours to find encompassing bounding box
all_points = np.concatenate(contours, axis=0)
@@ -416,9 +416,9 @@ class TRI3DExtractPartsBatch:
print(x, y, w, h, "x,y,w,h")
print(input_img.shape, "input_img.shape")
# Extract the region from the original image that contains both hands
hand_region = input_img[y:y + h, x:x + w]
return hand_region
# hand_region = input_img[y:y + h, x:x + w]
return input_img, [x,y,w,h]
def tensor_to_cv2_img(tensor, remove_alpha=False):
# This will give us (H, W, C)
@@ -436,6 +436,9 @@ class TRI3DExtractPartsBatch:
batch_results = []
images = []
secondaries = []
images_xywh = []
secondaries_xywh = []
# cv2_secondary = tensor_to_cv2_img(batch_secondaries)
for i in range(batch_images.shape[0]):
image = batch_images[i]
@@ -484,18 +487,20 @@ class TRI3DExtractPartsBatch:
if scarf:
color_code_list.append([128, 64, 0])
bimage = bounded_image(cv2_seg, color_code_list, cv2_image)
bsecondary = bounded_image(cv2_seg, color_code_list, cv2_secondary)
bimage, img_xywh = bounded_image(cv2_seg, color_code_list, cv2_image)
bsecondary, sec_xywh = bounded_image(cv2_seg, color_code_list, cv2_secondary)
# Handle case when bimage is None to avoid error during conversion to tensor
if bimage is not None:
images.append(bimage)
images_xywh.append(img_xywh)
else:
num_channels = cv2_image.shape[2] if len(
cv2_image.shape) > 2 else 1
black_img = np.zeros((10, 10, num_channels),
dtype=cv2_image.dtype)
images.append(black_img)
images_xywh.append([0, 0, 10, 10])
if bsecondary is not None:
secondaries.append(bsecondary)
@@ -507,28 +512,43 @@ class TRI3DExtractPartsBatch:
secondaries.append(black_img)
# Get max height and width
max_height = max(img.shape[0] for img in images)
max_width = max(img.shape[1] for img in images)
max_height = max([xywh[-1] for xywh in images_xywh])
max_width = max([xywh[-2] for xywh in images_xywh])
for i,img in enumerate(images): #this takes care of edge case where max crop height/width of batch
#exceeds image size of few images
x,y,_,_ = images_xywh[i]
h,w,_ = img.shape
if x+max_width > w:
max_width = w - x
if y+max_height > h:
max_height = h - y
max_height = max_height - max_height % 8 ##Making it divisible by 8
max_width = max_width - max_width % 8
batch_results = []
batch_secondaries = []
for img in images:
# Resize the image to max height and width
resized_img = cv2.resize(img, (max_width, max_height),
interpolation=cv2.INTER_CUBIC)
tensor_img = cv2_img_to_tensor(resized_img)
batch_results.append(tensor_img.squeeze(0))
for sec in secondaries:
# Resize the image to max height and width
resized_sec = cv2.resize(sec, (max_width, max_height),
interpolation=cv2.INTER_NEAREST)
tensor_sec = cv2_img_to_tensor(resized_sec)
batch_secondaries.append(tensor_sec.squeeze(0))
for i,img in enumerate(images):
x,y,w,h = images_xywh[i]
img = img[y:y+max_height, x:x+max_width]
img = cv2_img_to_tensor(img)
batch_results.append(img.squeeze(0))
for i,sec in enumerate(secondaries):
x,y,w,h = images_xywh[i]
sec = sec[y:y+max_height, x:x+max_width]
sec = cv2_img_to_tensor(sec)
batch_secondaries.append(sec.squeeze(0))
batch_results = torch.stack(batch_results)
batch_secondaries = torch.stack(batch_secondaries)
print(batch_results.shape, "batch_results.shape")
return (batch_results, batch_secondaries)
Binary file not shown.

After

Width:  |  Height:  |  Size: 87 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 86 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 86 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.8 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.8 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.8 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 134 KiB

@@ -0,0 +1,872 @@
{
"last_node_id": 28,
"last_link_id": 35,
"nodes": [
{
"id": 2,
"type": "LoadImage",
"pos": [
1968,
-883
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
1
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"title": "Image1",
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"image1.png",
"image"
]
},
{
"id": 3,
"type": "LoadImage",
"pos": [
1969,
-514
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
2
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"title": "Image2",
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"image2.png",
"image"
]
},
{
"id": 4,
"type": "LoadImage",
"pos": [
1969,
-151
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 2,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
3
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"title": "Image3",
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"image3.png",
"image"
]
},
{
"id": 5,
"type": "LoadImage",
"pos": [
2370,
-600
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 3,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
5
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"title": "ATR1",
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"atr1.png",
"image"
]
},
{
"id": 6,
"type": "LoadImage",
"pos": [
2367,
-236
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 4,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
6
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"title": "ATR2",
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"atr2.png",
"image"
]
},
{
"id": 14,
"type": "ImageBatch",
"pos": [
2745,
-584
],
"size": {
"0": 210,
"1": 46
},
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "image1",
"type": "IMAGE",
"link": 5
},
{
"name": "image2",
"type": "IMAGE",
"link": 6
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
7
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ImageBatch"
}
},
{
"id": 7,
"type": "LoadImage",
"pos": [
2372,
129
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 5,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
8
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"title": "ATR3",
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"atr3.png",
"image"
]
},
{
"id": 11,
"type": "ImageBatch",
"pos": [
2390,
-848
],
"size": {
"0": 210,
"1": 46
},
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "image1",
"type": "IMAGE",
"link": 1
},
{
"name": "image2",
"type": "IMAGE",
"link": 2
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
4
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ImageBatch"
}
},
{
"id": 12,
"type": "ImageBatch",
"pos": [
2981,
-582
],
"size": {
"0": 210,
"1": 46
},
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "image1",
"type": "IMAGE",
"link": 7
},
{
"name": "image2",
"type": "IMAGE",
"link": 8
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
10
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ImageBatch"
}
},
{
"id": 8,
"type": "LoadImage",
"pos": [
2786,
-268
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 6,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
11
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"title": "Mask",
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"mask (30).png",
"image"
]
},
{
"id": 15,
"type": "RepeatImageBatch",
"pos": [
2793,
-367
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 11
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
18
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "RepeatImageBatch"
},
"widgets_values": [
3
]
},
{
"id": 13,
"type": "ImageBatch",
"pos": [
2636,
-847
],
"size": {
"0": 210,
"1": 46
},
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "image1",
"type": "IMAGE",
"link": 4
},
{
"name": "image2",
"type": "IMAGE",
"link": 3
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
16
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ImageBatch"
}
},
{
"id": 26,
"type": "Image To Mask",
"pos": [
3667,
-11
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 15,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 32
}
],
"outputs": [
{
"name": "MASK",
"type": "MASK",
"links": [
33
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "Image To Mask"
},
"widgets_values": [
"intensity"
]
},
{
"id": 27,
"type": "InpaintPreprocessor",
"pos": [
3678,
99
],
"size": {
"0": 210,
"1": 46
},
"flags": {},
"order": 16,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 34
},
{
"name": "mask",
"type": "MASK",
"link": 33
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
35
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "InpaintPreprocessor"
}
},
{
"id": 17,
"type": "PreviewImage",
"pos": [
3679,
-294
],
"size": {
"0": 686.6637573242188,
"1": 246
},
"flags": {},
"order": 14,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 19
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 28,
"type": "PreviewImage",
"pos": [
3675,
196
],
"size": {
"0": 786.8681640625,
"1": 246
},
"flags": {},
"order": 17,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 35
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 16,
"type": "PreviewImage",
"pos": [
3661,
-574
],
"size": {
"0": 714.0955810546875,
"1": 246
},
"flags": {},
"order": 13,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 17
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 1,
"type": "tri3d-extract-parts-batch",
"pos": [
3280,
-482
],
"size": {
"0": 315,
"1": 530
},
"flags": {},
"order": 12,
"mode": 0,
"inputs": [
{
"name": "batch_images",
"type": "IMAGE",
"link": 16
},
{
"name": "batch_segs",
"type": "IMAGE",
"link": 10
},
{
"name": "batch_secondaries",
"type": "IMAGE",
"link": 18
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
17,
34
],
"shape": 3,
"slot_index": 0
},
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
19,
32
],
"shape": 3,
"slot_index": 1
}
],
"properties": {
"Node name for S&R": "tri3d-extract-parts-batch"
},
"widgets_values": [
20,
false,
false,
true,
true,
false,
false,
false,
false,
false,
false,
false,
false,
false,
false,
false,
false,
false,
false
]
}
],
"links": [
[
1,
2,
0,
11,
0,
"IMAGE"
],
[
2,
3,
0,
11,
1,
"IMAGE"
],
[
3,
4,
0,
13,
1,
"IMAGE"
],
[
4,
11,
0,
13,
0,
"IMAGE"
],
[
5,
5,
0,
14,
0,
"IMAGE"
],
[
6,
6,
0,
14,
1,
"IMAGE"
],
[
7,
14,
0,
12,
0,
"IMAGE"
],
[
8,
7,
0,
12,
1,
"IMAGE"
],
[
10,
12,
0,
1,
1,
"IMAGE"
],
[
11,
8,
0,
15,
0,
"IMAGE"
],
[
16,
13,
0,
1,
0,
"IMAGE"
],
[
17,
1,
0,
16,
0,
"IMAGE"
],
[
18,
15,
0,
1,
2,
"IMAGE"
],
[
19,
1,
1,
17,
0,
"IMAGE"
],
[
32,
1,
1,
26,
0,
"IMAGE"
],
[
33,
26,
0,
27,
1,
"MASK"
],
[
34,
1,
0,
27,
0,
"IMAGE"
],
[
35,
27,
0,
28,
0,
"IMAGE"
]
],
"groups": [],
"config": {},
"extra": {},
"version": 0.4
}