inpainting_enhancement.mbt

// 图像修复增强（v25.0.0）
// 图像修复、图像去雾、图像去雨、图像去噪

///|
/// 修复算法类型
pub(all) enum InpaintingAlgorithm {
  Telea
  NavierStokes
  FastMarching
  PatchMatch
  DeepLearning
  Unknown
} derive(Eq)

///|
/// 修复结果
pub(all) struct InpaintingResult {
  image : @types.Image
  mask : @types.Image
  algorithm : InpaintingAlgorithm
  inpainted_pixels : Int
  total_mask_pixels : Int
  processing_time_ms : Double
  iterations : Int
}

///|
/// 修复配置
pub(all) struct InpaintingConfig {
  algorithm : InpaintingAlgorithm
  radius : Int
  iterations : Int
  use_edge_preservation : Bool
  tolerance : Double
}

///|
/// 去雾算法类型
pub(all) enum DehazeAlgorithm {
  DarkChannelPrior
  ColorAttenuation
  HomomorphicFiltering
  Retinex
  CLAHE
  Unknown
} derive(Eq)

///|
/// 去雾结果
pub(all) struct DehazeResult {
  image : @types.Image
  algorithm : DehazeAlgorithm
  atmospheric_light : Double
  transmission_map : Array[Double]
  processing_time_ms : Double
}

///|
/// 去雨算法类型
pub(all) enum DerainAlgorithm {
  DictionaryLearning
  ConvolutionalNeuralNetwork
  GAN
  Morphological
  Unknown
} derive(Eq)

///|
/// 去雨结果
pub(all) struct DerainResult {
  image : @types.Image
  algorithm : DerainAlgorithm
  rain_streaks_removed : Int
  processing_time_ms : Double
}

///|
/// 创建默认修复配置
pub fn create_default_inpainting_config() -> InpaintingConfig {
  InpaintingConfig::{
    algorithm: InpaintingAlgorithm::Telea,
    radius: 3,
    iterations: 10,
    use_edge_preservation: true,
    tolerance: 0.1,
  }
}

///|
/// 创建修复掩膜（矩形区域）
pub fn create_rect_mask(
  width : Int,
  height : Int,
  x : Int,
  y : Int,
  rect_width : Int,
  rect_height : Int,
) -> @types.Image {
  let total = width * height
  let mask_data : Array[Byte] = []
  for _ in 0..<total {
    mask_data.push(b'\x00')
  }

  for my in 0..<rect_height {
    for mx in 0..<rect_width {
      let px = x + mx
      let py = y + my
      if px >= 0 && px < width && py >= 0 && py < height {
        mask_data[py * width + px] = b'\xFF'
      }
    }
  }

  @types.Image::{
    width,
    height,
    channels: 1,
    data: Bytes::from_array(mask_data),
  }
}

///|
/// 简单修复（邻域平均法）
pub fn inpaint_simple(
  img : @types.Image,
  mask : @types.Image,
  radius : Int,
) -> @types.Image {
  let w = img.width
  let h = img.height
  let total = w * h
  let out : Array[Byte] = []

  for i in 0..<total {
    out.push(img.data[i])
  }

  // 迭代修复
  for _ in 0..<5 {
    for y in 0..<h {
      for x in 0..<w {
        let idx = y * w + x
        if mask.data[idx] == b'\xFF' {
          // 计算邻域平均值
          let mut sum = 0.0
          let mut count = 0
          for dy in -radius..<(radius + 1) {
            for dx in -radius..<(radius + 1) {
              if dx == 0 && dy == 0 {
                continue
              }
              let nx = x + dx
              let ny = y + dy
              if nx >= 0 && nx < w && ny >= 0 && ny < h {
                let nidx = ny * w + nx
                if mask.data[nidx] == b'\x00' || out[nidx] != img.data[nidx] {
                  sum += out[nidx].to_double()
                  count += 1
                }
              }
            }
          }
          if count > 0 {
            let avg = sum / count.to_double()
            let clamped = if avg < 0.0 {
              0.0
            } else if avg > 255.0 {
              255.0
            } else {
              avg
            }
            out[idx] = clamped.to_int().to_byte()
          }
        }
      }
    }
  }

  @types.Image::{
    width: w,
    height: h,
    channels: img.channels,
    data: Bytes::from_array(out),
  }
}

///|
/// 基于扩散的修复（简化版 Telea）
pub fn inpaint_diffusion(
  img : @types.Image,
  mask : @types.Image,
  radius : Int,
) -> @types.Image {
  let w = img.width
  let h = img.height
  let total = w * h
  let out : Array[Byte] = []

  for i in 0..<total {
    out.push(img.data[i])
  }

  // 从边界向内部扩散
  for _ in 0..<10 {
    for y in 0..<h {
      for x in 0..<w {
        let idx = y * w + x
        if mask.data[idx] == b'\xFF' {
          // 检查是否有已修复的邻域
          let mut has_known = false
          for dy in -1..<2 {
            for dx in -1..<2 {
              if dx == 0 && dy == 0 {
                continue
              }
              let nx = x + dx
              let ny = y + dy
              if nx >= 0 && nx < w && ny >= 0 && ny < h {
                let nidx = ny * w + nx
                if mask.data[nidx] == b'\x00' || out[nidx] != img.data[nidx] {
                  has_known = true
                }
              }
            }
          }

          if has_known {
            // 加权平均（距离加权）
            let mut sum = 0.0
            let mut weight_sum = 0.0
            for dy in -radius..<(radius + 1) {
              for dx in -radius..<(radius + 1) {
                if dx == 0 && dy == 0 {
                  continue
                }
                let nx = x + dx
                let ny = y + dy
                if nx >= 0 && nx < w && ny >= 0 && ny < h {
                  let nidx = ny * w + nx
                  if mask.data[nidx] == b'\x00' || out[nidx] != img.data[nidx] {
                    let dist = (dx * dx + dy * dy).to_double().sqrt()
                    let weight = 1.0 / (1.0 + dist)
                    sum += out[nidx].to_double() * weight
                    weight_sum += weight
                  }
                }
              }
            }
            if weight_sum > 0.0 {
              let avg = sum / weight_sum
              let clamped = if avg < 0.0 {
                0.0
              } else if avg > 255.0 {
                255.0
              } else {
                avg
              }
              out[idx] = clamped.to_int().to_byte()
            }
          }
        }
      }
    }
  }

  @types.Image::{
    width: w,
    height: h,
    channels: img.channels,
    data: Bytes::from_array(out),
  }
}

///|
/// 图像修复
pub fn inpaint(
  img : @types.Image,
  mask : @types.Image,
  config : InpaintingConfig,
) -> InpaintingResult {
  let result = if config.algorithm == InpaintingAlgorithm::Telea {
    inpaint_diffusion(img, mask, config.radius)
  } else {
    inpaint_simple(img, mask, config.radius)
  }

  // 统计掩膜像素数
  let total = mask.width * mask.height
  let mut mask_pixels = 0
  for i in 0..<total {
    if mask.data[i] == b'\xFF' {
      mask_pixels += 1
    }
  }

  InpaintingResult::{
    image: result,
    mask,
    algorithm: config.algorithm,
    inpainted_pixels: mask_pixels,
    total_mask_pixels: mask_pixels,
    processing_time_ms: 0.0,
    iterations: config.iterations,
  }
}

///|
/// 暗通道先验去雾（简化版）
pub fn dehaze_dark_channel(
  img : @types.Image,
  omega : Double,
  t0 : Double,
) -> DehazeResult {
  let w = img.width
  let h = img.height
  let total = w * h
  let patch_size = 15

  // 计算暗通道
  let dark_channel : Array[Double] = []
  for _ in 0..<total {
    dark_channel.push(255.0)
  }

  for y in 0..<h {
    for x in 0..<w {
      let mut min_val = 255.0
      for dy in 0..<patch_size {
        for dx in 0..<patch_size {
          let nx = x + dx - patch_size / 2
          let ny = y + dy - patch_size / 2
          if nx >= 0 && nx < w && ny >= 0 && ny < h {
            let val = img.data[ny * w + nx].to_double()
            if val < min_val {
              min_val = val
            }
          }
        }
      }
      dark_channel[y * w + x] = min_val
    }
  }

  // 估计大气光（暗通道最亮的 0.1% 像素的平均值）
  let mut atmospheric_light = 0.0
  let mut max_dark = 0.0
  for i in 0..<total {
    if dark_channel[i] > max_dark {
      max_dark = dark_channel[i]
    }
  }
  let threshold = max_dark * 0.999
  let mut count = 0
  for i in 0..<total {
    if dark_channel[i] >= threshold {
      atmospheric_light += img.data[i].to_double()
      count += 1
    }
  }
  if count > 0 {
    atmospheric_light /= count.to_double()
  }

  // 计算透射率
  let transmission : Array[Double] = []
  for i in 0..<total {
    let t = 1.0 - omega * (dark_channel[i] / atmospheric_light)
    let clamped = if t < t0 { t0 } else if t > 1.0 { 1.0 } else { t }
    transmission.push(clamped)
  }

  // 去雾
  let out : Array[Byte] = []
  for i in 0..<total {
    let t = transmission[i]
    let pixel = img.data[i].to_double()
    let dehazed = (pixel - atmospheric_light) / t + atmospheric_light
    let clamped = if dehazed < 0.0 {
      0.0
    } else if dehazed > 255.0 {
      255.0
    } else {
      dehazed
    }
    out.push(clamped.to_int().to_byte())
  }

  DehazeResult::{
    image: @types.Image::{
      width: w,
      height: h,
      channels: img.channels,
      data: Bytes::from_array(out),
    },
    algorithm: DehazeAlgorithm::DarkChannelPrior,
    atmospheric_light,
    transmission_map: transmission,
    processing_time_ms: 0.0,
  }
}

///|
/// CLAHE 去雾（对比度受限自适应直方图均衡化）
pub fn dehaze_clahe(
  img : @types.Image,
  clip_limit : Double,
  tile_size : Int,
) -> @types.Image {
  ignore(clip_limit)
  ignore(tile_size)
  let w = img.width
  let h = img.height
  let total = w * h

  // 简化：全局直方图均衡化
  let histogram : Array[Int] = []
  for _ in 0..<256 {
    histogram.push(0)
  }
  for i in 0..<total {
    histogram[img.data[i].to_int()] += 1
  }

  // 累积分布函数
  let cdf : Array[Double] = []
  let mut cdf_min = total.to_double()
  let mut sum = 0.0
  for i in 0..<256 {
    sum += histogram[i].to_double()
    cdf.push(sum)
    if histogram[i] > 0 && sum < cdf_min {
      cdf_min = sum
    }
  }

  // 直方图均衡化
  let out : Array[Byte] = []
  for i in 0..<total {
    let pixel = img.data[i].to_int()
    let equalized = if total > 1 {
      ((cdf[pixel] - cdf_min) / (total.to_double() - cdf_min) * 255.0).to_int()
    } else {
      pixel
    }
    let clamped = if equalized < 0 {
      0
    } else if equalized > 255 {
      255
    } else {
      equalized
    }
    out.push(clamped.to_byte())
  }

  @types.Image::{
    width: w,
    height: h,
    channels: img.channels,
    data: Bytes::from_array(out),
  }
}

///|
/// 去雾
pub fn dehaze(img : @types.Image, algorithm : DehazeAlgorithm) -> DehazeResult {
  if algorithm == DehazeAlgorithm::DarkChannelPrior {
    dehaze_dark_channel(img, 0.95, 0.1)
  } else {
    let result = dehaze_clahe(img, 2.0, 8)
    DehazeResult::{
      image: result,
      algorithm,
      atmospheric_light: 0.0,
      transmission_map: [],
      processing_time_ms: 0.0,
    }
  }
}

///|
/// 形态学去雨（简化版）
pub fn derain_morphological(img : @types.Image) -> DerainResult {
  let w = img.width
  let h = img.height
  let _ = w * h

  // 简化：中值滤波去除雨线
  let out : Array[Byte] = []
  for y in 0..<h {
    for x in 0..<w {
      // 3x3 中值滤波
      let neighbors : Array[Int] = []
      for dy in -1..<2 {
        for dx in -1..<2 {
          let nx = x + dx
          let ny = y + dy
          if nx >= 0 && nx < w && ny >= 0 && ny < h {
            neighbors.push(img.data[ny * w + nx].to_int())
          }
        }
      }
      // 排序取中值
      for i in 0..<neighbors.length() {
        for j in 0..<(neighbors.length() - i - 1) {
          if neighbors[j] > neighbors[j + 1] {
            let temp = neighbors[j]
            neighbors[j] = neighbors[j + 1]
            neighbors[j + 1] = temp
          }
        }
      }
      let median = neighbors[neighbors.length() / 2]
      out.push(median.to_byte())
    }
  }

  DerainResult::{
    image: @types.Image::{
      width: w,
      height: h,
      channels: img.channels,
      data: Bytes::from_array(out),
    },
    algorithm: DerainAlgorithm::Morphological,
    rain_streaks_removed: 0,
    processing_time_ms: 0.0,
  }
}

///|
/// 去雨
pub fn derain(img : @types.Image, algorithm : DerainAlgorithm) -> DerainResult {
  ignore(algorithm)
  derain_morphological(img)
}

///|
/// 生成修复报告
pub fn inpainting_report_to_string(result : InpaintingResult) -> String {
  let mut report = "## 图像修复报告\n\n"
  report = report + "### 处理信息\n\n"
  report = report + "| 指标 | 值 |\n|------|------|\n"
  let algo_str = if result.algorithm == InpaintingAlgorithm::Telea {
    "Telea"
  } else if result.algorithm == InpaintingAlgorithm::NavierStokes {
    "Navier-Stokes"
  } else if result.algorithm == InpaintingAlgorithm::FastMarching {
    "快速行进"
  } else if result.algorithm == InpaintingAlgorithm::PatchMatch {
    "PatchMatch"
  } else if result.algorithm == InpaintingAlgorithm::DeepLearning {
    "深度学习"
  } else {
    "未知"
  }
  report = report + "| 算法 | " + algo_str + " |\n"
  report = report +
    "| 修复像素数 | " +
    result.inpainted_pixels.to_string() +
    " |\n"
  report = report +
    "| 掩膜像素总数 | " +
    result.total_mask_pixels.to_string() +
    " |\n"
  report = report + "| 迭代次数 | " + result.iterations.to_string() + " |\n"
  report = report +
    "| 处理耗时 | " +
    result.processing_time_ms.to_string() +
    "ms |\n"

  report
}
src\process\inpainting\inpainting_enhancement.mbt

86.16%
