08. 计算机视觉进阶

from ultralytics import YOLO

# 加载预训练模型
model = YOLO('yolov8n.pt')  # nano 版本，轻量快速

# 目标检测
results = model('image.jpg')
for r in results:
    boxes = r.boxes  # 边界框
    for box in boxes:
        cls = int(box.cls)          # 类别
        conf = float(box.conf)       # 置信度
        xyxy = box.xyxy[0].tolist()  # [x1, y1, x2, y2]
        print(f"类别: {model.names[cls]}, 置信度: {conf:.3f}, 位置: {xyxy}")

# 训练自定义数据集
model = YOLO('yolov8n.pt')
model.train(data='custom_data.yaml', epochs=100, imgsz=640)

# 导出 ONNX
model.export(format='onnx')

# COCO 格式
{
    "image_id": 1,
    "category_id": 1,  # person=1
    "bbox": [x, y, width, height],  # 左上角 + 宽高
    "score": 0.95,
    "area": width * height
}

# Pascal VOC 格式
# xmin, ymin, xmax, ymax

from ultralytics.utils.metrics import bbox_iou

iou = bbox_iou(pred_box, true_box, xywh=True, CIoU=True)

from torchvision.models.segmentation import fcn_resnet50, deeplabv3_resnet50

# DeepLabV3+
model = deeplabv3_resnet50(pretrained=True)
model.eval()

# 输入图像
from torchvision import transforms
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                        std=[0.229, 0.224, 0.225])
])

input_tensor = transform(image).unsqueeze(0)
output = model(input_tensor)['out']  # (1, 21, H, W)
predictions = output.argmax(1)       # 每个像素的类别

from torchvision.models.detection import maskrcnn_resnet50_fpn

model = maskrcnn_resnet50_fpn(pretrained=True)
model.eval()

predictions = model([transform(image)])
# predictions[0]['boxes']: 边界框
# predictions[0]['labels']: 类别
# predictions[0]['scores']: 置信度
# predictions[0]['masks']: 分割掩膜 (N, 1, H, W)

import torch.nn as nn

class UNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=1):
        super().__init__()
        
        # 编码器
        self.enc1 = self.conv_block(in_channels, 64)
        self.enc2 = self.conv_block(64, 128)
        self.enc3 = self.conv_block(128, 256)
        self.enc4 = self.conv_block(256, 512)
        self.pool = nn.MaxPool2d(2)
        
        # 瓶颈
        self.bottleneck = self.conv_block(512, 1024)
        
        # 解码器
        self.up4 = nn.ConvTranspose2d(1024, 512, 2, stride=2)
        self.dec4 = self.conv_block(1024, 512)
        self.up3 = nn.ConvTranspose2d(512, 256, 2, stride=2)
        self.dec3 = self.conv_block(512, 256)
        self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
        self.dec2 = self.conv_block(256, 128)
        self.up1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
        self.dec1 = self.conv_block(128, 64)
        
        self.final = nn.Conv2d(64, out_channels, 1)
    
    def conv_block(self, in_ch, out_ch):
        return nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )
    
    def forward(self, x):
        # 编码
        e1 = self.enc1(x)
        e2 = self.enc2(self.pool(e1))
        e3 = self.enc3(self.pool(e2))
        e4 = self.enc4(self.pool(e3))
        
        # 瓶颈
        b = self.bottleneck(self.pool(e4))
        
        # 解码 + Skip Connection
        d4 = self.up4(b)
        d4 = torch.cat([d4, e4], dim=1)
        d4 = self.dec4(d4)
        
        d3 = self.up3(d4)
        d3 = torch.cat([d3, e3], dim=1)
        d3 = self.dec3(d3)
        
        d2 = self.up2(d3)
        d2 = torch.cat([d2, e2], dim=1)
        d2 = self.dec2(d2)
        
        d1 = self.up1(d2)
        d1 = torch.cat([d1, e1], dim=1)
        d1 = self.dec1(d1)
        
        return self.final(d1)

from paddleocr import PaddleOCR

# 初始化
ocr = PaddleOCR(use_angle_cls=True, lang='ch', 
                det_model_dir='./models/det',
                rec_model_dir='./models/rec')

# 识别
result = ocr.ocr('document.jpg', cls=True)

for line in result[0]:
    box = line[0]      # 文字框坐标
    text = line[1][0]  # 识别文本
    confidence = line[1][1]
    print(f"文本: {text}, 置信度: {confidence:.3f}")

import easyocr

reader = easyocr.Reader(['ch_sim', 'en'])
results = reader.readtext('image.jpg')

for (bbox, text, conf) in results:
    print(f"{text} ({conf:.2f})")

# 使用 PaddleOCR 训练自定义数据集
# 数据格式：
# image_path\t[{transcription: "文字", points: [[x1,y1],...]}]

# 检测模型训练
!python tools/train.py -c configs/det/ch_PP-OCRv4/ch_PP-OCRv4_det_student.yml

# 识别模型训练
!python tools/train.py -c configs/rec/PP-OCRv4/ch_PP-OCRv4_rec.yml

import torch.nn as nn

# 生成器
class Generator(nn.Module):
    def __init__(self, latent_dim=100, img_size=28):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(latent_dim, 256),
            nn.LeakyReLU(0.2),
            nn.BatchNorm1d(256),
            nn.Linear(256, 512),
            nn.LeakyReLU(0.2),
            nn.BatchNorm1d(512),
            nn.Linear(512, 1024),
            nn.LeakyReLU(0.2),
            nn.BatchNorm1d(1024),
            nn.Linear(1024, img_size * img_size),
            nn.Tanh()
        )
    
    def forward(self, z):
        img = self.model(z)
        return img.view(img.size(0), 1, 28, 28)

# 判别器
class Discriminator(nn.Module):
    def __init__(self, img_size=28):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(img_size * img_size, 512),
            nn.LeakyReLU(0.2),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )
    
    def forward(self, img):
        flat = img.view(img.size(0), -1)
        return self.model(flat)

# 训练循环
def train_gan(generator, discriminator, dataloader, epochs=50):
    g_optimizer = torch.optim.Adam(generator.parameters(), lr=0.0002)
    d_optimizer = torch.optim.Adam(discriminator.parameters(), lr=0.0002)
    criterion = nn.BCELoss()
    
    for epoch in range(epochs):
        for real_imgs, _ in dataloader:
            batch_size = real_imgs.size(0)
            
            # 标签
            real = torch.ones(batch_size, 1)
            fake = torch.zeros(batch_size, 1)
            
            # 训练判别器
            d_optimizer.zero_grad()
            real_loss = criterion(discriminator(real_imgs), real)
            
            z = torch.randn(batch_size, 100)
            fake_imgs = generator(z)
            fake_loss = criterion(discriminator(fake_imgs.detach()), fake)
            
            d_loss = (real_loss + fake_loss) / 2
            d_loss.backward()
            d_optimizer.step()
            
            # 训练生成器
            g_optimizer.zero_grad()
            g_loss = criterion(discriminator(fake_imgs), real)
            g_loss.backward()
            g_optimizer.step()

# 将类别标签作为输入条件
class ConditionalGenerator(nn.Module):
    def __init__(self, latent_dim, num_classes, img_size):
        super().__init__()
        self.label_emb = nn.Embedding(num_classes, num_classes)
        self.model = nn.Sequential(
            nn.Linear(latent_dim + num_classes, 256),
            # ...
        )
    
    def forward(self, z, labels):
        label_input = self.label_emb(labels)
        x = torch.cat([z, label_input], dim=1)
        return self.model(x)

from diffusers import StableDiffusionPipeline
import torch

# 加载 Stable Diffusion
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16
).to("cuda")

# 文生图
prompt = "A beautiful sunset over mountains, digital art"
image = pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
image.save("output.png")

# 图生图
from diffusers import StableDiffusionImg2ImgPipeline

pipe_img2img = StableDiffusionImg2ImgPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16
).to("cuda")

init_image = Image.open("input.jpg").convert("RGB").resize((512, 512))
image = pipe_img2img(prompt="turn this into anime style", 
                     image=init_image, strength=0.7).images[0]

from torchvision.models.video import r3d_18, mc3_18

model = r3d_18(pretrained=True)
# 输入: (batch, channels, frames, height, width)
# r3d_18: (N, 3, 16, 112, 112)

from ultralytics import YOLO

model = YOLO('yolov8n.pt')
results = model.track(source='video.mp4', show=True, tracker="bytetrack.yaml")

# TensorRT 导出
import torch
from torch2trt import torch2trt

model = YOLO('yolov8n.pt').model
x = torch.ones((1, 3, 640, 640)).cuda()
model_trt = torch2trt(model, [x])
torch.save(model_trt.state_dict(), "yolov8n_trt.pth")

import torch.nn as nn

class PatchEmbed(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_ch=3, embed_dim=768):
        super().__init__()
        self.patch_size = patch_size
        self.n_patches = (img_size // patch_size) ** 2  # 14*14=196
        self.proj = nn.Conv2d(in_ch, embed_dim, kernel_size=patch_size, stride=patch_size)
    
    def forward(self, x):
        # x: (B, 3, 224, 224) -> (B, 768, 14, 14) -> (B, 196, 768)
        x = self.proj(x)  # (B, embed_dim, H', W')
        x = x.flatten(2)  # (B, embed_dim, n_patches)
        x = x.transpose(1, 2)  # (B, n_patches, embed_dim)
        return x

# 分类头加上 class token + position embedding
# ViT-Base: patch=16, layers=12, heads=12, dim=768

# 使用 timm 库加载预训练 ViT
import timm

model = timm.create_model('vit_base_patch16_224', pretrained=True, num_classes=10)
# 微调：冻结前 N 层，微调后几层
for name, param in model.named_parameters():
    if 'blocks.0' in name or 'patch_embed' in name:
        param.requires_grad = False

# DETR 核心思想：
# 1. CNN backbone 提取特征
# 2. Transformer Encoder 编码全局上下文
# 3. Transformer Decoder 的 N 个 object queries 并行解码
# 4. 每个 query 预测一个边界框 + 类别
# 5. 匈牙利匹配（bipartite matching）分配 ground truth

# 损失 = 分类损失 + 边界框 L1 损失 + GIoU 损失

from segment_anything import sam_model_registry, SamPredictor

sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h.pth")
predictor = SamPredictor(sam)

# 点提示分割
predictor.set_image(image)
masks, scores, logits = predictor.predict(
    point_coords=[[500, 375]],
    point_labels=[1],  # 1=前景, 0=背景
    multimask_output=True,
)
# 返回 3 个 mask（通常 granularities 不同）

# 框提示分割
masks, scores, _ = predictor.predict(
    box=[100, 100, 400, 400]
)

# 全自动分割：全图所有实例
from segment_anything import SamAutomaticMaskGenerator
mask_generator = SamAutomaticMaskGenerator(sam)
masks = mask_generator.generate(image)

import clip
import torch

model, preprocess = clip.load("ViT-B/32", device="cuda")

# 准备文本提示
classes = ["a photo of a cat", "a photo of a dog", "a photo of a car"]
text_tokens = clip.tokenize(classes).cuda()

with torch.no_grad():
    image_features = model.encode_image(preprocess(image).unsqueeze(0).cuda())
    text_features = model.encode_text(text_tokens)
    
    # 余弦相似度
    logits_per_image = image_features @ text_features.T
    probs = logits_per_image.softmax(dim=-1)
    
# probs: [[0.02, 0.95, 0.03]] -> 预测为 "dog"

# 多模态模型（如 BLIP / LLaVA）将视觉编码器与 LLM 对齐
# LLaVA 架构：CLIP ViT 视觉编码 + Vicuna/LLaMA2 语言解码 + 投影层对齐

# VQA 示例
question = "What is the color of the vehicle in the image?"
# 模型输出: "The vehicle is red."

# 图像描述
# 输入: 图片
# 输出: "A red fire truck with flashing lights parked on a busy street."