请求协助在Google Colab中实现基于图片/视频的TensorFlow Lite图像分割示例

一、适配静态图片的修改方案 #

以下代码基于原示例的模型逻辑，替换摄像头输入为本地图片，可直接在Colab中运行：

1. 安装依赖与下载模型 #

!pip install opencv-python numpy pillow matplotlib
!wget https://storage.googleapis.com/download.tensorflow.org/models/tflite/deeplabv3_mnv2_pascal_quant_2018_01_29.tflite
!wget https://raw.githubusercontent.com/tensorflow/examples/master/lite/examples/image_segmentation/raspberry_pi/labelmap.txt

2. 图片分割核心代码 #

import cv2
import numpy as np
import tensorflow as tf
from matplotlib import pyplot as plt

# 加载TFLite模型与标签
interpreter = tf.lite.Interpreter(model_path="deeplabv3_mnv2_pascal_quant_2018_01_29.tflite")
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

with open('labelmap.txt', 'r') as f:
    labels = [line.strip() for line in f.readlines()]

# 替换为你上传的图片路径（Colab左侧文件面板上传后复制路径）
image_path = "your_image.jpg"
image = cv2.imread(image_path)
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 预处理：调整图片尺寸匹配模型输入
input_shape = input_details[0]['shape']
resized_image = cv2.resize(image_rgb, (input_shape[1], input_shape[2]))
input_data = np.expand_dims(resized_image, axis=0)

# 执行推理
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()
output_data = interpreter.get_tensor(output_details[0]['index'])
segmentation_mask = np.argmax(output_data, axis=3)[0]

# 生成分割结果叠加图
def create_overlay(image, mask):
    color_map = np.array([
        [0,0,0],[128,0,0],[0,128,0],[128,128,0],[0,0,128],[128,0,128],
        [0,128,128],[128,128,128],[64,0,0],[192,0,0],[64,128,0],[192,128,0],
        [64,0,128],[192,0,128],[64,128,128],[192,128,128],[0,64,0],[128,64,0],
        [0,192,0],[128,192,0],[0,64,128]
    ], dtype=np.uint8)
    mask_colored = color_map[mask]
    return cv2.addWeighted(image, 0.5, mask_colored, 0.5, 0)

# 调整掩码尺寸匹配原图并生成结果
mask_resized = cv2.resize(segmentation_mask, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
overlay_image = create_overlay(image_rgb, mask_resized)

# 展示结果
plt.figure(figsize=(12,6))
plt.subplot(121)
plt.imshow(image_rgb)
plt.title('Original Image')
plt.subplot(122)
plt.imshow(overlay_image)
plt.title('Segmentation Result')
plt.show()

二、适配视频/GIF的修改方案 #

1. 视频分割代码 #

import cv2
import numpy as np
import tensorflow as tf

# 加载模型（同图片示例）
interpreter = tf.lite.Interpreter(model_path="deeplabv3_mnv2_pascal_quant_2018_01_29.tflite")
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

color_map = np.array([
    [0,0,0],[128,0,0],[0,128,0],[128,128,0],[0,0,128],[128,0,128],
    [0,128,128],[128,128,128],[64,0,0],[192,0,0],[64,128,0],[192,128,0],
    [64,0,128],[192,0,128],[64,128,128],[192,128,128],[0,64,0],[128,64,0],
    [0,192,0],[128,192,0],[0,64,128]
], dtype=np.uint8)

def process_frame(frame):
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    resized_frame = cv2.resize(frame_rgb, (input_details[0]['shape'][1], input_details[0]['shape'][2]))
    input_data = np.expand_dims(resized_frame, axis=0)
    
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()
    output_data = interpreter.get_tensor(output_details[0]['index'])
    mask = np.argmax(output_data, axis=3)[0]
    
    mask_resized = cv2.resize(mask, (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_NEAREST)
    mask_colored = color_map[mask_resized]
    overlay = cv2.addWeighted(frame_rgb, 0.5, mask_colored, 0.5, 0)
    return cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR)

# 替换为你上传的视频路径
video_path = "your_video.mp4"
cap = cv2.VideoCapture(video_path)
fps = cap.get(cv2.CAP_PROP_FPS)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

# 保存处理后的视频
output_path = "segmented_video.mp4"
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
    processed_frame = process_frame(frame)
    out.write(processed_frame)

cap.release()
out.release()
print(f"处理完成，结果已保存到 {output_path}")

2. GIF分割代码 #

from PIL import Image, ImageSequence
import cv2
import numpy as np

# 加载模型（同上述示例）
interpreter = tf.lite.Interpreter(model_path="deeplabv3_mnv2_pascal_quant_2018_01_29.tflite")
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

color_map = np.array([
    [0,0,0],[128,0,0],[0,128,0],[128,128,0],[0,0,128],[128,0,128],
    [0,128,128],[128,128,128],[64,0,0],[192,0,0],[64,128,0],[192,128,0],
    [64,0,128],[192,0,128],[64,128,128],[192,128,128],[0,64,0],[128,64,0],
    [0,192,0],[128,192,0],[0,64,128]
], dtype=np.uint8)

def process_frame(frame):
    frame_rgb = np.array(frame.convert('RGB'))
    resized_frame = cv2.resize(frame_rgb, (input_details[0]['shape'][1], input_details[0]['shape'][2]))
    input_data = np.expand_dims(resized_frame, axis=0)
    
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()
    output_data = interpreter.get_tensor(output_details[0]['index'])
    mask = np.argmax(output_data, axis=3)[0]
    
    mask_resized = cv2.resize(mask, (frame_rgb.shape[1], frame_rgb.shape[0]), interpolation=cv2.INTER_NEAREST)
    mask_colored = color_map[mask_resized]
    overlay = cv2.addWeighted(frame_rgb, 0.5, mask_colored, 0.5, 0)
    return Image.fromarray(overlay)

# 替换为你上传的GIF路径
gif_path = "your_gif.gif"
gif = Image.open(gif_path)
processed_frames = []

for frame in ImageSequence.Iterator(gif):
    processed_frames.append(process_frame(frame))

# 保存处理后的GIF
processed_frames[0].save('segmented_gif.gif', save_all=True, append_images=processed_frames[1:], loop=0)
print("GIF处理完成，结果已保存为 segmented_gif.gif")

三、关键修改说明 #

• 移除了原代码中依赖树莓派摄像头的picamera逻辑，替换为本地文件读取
• 保留原模型的推理核心，仅调整输入预处理（尺寸适配）与输出后处理（掩码可视化）
• 适配Colab环境，使用Matplotlib展示图片结果，用OpenCV/PIL处理视频/GIF并保存输出

内容的提问来源于stack exchange，提问作者Pepo Papo