mirror of
https://github.com/THU-MIG/yolov10.git
synced 2025-05-23 21:44:22 +08:00
aca8eb1fd4
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: Laughing-q <1182102784@qq.com>
104 lines
3.9 KiB
Python
104 lines
3.9 KiB
Python
from pathlib import Path
|
|
from typing import List
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from ultralytics.data.augment import LetterBox
|
|
from ultralytics.utils.ops import xyxy2xywh
|
|
from ultralytics.utils.plotting import plot_images
|
|
|
|
|
|
def get_table_schema(vector_size):
|
|
from lancedb.pydantic import LanceModel, Vector
|
|
|
|
class Schema(LanceModel):
|
|
im_file: str
|
|
labels: List[str]
|
|
cls: List[int]
|
|
bboxes: List[List[float]]
|
|
masks: List[List[List[int]]]
|
|
keypoints: List[List[List[float]]]
|
|
vector: Vector(vector_size)
|
|
|
|
return Schema
|
|
|
|
|
|
def get_sim_index_schema():
|
|
from lancedb.pydantic import LanceModel
|
|
|
|
class Schema(LanceModel):
|
|
idx: int
|
|
im_file: str
|
|
count: int
|
|
sim_im_files: List[str]
|
|
|
|
return Schema
|
|
|
|
|
|
def sanitize_batch(batch, dataset_info):
|
|
batch['cls'] = batch['cls'].flatten().int().tolist()
|
|
box_cls_pair = sorted(zip(batch['bboxes'].tolist(), batch['cls']), key=lambda x: x[1])
|
|
batch['bboxes'] = [box for box, _ in box_cls_pair]
|
|
batch['cls'] = [cls for _, cls in box_cls_pair]
|
|
batch['labels'] = [dataset_info['names'][i] for i in batch['cls']]
|
|
batch['masks'] = batch['masks'].tolist() if 'masks' in batch else [[[]]]
|
|
batch['keypoints'] = batch['keypoints'].tolist() if 'keypoints' in batch else [[[]]]
|
|
|
|
return batch
|
|
|
|
|
|
def plot_similar_images(similar_set, plot_labels=True):
|
|
"""
|
|
Plot images from the similar set.
|
|
|
|
Args:
|
|
similar_set (list): Pyarrow table containing the similar data points
|
|
plot_labels (bool): Whether to plot labels or not
|
|
"""
|
|
similar_set = similar_set.to_pydict()
|
|
empty_masks = [[[]]]
|
|
empty_boxes = [[]]
|
|
images = similar_set.get('im_file', [])
|
|
bboxes = similar_set.get('bboxes', []) if similar_set.get('bboxes') is not empty_boxes else []
|
|
masks = similar_set.get('masks') if similar_set.get('masks')[0] != empty_masks else []
|
|
kpts = similar_set.get('keypoints') if similar_set.get('keypoints')[0] != empty_masks else []
|
|
cls = similar_set.get('cls', [])
|
|
|
|
plot_size = 640
|
|
imgs, batch_idx, plot_boxes, plot_masks, plot_kpts = [], [], [], [], []
|
|
for i, imf in enumerate(images):
|
|
im = cv2.imread(imf)
|
|
im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
|
|
h, w = im.shape[:2]
|
|
r = min(plot_size / h, plot_size / w)
|
|
imgs.append(LetterBox(plot_size, center=False)(image=im).transpose(2, 0, 1))
|
|
if plot_labels:
|
|
if len(bboxes) > i and len(bboxes[i]) > 0:
|
|
box = np.array(bboxes[i], dtype=np.float32)
|
|
box[:, [0, 2]] *= r
|
|
box[:, [1, 3]] *= r
|
|
plot_boxes.append(box)
|
|
if len(masks) > i and len(masks[i]) > 0:
|
|
mask = np.array(masks[i], dtype=np.uint8)[0]
|
|
plot_masks.append(LetterBox(plot_size, center=False)(image=mask))
|
|
if len(kpts) > i and kpts[i] is not None:
|
|
kpt = np.array(kpts[i], dtype=np.float32)
|
|
kpt[:, :, :2] *= r
|
|
plot_kpts.append(kpt)
|
|
batch_idx.append(np.ones(len(np.array(bboxes[i], dtype=np.float32))) * i)
|
|
imgs = np.stack(imgs, axis=0)
|
|
masks = np.stack(plot_masks, axis=0) if len(plot_masks) > 0 else np.zeros(0, dtype=np.uint8)
|
|
kpts = np.concatenate(plot_kpts, axis=0) if len(plot_kpts) > 0 else np.zeros((0, 51), dtype=np.float32)
|
|
boxes = xyxy2xywh(np.concatenate(plot_boxes, axis=0)) if len(plot_boxes) > 0 else np.zeros(0, dtype=np.float32)
|
|
batch_idx = np.concatenate(batch_idx, axis=0)
|
|
cls = np.concatenate([np.array(c, dtype=np.int32) for c in cls], axis=0)
|
|
|
|
fname = 'temp_exp_grid.jpg'
|
|
plot_images(imgs, batch_idx, cls, bboxes=boxes, masks=masks, kpts=kpts, fname=fname,
|
|
max_subplots=len(images)).join()
|
|
img = cv2.imread(fname, cv2.IMREAD_COLOR)
|
|
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
Path(fname).unlink()
|
|
return img_rgb
|