import torch
import [Link] as nn
from torchvision import models, transforms
from [Link] import Dataset, DataLoader
from [Link] import DataParallel
from [Link] import Sampler
from PIL import Image
import numpy as np
from [Link] import gaussian_filter1d
import os
import cv2
import mstcn
from transformer2_3_1 import Transformer2_3_1
import multiprocessing
import argparse
import json
import sys
sequence_length = 1
val_batch_size = 100
workers = min(4, multiprocessing.cpu_count() - 1) # 保留 1 个核心供其他任务使用
num_gpus = [Link].device_count()
gpus = ",".join(map(str, range(num_gpus)))
[Link]["CUDA_VISIBLE_DEVICES"] = gpus
def pil_loader(path):
with open(path, 'rb') as f:
with [Link](f) as img:
return [Link]('RGB')
def change_size(image):
binary_image = [Link](image, cv2.COLOR_BGR2GRAY)
_, binary_image2 = [Link](binary_image, 15, 255, cv2.THRESH_BINARY)
binary_image2 = [Link](binary_image2, 19)
x = binary_image2.shape[0]
y = binary_image2.shape[1]
edges_x = []
edges_y = []
for i in range(x):
for j in range(10, y-10):
if binary_image2.item(i, j) != 0:
edges_x.append(i)
edges_y.append(j)
if not edges_x:
return image
left = min(edges_x)
right = max(edges_x)
width = right - left
bottom = min(edges_y)
top = max(edges_y)
height = top - bottom
pre1_picture = image[left:left + width, bottom:bottom + height]
return pre1_picture
def get_data_from_video(video_path):
cap = [Link](video_path)
fps = [Link](cv2.CAP_PROP_FPS)
frame_count = int([Link](cv2.CAP_PROP_FRAME_COUNT))
duration = frame_count / fps
frames = []
for sec in range(int(duration)):
[Link](cv2.CAP_PROP_POS_MSEC, sec * 1000)
success, frame = [Link]()
if success:
dim = (int([Link][1] / [Link][0] * 300), 300)
frame = [Link](frame, dim)
frame = change_size(frame)
frame = [Link](frame, (250, 250))
frame = [Link](frame, cv2.COLOR_BGR2RGB)
frame = [Link](frame)
[Link](frame)
else:
break
[Link]()
return frames
def pil_loader_from_video(video_path):
frames = get_data_from_video(video_path)
return frames
class VideoDataset(Dataset):
def __init__(self, video_path, transform=None):
[Link] = pil_loader_from_video(video_path)
[Link] = transform
def __getitem__(self, index):
img = [Link][index]
if [Link] is not None:
img = [Link](img)
return img, index
def __len__(self):
return len([Link])
class resnet_lstm([Link]):
def __init__(self):
super(resnet_lstm, self).__init__()
resnet = models.resnet50(pretrained=True)
[Link] = [Link](
resnet.conv1,
resnet.bn1,
[Link],
[Link],
resnet.layer1,
resnet.layer2,
resnet.layer3,
resnet.layer4,
[Link]
)
[Link] = [Link]([Link](2048, 512),
[Link](),
[Link](512, 7))
def forward(self, x):
x = [Link](-1, 3, 224, 224)
x = [Link](x)
x = [Link](-1, 2048)
return x
def get_useful_start_idx(sequence_length, list_each_length):
count = 0
idx = []
for i in range(len(list_each_length)):
for j in range(count, count + (list_each_length[i] + 1 - sequence_length)):
[Link](j)
count += list_each_length[i]
return idx
def get_useful_start_idx_LFB(sequence_length, list_each_length):
count = 0
idx = []
for i in range(len(list_each_length)):
for j in range(count, count + (list_each_length[i] + 1 - sequence_length)):
[Link](j)
count += list_each_length[i]
return idx
def get_data(video_path):
test_transforms = None
test_transforms = [Link]([
[Link]((250, 250)),
[Link](224),
[Link](),
[Link]([0.41757566, 0.26098573, 0.25888634], [0.21938758,
0.1983, 0.19342837])
])
video_dataset = VideoDataset(video_path, test_transforms)
return video_dataset
class SeqSampler(Sampler):
def __init__(self, data_source, idx):
super().__init__(data_source)
self.data_source = data_source
[Link] = idx
def __iter__(self):
return iter([Link])
def __len__(self):
return len([Link])
sig_f = [Link]()
g_LFB_test = [Link](shape=(0, 2048))
def first_stage_inference(test_dataset):
test_num_each = [len(test_dataset)]
test_useful_start_idx = get_useful_start_idx(sequence_length, test_num_each)
test_useful_start_idx_LFB = get_useful_start_idx_LFB(sequence_length,
test_num_each)
num_test_we_use = len(test_useful_start_idx)
num_test_we_use_LFB = len(test_useful_start_idx_LFB)
test_we_use_start_idx = test_useful_start_idx
test_we_use_start_idx_LFB = test_useful_start_idx_LFB
test_idx = []
for i in range(num_test_we_use):
for j in range(sequence_length):
test_idx.append(test_we_use_start_idx[i] + j)
test_idx_LFB = []
for i in range(num_test_we_use_LFB):
for j in range(sequence_length):
test_idx_LFB.append(test_we_use_start_idx_LFB[i] + j)
global g_LFB_test
test_feature_loader = DataLoader(
test_dataset,
batch_size=val_batch_size,
sampler=SeqSampler(test_dataset, test_idx_LFB),
num_workers=workers,
pin_memory=False
)
model_LFB = resnet_lstm()
model_LFB = DataParallel(model_LFB)
model_LFB = [Link]("[Link]")
def get_parameter_number(net):
trainable_num = sum([Link]() for p in [Link]() if p.requires_grad)
return trainable_num
total_papa_num = 0
total_papa_num += get_parameter_number(model_LFB)
model_LFB.cuda()
for params in model_LFB.parameters():
params.requires_grad = False
model_LFB.eval()
with torch.no_grad():
for data in test_feature_loader:
inputs, _ = data[0].cuda(), data[1].cuda()
inputs = [Link](-1, sequence_length, 3, 224, 224)
outputs_feature = model_LFB.forward(inputs).[Link]().numpy()
g_LFB_test = [Link]((g_LFB_test, outputs_feature), axis=0)
g_LFB_test = [Link](g_LFB_test)
class Transformer([Link]):
def __init__(self, mstcn_f_maps, mstcn_f_dim, out_features, len_q):
super(Transformer, self).__init__()
self.num_f_maps = mstcn_f_maps
[Link] = mstcn_f_dim
self.num_classes = out_features
self.len_q = len_q
[Link] = Transformer2_3_1(d_model=out_features,
d_ff=mstcn_f_maps, d_k=mstcn_f_maps,
d_v=mstcn_f_maps, n_layers=1,
n_heads=8, len_q=sequence_length)
[Link] = [Link](mstcn_f_dim, out_features, bias=False)
def forward(self, x, long_feature):
out_features = [Link](1, 2)
inputs = []
for i in range(out_features.size(1)):
if i < self.len_q - 1:
input = [Link]((1, self.len_q - 1 - i,
self.num_classes)).cuda()
input = [Link]([input, out_features[:, 0:i + 1]], dim=1)
else:
input = out_features[:, i - self.len_q + 1:i + 1]
[Link](input)
inputs = [Link](inputs, dim=0).squeeze(1)
feas = [Link]([Link](long_feature).transpose(0, 1))
output = [Link](inputs, feas)
return output
def smooth_labels(labels, sigma=2):
return gaussian_filter1d(labels, sigma=sigma, mode='nearest')
def second_stage_inference():
def get_long_feature(start_index, lfb, LFB_length):
long_feature = []
long_feature_each = []
for k in range(LFB_length):
LFB_index = (start_index + k)
LFB_index = int(LFB_index)
long_feature_each.append(lfb[LFB_index])
long_feature.append(long_feature_each)
return long_feature
test_num_each_80, test_start_vidx = [len(g_LFB_test)], [0]
out_features = 7
mstcn_causal_conv = True
mstcn_layers = 8
mstcn_f_maps = 32
mstcn_f_dim= 2048
mstcn_stages = 2
sequence_length = 30
model = [Link](mstcn_stages, mstcn_layers, mstcn_f_maps,
mstcn_f_dim, out_features, mstcn_causal_conv)
model = DataParallel(model)
model = [Link]('[Link]')
[Link]()
[Link]()
model1 = Transformer(mstcn_f_maps, mstcn_f_dim, out_features, sequence_length)
model1 = DataParallel(model1)
model1 = [Link]('[Link]')
[Link]()
[Link]()
[Link].empty_cache()
with torch.no_grad():
long_feature = get_long_feature(start_index=test_start_vidx[0],
lfb=g_LFB_test, LFB_length=test_num_each_80[0])
long_feature = ([Link](long_feature)).cuda()
video_fe = long_feature.transpose(2, 1)
out_features = [Link](video_fe)[-1]
out_features = out_features.squeeze(1)
p_classes1 = model1(out_features, long_feature)
p_classes = p_classes1.squeeze()
_, preds_phase = [Link](p_classes.data, 1)
preds_phase = preds_phase.cpu().numpy()
results = smooth_labels(preds_phase, sigma=2)
return results
def is_mp4_file(file_path):
return file_path.lower().endswith('.mp4')
def check_mp4_file(file_path):
if not is_mp4_file(file_path):
print(f"{file_path} is not an MP4 file.")
return False
try:
cap = [Link](file_path)
if not [Link]():
print(f"{file_path} might be corrupted.")
return False
# 尝试读取第一帧
ret, frame = [Link]()
if not ret:
print(f"{file_path} might be corrupted.")
return False
print(f"{file_path} is a valid MP4 file.")
return True
except Exception as e:
print(f"Error checking {file_path}: {str(e)}")
return False
finally:
[Link]()
def convert_seconds_to_timecode(seconds):
# 将秒数转换为时:分:秒格式，确保每部分为两位数
hours = seconds // 3600
minutes = (seconds % 3600) // 60
seconds = seconds % 60
return f"{hours:02}:{minutes:02}:{seconds:02}"
def process_labels(labels):
label_names = {
0: "术前准备",
1: "Calot 三角解剖",
2: "剪切与夹闭",
3: "胆囊解剖",
4: "胆囊打包",
5: "清理与凝血",
6: "胆囊牵引"
}
processed_labels = []
if not [Link]():
return processed_labels
current_label = labels[0]
start_time = 0
for i in range(1, len(labels)):
if labels[i] != current_label:
end_time = i
processed_labels.append({
"startTime": convert_seconds_to_timecode(start_time),
"labelName": label_names[current_label],
"endTime": convert_seconds_to_timecode(end_time)
})
current_label = labels[i]
start_time = i
# 最后一个阶段
processed_labels.append({
"startTime": convert_seconds_to_timecode(start_time),
"labelName": label_names[current_label],
"endTime": convert_seconds_to_timecode(len(labels))
})
return processed_labels
def main(video_path):
video_dataset = get_data(video_path)
first_stage_inference(video_dataset)
results = second_stage_inference()
return results
if __name__ == "__main__":
parser = [Link]()
parser.add_argument('video_path', type=str, help="Path to the video file")
args = parser.parse_args()
status = 0
b_mp4 = True # 假设判断是否是 mp4 的结果
b_valid = True # 假设判断文件是否有效的结果
video_path = args.video_path
b_mp4 = is_mp4_file(video_path)
b_valid = check_mp4_file(video_path)
if not b_mp4:
status = 1
elif not b_valid:
status = 2
if status == 0:
labels = main(video_path)
processed_labels = process_labels(labels)
else:
processed_labels = []
output = {
"status": status,
"labels": processed_labels
}
# 输出 JSON 结果
with open('[Link]', 'w', encoding='utf-8') as f:
[Link](output, f, ensure_ascii=False, indent=4)
print("status:", status)
# 退出并返回 status
[Link](status)