120 lines
4.3 KiB
Python
120 lines
4.3 KiB
Python
import warnings
|
|
import os
|
|
import json
|
|
import numpy as np
|
|
import cv2
|
|
from tqdm import tqdm
|
|
import mediapipe as mp
|
|
from PIL import Image
|
|
from mediapipe.tasks import python
|
|
from mediapipe.tasks.python import vision
|
|
from mediapipe.framework.formats import landmark_pb2
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
import torch
|
|
import requests
|
|
from PIL import Image
|
|
from transformers import BlipProcessor, BlipForConditionalGeneration
|
|
|
|
warnings.filterwarnings("ignore")
|
|
|
|
|
|
class MediapipeEngine:
|
|
def __init__(self, model_asset_path='assets/gesture_recognizer.task'):
|
|
self.mp_hands = mp.solutions.hands
|
|
self.hands = self.mp_hands.Hands(static_image_mode=True, max_num_hands=10, min_detection_confidence=0.5)
|
|
self.mp_drawing = mp.solutions.drawing_utils
|
|
|
|
def __call__(self, image):
|
|
image = np.array(image)
|
|
# image = self.resize_image(image)
|
|
|
|
annotations = self.detect(image.copy())
|
|
mask = self.prepare_mask(image, annotations)
|
|
return Image.fromarray(image), Image.fromarray(mask)
|
|
|
|
def detect(self, image):
|
|
if isinstance(image, Image.Image):
|
|
image = np.array(image)
|
|
|
|
# 进行手部检测
|
|
results = self.hands.process(image)
|
|
annotations = []
|
|
if results.multi_hand_landmarks:
|
|
for hand_landmarks in results.multi_hand_landmarks:
|
|
# self.mp_drawing.draw_landmarks(
|
|
# image, hand_landmarks, self.mp_hands.HAND_CONNECTIONS)
|
|
# 获取所有关键点的坐标
|
|
coords = np.array([(lm.x, lm.y, lm.z) for lm in hand_landmarks.landmark])
|
|
|
|
# 计算边界框
|
|
x_min, y_min, _ = np.min(coords, axis=0)
|
|
x_max, y_max, _ = np.max(coords, axis=0)
|
|
|
|
# 转换为图像坐标
|
|
H, W, _ = image.shape
|
|
x_min, y_min = int(x_min * W), int(y_min * H)
|
|
x_max, y_max = int(x_max * W), int(y_max * H)
|
|
|
|
# loosen bbox
|
|
dynamic_resize = 0.15
|
|
padding = 30
|
|
|
|
bb_xpad = max(int((x_max - x_min + 1) * dynamic_resize), padding)
|
|
bb_ypad = max(int((y_max - y_min + 1) * dynamic_resize), padding)
|
|
bbx_min = max(int(x_min - bb_xpad), 0)
|
|
bbx_max = min(int(x_max + bb_xpad), W-1)
|
|
bby_min = max(int(y_min - bb_ypad), 0)
|
|
bby_max = min(int(y_max + bb_ypad), H-1)
|
|
|
|
annotations.append({
|
|
'bbox': [bbx_min, bby_min, bbx_max, bby_max],
|
|
'landmarks': coords
|
|
})
|
|
return annotations
|
|
|
|
def prepare_mask(self, image, annotations):
|
|
mask = np.zeros(image.shape[:2], dtype=np.uint8)
|
|
for annotation in annotations:
|
|
bbox = annotation['bbox']
|
|
cv2.rectangle(mask, (bbox[0], bbox[1]), (bbox[2], bbox[3]), 255, -1)
|
|
return mask
|
|
|
|
def resize_image(self, image, target_size=1024):
|
|
h, w = image.shape[:2]
|
|
|
|
# 计算长边
|
|
long_side = max(h, w)
|
|
|
|
# 计算缩放比例
|
|
scale = target_size / long_side
|
|
|
|
# 计算新的尺寸
|
|
new_h = int(h * scale)
|
|
new_w = int(w * scale)
|
|
|
|
# 使用cv2.resize进行缩放
|
|
resized_image = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)
|
|
|
|
return resized_image
|
|
|
|
|
|
class ImageCaptioner:
|
|
def __init__(self, model_name="Salesforce/blip-image-captioning-base"):
|
|
self.processor = BlipProcessor.from_pretrained(model_name)
|
|
self.model = BlipForConditionalGeneration.from_pretrained(model_name, torch_dtype=torch.float16).to("cuda")
|
|
|
|
def load_image(self, image_path):
|
|
if image_path.startswith('http'):
|
|
return Image.open(requests.get(image_path, stream=True).raw).convert('RGB')
|
|
else:
|
|
return Image.open(image_path).convert('RGB')
|
|
|
|
def generate_caption(self, image, conditional_text="a photography of"):
|
|
if conditional_text:
|
|
inputs = self.processor(image, conditional_text, return_tensors="pt").to("cuda", torch.float16)
|
|
else:
|
|
inputs = self.processor(image, return_tensors="pt").to("cuda", torch.float16)
|
|
|
|
out = self.model.generate(**inputs)
|
|
return self.processor.decode(out[0], skip_special_tokens=True)
|