Spaces:
Running
Running
File size: 2,185 Bytes
d228e38 |
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 |
from ultralytics import YOLO
from PIL import Image
import numpy as np
import cv2
def load_models():
classification_model = YOLO('models/classification_model.pt')
segmentation_model = YOLO('models/segmentation_model.pt')
return classification_model, segmentation_model
def classify_image(image, classification_model):
results = classification_model(image, verbose=True)
category = results[0].probs.top5
return category
def segment_image(image, segmentation_model):
results = segmentation_model(image, agnostic_nms=True, retina_masks=True, verbose=True)
return results
def tile_pattern(image_size, pattern_image):
pattern_np = np.array(pattern_image)
pattern_height, pattern_width = pattern_np.shape[:2]
tiled_pattern = np.tile(pattern_np, (
image_size[1] // pattern_height + 1,
image_size[0] // pattern_width + 1,
1
))
return Image.fromarray(tiled_pattern[:image_size[1], :image_size[0]])
def apply_mask(image, mask, pattern_image, head_image):
image_np = np.array(image)
mask_np = np.array(mask)
tiled_pattern = tile_pattern(image.size, pattern_image)
tiled_pattern_np = np.array(tiled_pattern)
mask_indices = np.argwhere(mask_np > 0)
if mask_indices.size > 0:
min_y, min_x = mask_indices.min(axis=0)
max_y, max_x = mask_indices.max(axis=0)
dst_points = np.float32([
[min_x, min_y],
[max_x, min_y],
[max_x, max_y],
[min_x, max_y]
])
head_width, head_height = head_image.size
src_points = np.float32([
[0, 0],
[head_width, 0],
[head_width, head_height],
[0, head_height]
])
matrix = cv2.getPerspectiveTransform(src_points, dst_points)
head_np = np.array(head_image)
transformed_head = cv2.warpPerspective(head_np, matrix, (image.size[0], image.size[1]))
overlay = Image.fromarray(transformed_head, 'RGBA')
image_with_head = Image.alpha_composite(Image.fromarray(image_np).convert("RGBA"), overlay)
return image_with_head.convert("RGB")
return Image.fromarray(image_np)
|