Spaces:
Build error
Build error
File size: 2,160 Bytes
a75bdfa | 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 | 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.):
# transform a bbox(xmin, ymin, xmax, ymax) to (H, W) square
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)
|