diff --git a/__init__.py b/__init__.py index e9e5ef6..e3e0f93 100644 --- a/__init__.py +++ b/__init__.py @@ -18,7 +18,7 @@ from qcloud_cos import CosS3Client from .test_single_image import test_node import ffmpy -video_extensions = ['webm', 'mp4', 'mkv', 'gif', 'mov'] +video_extensions = ["webm", "mp4", "mkv", "gif", "mov"] class FaceDetect: @@ -31,30 +31,67 @@ class FaceDetect: return { "required": { "image": ("IMAGE",), - "main_seed": ("INT:seed", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "main_seed": ( + "INT:seed", + {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}, + ), "model": (["convnext_tiny", "convnext_base"],), "length": ("INT", {"default": 10, "min": 3, "max": 60, "step": 1}), - "threshold": ("FLOAT", {"default": 94, "min": 55, "max": 99, "step": 0.1}) + "threshold": ( + "FLOAT", + {"default": 94, "min": 55, "max": 99, "step": 0.1}, + ), }, } - RETURN_TYPES = ("IMAGE", "IMAGE", "STRING", "STRING", "STRING", "STRING", "STRING", "INT", "INT") - RETURN_NAMES = ("图像", "选中人脸", "分类", "概率", "采用帧序号", "全部帧序列", "剪辑配置", "起始帧序号", "帧数量") + RETURN_TYPES = ( + "IMAGE", + "IMAGE", + "STRING", + "STRING", + "STRING", + "STRING", + "STRING", + "INT", + "INT", + ) + RETURN_NAMES = ( + "图像", + "选中人脸", + "分类", + "概率", + "采用帧序号", + "全部帧序列", + "剪辑配置", + "起始帧序号", + "帧数量", + ) FUNCTION = "predict" CATEGORY = "不忘科技-自定义节点🚩" def predict(self, image, main_seed, model, length, threshold): - image, image_selected, cls, prob, nums, period = test_node(image, length=length, thres=threshold, - model_name=model) + image, image_selected, cls, prob, nums, period = test_node( + image, length=length, thres=threshold, model_name=model + ) print("全部帧序列", period) if len(period) > 0: start, end = period[main_seed % len(period)] config = {"start": start, "end": end} else: raise RuntimeError("未找到符合要求的视频片段") - return (image, image_selected, cls, prob, nums, str(period), json.dumps(config), start, end - start+1) + return ( + image, + image_selected, + cls, + prob, + nums, + str(period), + json.dumps(config), + start, + end - start + 1, + ) class FaceExtract: @@ -77,8 +114,14 @@ class FaceExtract: def crop(self, image): device = model_management.get_torch_device() - image_np = 255. * image.cpu().numpy() - model = YOLO(model=os.path.join(os.path.dirname(os.path.abspath(__file__)), "model", "yolov8n-face-lindevs.pt")) + image_np = 255.0 * image.cpu().numpy() + model = YOLO( + model=os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "model", + "yolov8n-face-lindevs.pt", + ) + ) total_images = image_np.shape[0] out_images = np.ndarray(shape=(total_images, 512, 512, 3)) print("shape", image_np.shape) @@ -86,12 +129,7 @@ class FaceExtract: idx = 0 for image_item in image_np: results = model.predict( - image_item, - imgsz=640, - conf=0.75, - iou=0.7, - device=device, - verbose=False + image_item, imgsz=640, conf=0.75, iou=0.7, device=device, verbose=False ) n = 512 r = results[0] @@ -99,13 +137,21 @@ class FaceExtract: y1, x1, y2, x2, p, cls = r.boxes.data.cpu().numpy()[0] face_size = int(max(y2 - y1, x2 - x1)) center = (x1 + x2) // 2, (y1 + y2) // 2 - x1, x2, y1, y2 = center[0] - face_size // 2, center[0] + face_size // 2, center[1] - face_size // 2, \ - center[1] + face_size // 2 + x1, x2, y1, y2 = ( + center[0] - face_size // 2, + center[0] + face_size // 2, + center[1] - face_size // 2, + center[1] + face_size // 2, + ) template = np.ndarray(shape=(face_size, face_size, 3)) template.fill(20) for a, a1 in zip(list(range(int(x1), int(x2))), list(range(face_size))): - for b, b1 in zip(list(range(int(y1), int(y2))), list(range(face_size))): - if (a >= 0 and a <= r.orig_img.shape[1]) and (b >= 0 and b <= r.orig_img.shape[0]): + for b, b1 in zip( + list(range(int(y1), int(y2))), list(range(face_size)) + ): + if (a >= 0 and a <= r.orig_img.shape[1]) and ( + b >= 0 and b <= r.orig_img.shape[0] + ): template[a1][b1] = r.orig_img[a][b] print(int(x1), int(x2), int(y1), int(y2)) img = cv2.resize(template, (n, n)) @@ -137,26 +183,51 @@ class COSDownload: def download(self, cos_key): if os.sep in cos_key or "/" in cos_key or "\\" in cos_key: - os.makedirs(os.path.join(os.path.dirname(os.path.abspath(__file__)), "download", os.path.dirname(cos_key)), - exist_ok=True) + os.makedirs( + os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "download", + os.path.dirname(cos_key), + ), + exist_ok=True, + ) for i in range(0, 10): try: - with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.yaml"), encoding="utf-8", - mode="r+") as f: + with open( + os.path.join( + os.path.dirname(os.path.abspath(__file__)), "config.yaml" + ), + encoding="utf-8", + mode="r+", + ) as f: yaml_config = yaml.load(f, Loader=yaml.FullLoader) - config = CosConfig(Region=yaml_config["region"], SecretId=yaml_config["secret_id"], - SecretKey=yaml_config["secret_key"]) + config = CosConfig( + Region=yaml_config["region"], + SecretId=yaml_config["secret_id"], + SecretKey=yaml_config["secret_key"], + ) client = CosS3Client(config) response = client.download_file( Bucket=yaml_config["bucket"], Key=cos_key, - DestFilePath=os.path.join(os.path.dirname(os.path.abspath(__file__)), "download", - os.path.dirname(cos_key), os.path.basename(cos_key))) + DestFilePath=os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "download", + os.path.dirname(cos_key), + os.path.basename(cos_key), + ), + ) break except CosClientError or CosServiceError as e: print(f"下载失败 {e}") - return (os.path.join(os.path.dirname(os.path.abspath(__file__)), "download", os.path.dirname(cos_key), - os.path.basename(cos_key)),) + return ( + os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "download", + os.path.dirname(cos_key), + os.path.basename(cos_key), + ), + ) class COSUpload: @@ -181,21 +252,45 @@ class COSUpload: def upload(self, path): for i in range(0, 10): try: - with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.yaml"), encoding="utf-8", - mode="r+") as f: + with open( + os.path.join( + os.path.dirname(os.path.abspath(__file__)), "config.yaml" + ), + encoding="utf-8", + mode="r+", + ) as f: yaml_config = yaml.load(f, Loader=yaml.FullLoader) - config = CosConfig(Region=yaml_config["region"], SecretId=yaml_config["secret_id"], - SecretKey=yaml_config["secret_key"]) + config = CosConfig( + Region=yaml_config["region"], + SecretId=yaml_config["secret_id"], + SecretKey=yaml_config["secret_key"], + ) client = CosS3Client(config) response = client.upload_file( Bucket=yaml_config["bucket"], Key="/".join( - [yaml_config["subfolder"], path.split("/")[-1] if "/" in path else path.split("\\")[-1]]), - LocalFilePath=path) + [ + yaml_config["subfolder"], + ( + path.split("/")[-1] + if "/" in path + else path.split("\\")[-1] + ), + ] + ), + LocalFilePath=path, + ) break except CosClientError or CosServiceError as e: print(e) - return ("/".join([yaml_config["subfolder"], path.split("/")[-1] if "/" in path else path.split("\\")[-1]]),) + return ( + "/".join( + [ + yaml_config["subfolder"], + path.split("/")[-1] if "/" in path else path.split("\\")[-1], + ] + ), + ) class VideoCut: @@ -209,7 +304,16 @@ class VideoCut: "video_path": ("STRING",), "mod": ("INT",), "fps": ("FLOAT",), - "period_length": ("INT", {"default": 10, "min": 4, "max": 100, "step": 1, "forceInput": True}) + "period_length": ( + "INT", + { + "default": 10, + "min": 4, + "max": 100, + "step": 1, + "forceInput": True, + }, + ), }, } @@ -234,31 +338,54 @@ class VideoCut: start, end = config["start"], config["end"] # 新文件名 复制改名适配ffmpeg uid = uuid.uuid1() - temp_fname = os.sep.join([*video_path.split(os.sep)[:-1], "%s.%s" % (str(uid), video_path.split(".")[-1])]) + temp_fname = os.sep.join( + [ + *video_path.split(os.sep)[:-1], + "%s.%s" % (str(uid), video_path.split(".")[-1]), + ] + ) try: shutil.copy(video_path, temp_fname) except: return ("请检查输入文件权限",) video_path = temp_fname # 组装输出文件名 - output_name = (".".join([*video_path.split(os.sep)[-1].split(".")[:-2], - video_path.split(os.sep)[-1].split(".")[-2] + - "_output_%%03d_%s" % datetime.now().strftime('%Y%m%d_%H%M%S'), - video_path.split(os.sep)[-1].split(".")[-1]])) - output = (os.sep.join([*video_path.split(os.sep)[:-1], output_name]) - .replace(os.sep.join(["ComfyUI", "input"]), os.sep.join(["ComfyUI", "output"])).replace(" ", "")) - #调用ffmpeg + output_name = ".".join( + [ + *video_path.split(os.sep)[-1].split(".")[:-2], + video_path.split(os.sep)[-1].split(".")[-2] + + "_output_%%03d_%s" % datetime.now().strftime("%Y%m%d_%H%M%S"), + video_path.split(os.sep)[-1].split(".")[-1], + ] + ) + output = ( + os.sep.join([*video_path.split(os.sep)[:-1], output_name]) + .replace( + os.sep.join(["ComfyUI", "input"]), os.sep.join(["ComfyUI", "output"]) + ) + .replace(" ", "") + ) + # 调用ffmpeg ff = ffmpy.FFmpeg( - inputs={video_path: ['-accurate_seek']}, - outputs={output: [ - '-f', 'segment', - '-ss', str(round(start * mul, 3)), - '-to', str(round(end * mul, 3)), - '-segment_times', str(period_length), - '-c', 'copy', - '-map', '0', - '-avoid_negative_ts', '1' - ]} + inputs={video_path: ["-accurate_seek"]}, + outputs={ + output: [ + "-f", + "segment", + "-ss", + str(round(start * mul, 3)), + "-to", + str(round(end * mul, 3)), + "-segment_times", + str(period_length), + "-c", + "copy", + "-map", + "0", + "-avoid_negative_ts", + "1", + ] + }, ) print(ff.cmd) ff.run() @@ -271,7 +398,9 @@ class VideoCut: files = glob.glob(output.replace("%03d", "*")) for file in files: shutil.move(file, file.replace(str(uid), origin_fname)) - files = glob.glob(output.replace(str(uid), origin_fname).replace("%03d", "*")) + files = glob.glob( + output.replace(str(uid), origin_fname).replace("%03d", "*") + ) return (str(files),) except: files = glob.glob(output.replace("%03d", "*")) @@ -289,8 +418,6 @@ async def get_hello(request): return web.json_response("hello") - - # 腾讯云 VOD from tencentcloud.common.profile.client_profile import ClientProfile @@ -301,13 +428,14 @@ import requests from pathlib import Path import tempfile + class VodToLocalNode: - + def __init__(self): self.secret_id = "AKIDsrihIyjZOBsjimt8TsN8yvv1AMh5dB44" self.secret_key = "CPZcxdk6W39Jd4cGY95wvupoyMd0YFqW" self.vod_client = self.init_vod_client() - + def init_vod_client(self): """初始化VOD客户端""" try: @@ -336,22 +464,22 @@ class VodToLocalNode: # 调用下载逻辑 local_path = self.download_vod(file_id, sub_app_id) return (local_path,) - + def _get_download_url(self, file_id, sub_app_id): """获取媒体文件下载地址""" try: req = models.DescribeMediaInfosRequest() req.FileIds = [file_id] req.SubAppId = int(sub_app_id) - + resp = self.vod_client.DescribeMediaInfos(req) if not resp.MediaInfoSet: raise ValueError("File not found") - + media_info = resp.MediaInfoSet[0] if not media_info.BasicInfo.MediaUrl: raise ValueError("No download URL available") - + return media_info.BasicInfo.MediaUrl except Exception as e: raise RuntimeError(f"Tencent API error: {e}") @@ -359,7 +487,9 @@ class VodToLocalNode: def create_directory(self, path): p = Path(path) if not p.exists(): - p.mkdir(parents=True, exist_ok=True) # parents=True会自动创建所有必需的父目录，exist_ok=True表示如果目录已存在则不会引发异常 + p.mkdir( + parents=True, exist_ok=True + ) # parents=True会自动创建所有必需的父目录，exist_ok=True表示如果目录已存在则不会引发异常 print(f"目录已创建: {path}") else: print(f"目录已存在: {path}") @@ -371,14 +501,23 @@ class VodToLocalNode: """ media_url = self._get_download_url(file_id=file_id, sub_app_id=sub_app_id) # 生成一个临时目录路径名并创建该目录 - self.create_directory(os.path.join(os.path.dirname(os.path.abspath(__file__)), "download", f"{sub_app_id}")) - output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "download", f"{sub_app_id}", f"{file_id}.mp4" ) + self.create_directory( + os.path.join( + os.path.dirname(os.path.abspath(__file__)), "download", f"{sub_app_id}" + ) + ) + output_dir = os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "download", + f"{sub_app_id}", + f"{file_id}.mp4", + ) # 判断文件是否存在 if os.path.exists(output_dir): return output_dir return self._download_file(url=media_url, save_path=output_dir, timeout=60 * 10) - - def _download_file(self, url: str, save_path: str, timeout: int=30): + + def _download_file(self, url: str, save_path: str, timeout: int = 30): """下载文件到本地""" try: with requests.get(url, stream=True, timeout=timeout) as response: @@ -391,6 +530,7 @@ class VodToLocalNode: except Exception as e: raise RuntimeError(f"Download error: {e}") + # A dictionary that contains all nodes you want to export with their names # NOTE: names should be globally unique NODE_CLASS_MAPPINGS = { @@ -399,7 +539,7 @@ NODE_CLASS_MAPPINGS = { "COSUpload": COSUpload, "COSDownload": COSDownload, "VideoCutCustom": VideoCut, - "VodToLocal": VodToLocalNode + "VodToLocal": VodToLocalNode, } # A dictionary that contains the friendly/humanly readable titles for the nodes @@ -409,6 +549,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "COSUpload": "COS上传", "COSDownload": "COS下载", "VideoCutCustom": "视频剪裁", - "VodToLocal": "腾讯云VOD下载" + "VodToLocal": "腾讯云VOD下载", } -