diff --git a/__init__.py b/__init__.py index 983dca9..b3bbf96 100644 --- a/__init__.py +++ b/__init__.py @@ -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) diff --git a/samples/3.1_resize_extract_parts/atr1.png b/samples/3.1_resize_extract_parts/atr1.png new file mode 100644 index 0000000..18858b4 Binary files /dev/null and b/samples/3.1_resize_extract_parts/atr1.png differ diff --git a/samples/3.1_resize_extract_parts/atr2.png b/samples/3.1_resize_extract_parts/atr2.png new file mode 100644 index 0000000..b7e1c8d Binary files /dev/null and b/samples/3.1_resize_extract_parts/atr2.png differ diff --git a/samples/3.1_resize_extract_parts/atr3.png b/samples/3.1_resize_extract_parts/atr3.png new file mode 100644 index 0000000..eb5b666 Binary files /dev/null and b/samples/3.1_resize_extract_parts/atr3.png differ diff --git a/samples/3.1_resize_extract_parts/image1.png b/samples/3.1_resize_extract_parts/image1.png new file mode 100644 index 0000000..ebbcd25 Binary files /dev/null and b/samples/3.1_resize_extract_parts/image1.png differ diff --git a/samples/3.1_resize_extract_parts/image2.png b/samples/3.1_resize_extract_parts/image2.png new file mode 100644 index 0000000..8a203a7 Binary files /dev/null and b/samples/3.1_resize_extract_parts/image2.png differ diff --git a/samples/3.1_resize_extract_parts/image3.png b/samples/3.1_resize_extract_parts/image3.png new file mode 100644 index 0000000..22a45d9 Binary files /dev/null and b/samples/3.1_resize_extract_parts/image3.png differ diff --git a/samples/3.1_resize_extract_parts/mask.png b/samples/3.1_resize_extract_parts/mask.png new file mode 100644 index 0000000..e91c0c6 Binary files /dev/null and b/samples/3.1_resize_extract_parts/mask.png differ diff --git a/samples/3.1_resize_extract_parts/test_workflow.json b/samples/3.1_resize_extract_parts/test_workflow.json new file mode 100644 index 0000000..edd2888 --- /dev/null +++ b/samples/3.1_resize_extract_parts/test_workflow.json @@ -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 +} \ No newline at end of file