| import os |
| import json |
| import random |
| from typing import List |
| import csv |
| import glob |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| import torchvision.transforms as transforms |
| from decord import VideoReader |
| from PIL import Image |
| from torch.utils.data import Dataset |
| from transformers import CLIPImageProcessor |
| from tqdm import tqdm |
|
|
|
|
| def process_bbox(bbox, H, W, scale=1.): |
| |
| x_min, y_min, x_max, y_max = bbox |
| width = x_max - x_min |
| height = y_max - y_min |
| |
| side_length = max(width, height) |
| |
| center_x = (x_min + x_max) / 2 |
| center_y = (y_min + y_max) / 2 |
| |
| scaled_side_length = side_length * scale |
| scaled_xmin = center_x - scaled_side_length / 2 |
| scaled_xmax = center_x + scaled_side_length / 2 |
| scaled_ymin = center_y - scaled_side_length / 2 |
| scaled_ymax = center_y + scaled_side_length / 2 |
|
|
| scaled_xmin = int(max(0, scaled_xmin)) |
| scaled_xmax = int(min(W, scaled_xmax)) |
| scaled_ymin = int(max(0, scaled_ymin)) |
| scaled_ymax = int(min(H, scaled_ymax)) |
| |
| return scaled_xmin, scaled_ymin, scaled_xmax, scaled_ymax |
|
|
| def crop_bbox(img, bbox, do_resize=False, size=512): |
| |
| if isinstance(img, (Path, str)): |
| img = Image.open(img) |
| cropped_img = img.crop(bbox) |
| if do_resize: |
| cropped_W, cropped_H = cropped_img.size |
| ratio = size / max(cropped_W, cropped_H) |
| new_W = cropped_W * ratio |
| new_H = cropped_H * ratio |
| cropped_img = cropped_img.resize((new_W, new_H)) |
| |
| return cropped_img |
|
|
| def mask_to_bbox(mask_path): |
| mask = np.array(Image.open(mask_path))[..., 0] |
| rows = np.any(mask, axis=1) |
| cols = np.any(mask, axis=0) |
| |
| ymin, ymax = np.where(rows)[0][[0, -1]] |
| xmin, xmax = np.where(cols)[0][[0, -1]] |
| return xmin, ymin, xmax, ymax |
|
|
| def mask_to_bkgd(img_path, mask_path): |
| img = Image.open(img_path) |
| img_array = np.array(img) |
| |
| mask = Image.open(mask_path).convert("RGB") |
| mask_array = np.array(mask) |
| |
| img_array = np.where(mask_array > 0, img_array, 0) |
| return Image.fromarray(img_array) |
| |
