Spaces:
Running
Running
File size: 2,099 Bytes
95eae85 |
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 |
import os
import random
from typing import List, Union
import numpy as np
from PIL import Image, ImageOps
from ultralytics import YOLO
from hp.visualizer import visualizer
def resize_image_pil(image_pil, max_size=1024):
# Ensure image is in RGB
if image_pil.mode != "RGB":
image_pil = image_pil.convert("RGB")
# Calculate new dimensions preserving aspect ratio
width, height = image_pil.size
scale = min(max_size / width, max_size / height)
new_width = int(width * scale)
new_height = int(height * scale)
image_pil = image_pil.resize((new_width, new_height), Image.LANCZOS)
# Calculate padding needed to reach 1024x1024
pad_width = (max_size - new_width) // 2
pad_height = (max_size - new_height) // 2
# Apply padding symmetrically
image_pil = ImageOps.expand(
image_pil,
border=(
pad_width,
pad_height,
max_size - new_width - pad_width,
max_size - new_height - pad_height,
),
fill=(0, 0, 0),
)
return image_pil
def load_resize_image(image_path: str | Image.Image, size: int) -> Image.Image:
if isinstance(image_path, str):
image_pil = Image.open(image_path).convert("RGB")
else:
image_pil = image_path.convert("RGB")
image_pil = resize_image_pil(image_pil, size)
return image_pil
def unload_mask(mask):
mask = mask.cpu().numpy().squeeze()
mask = mask.astype(np.uint8) * 255
return Image.fromarray(mask)
def unload_masks(masks):
return [unload_mask(mask) for mask in masks]
def unload_box(box):
return box.cpu().numpy().tolist()
def unload_boxes(boxes):
return [unload_box(box) for box in boxes]
def format_scores(scores):
return scores.squeeze().cpu().numpy().tolist()
def format_results(labels, scores, boxes, masks):
results_dict = []
for row in zip(labels, scores, boxes, masks):
label, score, box, mask = row
results_row = dict(label=label, score=score, mask=mask, box=box)
results_dict.append(results_row)
return results_dict
|