File size: 4,522 Bytes
45a63fb | 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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | 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)
|