Spaces:
Build error
Build error
File size: 3,281 Bytes
11d7b39 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
import cv2
import numpy as np
from pycocotools import mask as coco_mask
from pytoshop import layers
import pytoshop
from pytoshop.enums import BlendMode
from datetime import datetime
def generate_random_color():
return np.random.randint(0, 256), np.random.randint(0, 256), np.random.randint(0, 256)
def create_base_layer(image):
rgba_image = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
return [rgba_image]
def create_mask_layers(image, masks):
layer_list = []
for result in masks:
rle = result['segmentation']
mask = coco_mask.decode(rle).astype(np.uint8)
rgba_image = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
rgba_image[..., 3] = cv2.bitwise_and(rgba_image[..., 3], rgba_image[..., 3], mask=mask)
layer_list.append(rgba_image)
return layer_list
def create_mask_gallery(image, masks):
mask_array_list = []
label_list = []
for index, result in enumerate(masks):
rle = result['segmentation']
mask = coco_mask.decode(rle).astype(np.uint8)
rgba_image = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
rgba_image[..., 3] = cv2.bitwise_and(rgba_image[..., 3], rgba_image[..., 3], mask=mask)
mask_array_list.append(rgba_image)
label_list.append(f'Part {index}')
return [[img, label] for img, label in zip(mask_array_list, label_list)]
def create_mask_combined_images(image, masks):
final_result = np.zeros_like(image)
for result in masks:
rle = result['segmentation']
mask = coco_mask.decode(rle).astype(np.uint8)
color = generate_random_color()
colored_mask = np.zeros_like(image)
colored_mask[mask == 1] = color
final_result = cv2.addWeighted(final_result, 1, colored_mask, 0.5, 0)
combined_image = cv2.addWeighted(image, 1, final_result, 0.5, 0)
return [combined_image, "masked"]
def insert_psd_layer(psd, image_data, layer_name, blending_mode):
channel_data = [layers.ChannelImageData(image=image_data[:, :, i], compression=1) for i in range(4)]
layer_record = layers.LayerRecord(
channels={-1: channel_data[3], 0: channel_data[0], 1: channel_data[1], 2: channel_data[2]},
top=0, bottom=image_data.shape[0], left=0, right=image_data.shape[1],
blend_mode=blending_mode,
name=layer_name,
opacity=255,
)
psd.layer_and_mask_info.layer_info.layer_records.append(layer_record)
return psd
def save_psd(input_image_data, layer_data, layer_names, blending_modes):
psd_file = pytoshop.core.PsdFile(num_channels=3, height=input_image_data.shape[0], width=input_image_data.shape[1])
for index, layer in enumerate(layer_data):
psd_file = insert_psd_layer(psd_file, layer, layer_names[index], blending_modes[index])
timestamp = datetime.now().strftime("%m%d%H%M%S")
with open(f"outputs/psd/result-{timestamp}.psd", 'wb') as output_file:
psd_file.write(output_file)
def save_psd_with_masks(image, masks):
original_layer = create_base_layer(image)
mask_layers = create_mask_layers(image, masks)
names = [f'Part {i}' for i in range(len(mask_layers))]
modes = [BlendMode.normal] * (len(mask_layers)+1)
save_psd(image, original_layer+mask_layers, ['Original_Image']+names, modes)
|