82 lines
3.2 KiB
Python
82 lines
3.2 KiB
Python
import cv2
|
|
import numpy as np
|
|
from typing import Union, List
|
|
import torch
|
|
from PIL import Image, ImageDraw
|
|
|
|
def tensor2np(tensor: torch.Tensor):
|
|
if len(tensor.shape) == 3: # Single image
|
|
return np.clip(255.0 * tensor.cpu().numpy(), 0, 255).astype(np.uint8)
|
|
else: # Batch of images
|
|
return [np.clip(255.0 * t.cpu().numpy(), 0, 255).astype(np.uint8) for t in tensor]
|
|
|
|
def np2tensor(img_np: Union[np.ndarray, List[np.ndarray]]) -> torch.Tensor:
|
|
if isinstance(img_np, list):
|
|
if len(img_np) == 0:
|
|
return torch.tensor([])
|
|
return torch.cat([np2tensor(img) for img in img_np], dim=0)
|
|
return torch.from_numpy(img_np.astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
def tensor2pil(t_image: torch.Tensor) -> Image:
|
|
return Image.fromarray(np.clip(255.0 * t_image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
|
|
|
|
def pil2tensor(image:Image) -> torch.Tensor:
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
def image2mask(image:Image) -> torch.Tensor:
|
|
if image.mode == 'L':
|
|
return torch.tensor([pil2tensor(image)[0, :, :].tolist()])
|
|
else:
|
|
image = image.convert('RGB').split()[0]
|
|
return torch.tensor([pil2tensor(image)[0, :, :].tolist()])
|
|
|
|
def mask2image(mask:torch.Tensor) -> Image:
|
|
masks = tensor2np(mask)
|
|
for m in masks:
|
|
_mask = Image.fromarray(m).convert("L")
|
|
_image = Image.new("RGBA", _mask.size, color='white')
|
|
_image = Image.composite(
|
|
_image, Image.new("RGBA", _mask.size, color='black'), _mask)
|
|
return _image
|
|
|
|
def pil2cv2(pil_img:Image) -> np.array:
|
|
np_img_array = np.asarray(pil_img)
|
|
return cv2.cvtColor(np_img_array, cv2.COLOR_RGB2BGR)
|
|
|
|
def min_bounding_rect(image:Image) -> tuple:
|
|
cv2_image = pil2cv2(image)
|
|
gray = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2GRAY)
|
|
ret, thresh = cv2.threshold(gray, 127, 255, 0)
|
|
contours, _ = cv2.findContours(thresh, 1, 2)
|
|
x, y, width, height = 0, 0, 0, 0
|
|
area = 0
|
|
for contour in contours:
|
|
_x, _y, _w, _h = cv2.boundingRect(contour)
|
|
_area = _w * _h
|
|
if _area > area:
|
|
area = _area
|
|
x, y, width, height = _x, _y, _w, _h
|
|
return (x, y, width, height)
|
|
|
|
def mask_area(image:Image) -> tuple:
|
|
cv2_image = pil2cv2(image.convert('RGBA'))
|
|
gray = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2GRAY)
|
|
_, thresh = cv2.threshold(gray, 127, 255, 0)
|
|
locs = np.where(thresh == 255)
|
|
x1 = np.min(locs[1]) if len(locs[1]) > 0 else 0
|
|
x2 = np.max(locs[1]) if len(locs[1]) > 0 else image.width
|
|
y1 = np.min(locs[0]) if len(locs[0]) > 0 else 0
|
|
y2 = np.max(locs[0]) if len(locs[0]) > 0 else image.height
|
|
x1, y1, x2, y2 = min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2)
|
|
return (x1, y1, x2 - x1, y2 - y1)
|
|
|
|
def draw_rect(image:Image, x:int, y:int, width:int, height:int, line_color:str, line_width:int,
|
|
box_color:str=None) -> Image:
|
|
draw = ImageDraw.Draw(image)
|
|
draw.rectangle((x, y, x + width, y + height), fill=box_color, outline=line_color, width=line_width, )
|
|
return image
|
|
|
|
# Tensor to cv2
|
|
def tensor2cv(image):
|
|
image_np = np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
|
|
return cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR) |