feat: add predict api and load model

app.py

@@ -1,11 +1,80 @@
 from fastapi import FastAPI, Request, Response
 from fastapi.responses import JSONResponse
 
-app = FastAPI()
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from PIL import Image
+import io
+import numpy as np
+import os
 
+from models.fpn_inception import FPNInception  # 你自己的模型類別
+
+# =====================
+# 初始化模型
+# =====================
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"🔹 Using device: {device}")
+
+checkpoint_path = os.path.join("model", "deblurgan_v2_latest.pth")
+
+G = FPNInception(norm_layer=nn.InstanceNorm2d).to(device)
+checkpoint = torch.load(checkpoint_path, map_location=device)
+G.load_state_dict(checkpoint["G"], strict=False)
+G.eval()
+print("✅ Model loaded from", checkpoint_path)
+
+# =====================
+# Tile-based 推論函式
+# =====================
+def deblur_image_tiled(model, img, device, tile_size=512, overlap=32):
+    model.eval()
+    w, h = img.size
+    new_w = (w // 32) * 32
+    new_h = (h // 32) * 32
+    if new_w != w or new_h != h:
+        img = img.resize((new_w, new_h), Image.BICUBIC)
+        w, h = new_w, new_h
+
+    img_np = np.array(img).astype(np.float32) / 255.0
+    img_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0).to(device)
+
+    stride = tile_size - overlap
+    tiles_x = list(range(0, w, stride))
+    tiles_y = list(range(0, h, stride))
+    if tiles_x[-1] + tile_size > w:
+        tiles_x[-1] = w - tile_size
+    if tiles_y[-1] + tile_size > h:
+        tiles_y[-1] = h - tile_size
+
+    output = torch.zeros_like(img_tensor)
+    weight = torch.zeros_like(img_tensor)
+
+    with torch.no_grad():
+        for y in tiles_y:
+            for x in tiles_x:
+                patch = img_tensor[:, :, y:y+tile_size, x:x+tile_size]
+                pred = model(patch)
+                output[:, :, y:y+tile_size, x:x+tile_size] += pred
+                weight[:, :, y:y+tile_size, x:x+tile_size] += 1.0
+
+    output /= weight
+    output = torch.clamp(output, 0, 1)
+    out_np = (output.squeeze().permute(1, 2, 0).cpu().numpy() * 255.0).astype(np.uint8)
+    return Image.fromarray(out_np)
+
+# =====================
+# 初始化 FastAPI
+# =====================
+app = FastAPI(title="DeblurGANv2 API")
+
+# =====================
+# API 路由
+# =====================
 @app.get("/")
 def root():
-    return {"Hello": "World!"}
+    return {"message": "DeblurGANv2 API ready!"}
 
 @app.get("/greetjson")
 def greet_json(request: Request, response: Response):
@@ -17,4 +86,25 @@ def greet_json(request: Request, response: Response):
     response.headers["X-Custom-Header"] = "HelloHeader"
 
     # 回傳 JSON
-    return JSONResponse(content={"message": "Hello World", "client": client_host})
+    return JSONResponse(content={"message": "Hello World", "client": client_host})
+
+@app.post("/predict")
+async def predict(file: UploadFile = File(...)):
+    try:
+        # 讀取上傳圖片
+        contents = await file.read()
+        img = Image.open(io.BytesIO(contents)).convert("RGB")
+
+        # 去模糊
+        result = deblur_image_tiled(G, img, device)
+
+        # 輸出為 bytes
+        img_byte_arr = io.BytesIO()
+        result.save(img_byte_arr, format="PNG")
+        img_byte_arr.seek(0)
+
+        # 直接回傳圖片
+        return StreamingResponse(img_byte_arr, media_type="image/png")
+
+    except Exception as e:
+        return JSONResponse({"status": "error", "message": str(e)}, status_code=500)

app_DeblurGan_PyTorch.py

@@ -0,0 +1,159 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Thu Oct 16 12:05:42 2025
+
+@author: ittraining
+"""
+
+# -*- coding: utf-8 -*-
+"""
+Use PyTorch DeblurGAN-v2 (.pth) to deblur images with Tkinter UI
+"""
+
+import os
+import torch
+import torch.nn as nn
+import numpy as np
+from PIL import Image, ImageTk
+from torchvision import transforms
+import tkinter as tk
+from tkinter import filedialog
+
+# ======== 模型定義區 ========
+from models.fpn_inception import FPNInception  # 你需確認這個檔案存在
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"🔹 Using device: {device}")
+
+# 模型 checkpoint 路徑
+checkpoint_dir = os.path.join(os.getcwd(), "model")
+ckpt_path = os.path.join(checkpoint_dir, "deblurgan_v2_latest.pth")
+
+# 初始化模型
+G = FPNInception(norm_layer=nn.InstanceNorm2d).to(device)
+checkpoint = torch.load(ckpt_path, map_location=device)
+G.load_state_dict(checkpoint["G"], strict=False)
+G.eval()
+print("✅ Model loaded from", ckpt_path)
+
+
+# ======== Tile-based 推論函式 ========
+def deblur_image_tiled(model, img, device, tile_size=512, overlap=32):
+    """
+    用 tile-based 方法在 GPU 記憶體有限時推論整張大圖。
+    Args:
+        model: 已載入權重的 DeblurGAN-v2 Generator
+        img: 要處理的影像
+        device: torch.device("cuda" or "cpu")
+        tile_size: 每塊大小（建議 512）
+        overlap: 重疊區域像素數（建議 16~64）
+    """
+    model.eval()
+    
+    # ---- 預處理 ----
+    w, h = img.size
+    
+    # 確保為 32 倍數
+    new_w = (w // 32) * 32
+    new_h = (h // 32) * 32
+    if new_w != w or new_h != h:
+        img = img.resize((new_w, new_h), Image.BICUBIC)
+        w, h = new_w, new_h
+
+    img_np = np.array(img).astype(np.float32) / 255.0
+    img_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0).to(device)
+
+    # ---- 計算 tile 網格 ----
+    stride = tile_size - overlap
+    tiles_x = list(range(0, w, stride))
+    tiles_y = list(range(0, h, stride))
+    if tiles_x[-1] + tile_size > w:
+        tiles_x[-1] = w - tile_size
+    if tiles_y[-1] + tile_size > h:
+        tiles_y[-1] = h - tile_size
+
+    # ---- 準備空白輸出與權重 ----
+    output = torch.zeros_like(img_tensor)
+    weight = torch.zeros_like(img_tensor)
+
+    with torch.no_grad():
+        for y in tiles_y:
+            for x in tiles_x:
+                patch = img_tensor[:, :, y:y+tile_size, x:x+tile_size]
+                pred = model(patch)
+                
+                # 疊加到對應位置
+                output[:, :, y:y+tile_size, x:x+tile_size] += pred
+                weight[:, :, y:y+tile_size, x:x+tile_size] += 1.0
+
+    # ---- 平均化（避免重疊區域過曝）----
+    output /= weight
+    output = torch.clamp(output, 0, 1)
+
+    # ---- 轉回圖片 ----
+    out_np = (output.squeeze().permute(1, 2, 0).cpu().numpy() * 255.0).astype(np.uint8)
+    return Image.fromarray(out_np)
+
+
+# ======== 封裝成 DeblurModel 類 ========
+class DeblurModel:
+    def __init__(self, model):
+        self.model = model
+
+    def predict(self, image_path):
+        img = Image.open(image_path).convert("RGB")
+        out_img = deblur_image_tiled(self.model, img, device, tile_size=512, overlap=32)
+        return out_img
+
+
+# ======== Tkinter GUI ========
+class ImageViewerApp:
+    def __init__(self, root):
+        self.root = root
+        self.root.title("AI Image Deblurring Viewer (PyTorch)")
+        self.root.geometry("1500x700")
+        self.create_gui()
+        self.model = DeblurModel(G)
+
+    def create_gui(self):
+        label_font = ("Helvetica", 16)
+        self.browse_button = tk.Button(
+            self.root, text="Browse Image", command=self.browse_image, font=label_font
+        )
+
+        self.canvas_original = tk.Canvas(self.root, width=480, height=420, bg="lightgray")
+        self.canvas_result = tk.Canvas(self.root, width=480, height=420, bg="lightgray")
+        self.result_label = tk.Label(self.root, text="", font=("Helvetica", 18, "bold"), fg="blue")
+
+        self.browse_button.grid(row=0, column=0, columnspan=2, pady=10)
+        self.canvas_original.grid(row=1, column=0, padx=10, pady=10)
+        self.canvas_result.grid(row=1, column=1, padx=10, pady=10)
+        self.result_label.grid(row=2, column=0, columnspan=2, pady=10)
+
+    def browse_image(self):
+        file_path = filedialog.askopenfilename(
+            filetypes=[("Image files", "*.jpg *.jpeg *.png *.gif *.bmp *.tif")]
+        )
+        if file_path:
+            self.display_images(file_path)
+
+    def display_images(self, image_path):
+        img = Image.open(image_path)
+        img.thumbnail((480, 420))
+        photo = ImageTk.PhotoImage(img)
+        self.canvas_original.create_image(0, 0, anchor="nw", image=photo)
+        self.canvas_original.image = photo
+
+        result_img = self.model.predict(image_path)
+        result_img.thumbnail((480, 420))
+        photo_result = ImageTk.PhotoImage(result_img)
+        self.canvas_result.create_image(0, 0, anchor="nw", image=photo_result)
+        self.canvas_result.image = photo_result
+
+        self.result_label.config(text=f"File: {os.path.basename(image_path)} → Deblurred by DeblurGAN-v2")
+
+
+if __name__ == "__main__":
+    root = tk.Tk()
+    app = ImageViewerApp(root)
+    root.mainloop()

models/fpn_inception.py

@@ -0,0 +1,167 @@
+import torch
+import torch.nn as nn
+from torchsummary import summary
+from pretrainedmodels import inceptionresnetv2
+import torch.nn.functional as F
+
+class FPNHead(nn.Module):
+    def __init__(self, num_in, num_mid, num_out):
+        super().__init__()
+
+        self.block0 = nn.Conv2d(num_in, num_mid, kernel_size=3, padding=1, bias=False)
+        self.block1 = nn.Conv2d(num_mid, num_out, kernel_size=3, padding=1, bias=False)
+
+    def forward(self, x):
+        x = nn.functional.relu(self.block0(x), inplace=True)
+        x = nn.functional.relu(self.block1(x), inplace=True)
+        return x
+
+class ConvBlock(nn.Module):
+    def __init__(self, num_in, num_out, norm_layer):
+        super().__init__()
+
+        self.block = nn.Sequential(nn.Conv2d(num_in, num_out, kernel_size=3, padding=1),
+                                 norm_layer(num_out),
+                                 nn.ReLU(inplace=True))
+
+    def forward(self, x):
+        x = self.block(x)
+        return x
+
+
+class FPNInception(nn.Module):
+
+    def __init__(self, norm_layer=nn.InstanceNorm2d, output_ch=3, num_filters=128, num_filters_fpn=256):
+        super().__init__()
+
+        # Feature Pyramid Network (FPN) with four feature maps of resolutions
+        # 1/4, 1/8, 1/16, 1/32 and `num_filters` filters for all feature maps.
+        self.fpn = FPN(num_filters=num_filters_fpn, norm_layer=norm_layer)
+
+        # The segmentation heads on top of the FPN
+
+        self.head1 = FPNHead(num_filters_fpn, num_filters, num_filters)
+        self.head2 = FPNHead(num_filters_fpn, num_filters, num_filters)
+        self.head3 = FPNHead(num_filters_fpn, num_filters, num_filters)
+        self.head4 = FPNHead(num_filters_fpn, num_filters, num_filters)
+
+        self.smooth = nn.Sequential(
+            nn.Conv2d(4 * num_filters, num_filters, kernel_size=3, padding=1),
+            norm_layer(num_filters),
+            nn.ReLU(),
+        )
+
+        self.smooth2 = nn.Sequential(
+            nn.Conv2d(num_filters, num_filters // 2, kernel_size=3, padding=1),
+            norm_layer(num_filters // 2),
+            nn.ReLU(),
+        )
+
+        self.final = nn.Conv2d(num_filters // 2, output_ch, kernel_size=3, padding=1)
+
+    def unfreeze(self):
+        self.fpn.unfreeze()
+
+    def forward(self, x):
+        map0, map1, map2, map3, map4 = self.fpn(x)
+
+        map4 = nn.functional.interpolate(self.head4(map4), scale_factor=8, mode="nearest")
+        map3 = nn.functional.interpolate(self.head3(map3), scale_factor=4, mode="nearest")
+        map2 = nn.functional.interpolate(self.head2(map2), scale_factor=2, mode="nearest")
+        map1 = nn.functional.interpolate(self.head1(map1), scale_factor=1, mode="nearest")
+
+        smoothed = self.smooth(torch.cat([map4, map3, map2, map1], dim=1))
+        smoothed = nn.functional.interpolate(smoothed, scale_factor=2, mode="nearest")
+        smoothed = self.smooth2(smoothed + map0)
+        smoothed = nn.functional.interpolate(smoothed, scale_factor=2, mode="nearest")
+
+        final = self.final(smoothed)
+        res = torch.tanh(final) + x
+
+        return torch.clamp(res, min = -1,max = 1)
+
+
+class FPN(nn.Module):
+
+    def __init__(self, norm_layer, num_filters=256):
+        """Creates an `FPN` instance for feature extraction.
+        Args:
+          num_filters: the number of filters in each output pyramid level
+          pretrained: use ImageNet pre-trained backbone feature extractor
+        """
+
+        super().__init__()
+        self.inception = inceptionresnetv2(num_classes=1000, pretrained='imagenet')
+
+        self.enc0 = self.inception.conv2d_1a
+        self.enc1 = nn.Sequential(
+            self.inception.conv2d_2a,
+            self.inception.conv2d_2b,
+            self.inception.maxpool_3a,
+        ) # 64
+        self.enc2 = nn.Sequential(
+            self.inception.conv2d_3b,
+            self.inception.conv2d_4a,
+            self.inception.maxpool_5a,
+        )  # 192
+        self.enc3 = nn.Sequential(
+            self.inception.mixed_5b,
+            self.inception.repeat,
+            self.inception.mixed_6a,
+        )   # 1088
+        self.enc4 = nn.Sequential(
+            self.inception.repeat_1,
+            self.inception.mixed_7a,
+        ) #2080
+        self.td1 = nn.Sequential(nn.Conv2d(num_filters, num_filters, kernel_size=3, padding=1),
+                                 norm_layer(num_filters),
+                                 nn.ReLU(inplace=True))
+        self.td2 = nn.Sequential(nn.Conv2d(num_filters, num_filters, kernel_size=3, padding=1),
+                                 norm_layer(num_filters),
+                                 nn.ReLU(inplace=True))
+        self.td3 = nn.Sequential(nn.Conv2d(num_filters, num_filters, kernel_size=3, padding=1),
+                                 norm_layer(num_filters),
+                                 nn.ReLU(inplace=True))
+        self.pad = nn.ReflectionPad2d(1)
+        self.lateral4 = nn.Conv2d(2080, num_filters, kernel_size=1, bias=False)
+        self.lateral3 = nn.Conv2d(1088, num_filters, kernel_size=1, bias=False)
+        self.lateral2 = nn.Conv2d(192, num_filters, kernel_size=1, bias=False)
+        self.lateral1 = nn.Conv2d(64, num_filters, kernel_size=1, bias=False)
+        self.lateral0 = nn.Conv2d(32, num_filters // 2, kernel_size=1, bias=False)
+
+        for param in self.inception.parameters():
+            param.requires_grad = False
+
+    def unfreeze(self):
+        for param in self.inception.parameters():
+            param.requires_grad = True
+
+    def forward(self, x):
+
+        # Bottom-up pathway, from ResNet
+        enc0 = self.enc0(x)
+
+        enc1 = self.enc1(enc0) # 256
+
+        enc2 = self.enc2(enc1) # 512
+
+        enc3 = self.enc3(enc2) # 1024
+
+        enc4 = self.enc4(enc3) # 2048
+
+        # Lateral connections
+
+        lateral4 = self.pad(self.lateral4(enc4))
+        lateral3 = self.pad(self.lateral3(enc3))
+        lateral2 = self.lateral2(enc2)
+        lateral1 = self.pad(self.lateral1(enc1))
+        lateral0 = self.lateral0(enc0)
+
+        # Top-down pathway
+        pad = (1, 2, 1, 2)  # pad last dim by 1 on each side
+        pad1 = (0, 1, 0, 1)
+        map4 = lateral4
+        map3 = self.td1(lateral3 + nn.functional.interpolate(map4, scale_factor=2, mode="nearest"))
+        map2 = self.td2(F.pad(lateral2, pad, "reflect") + nn.functional.interpolate(map3, scale_factor=2, mode="nearest"))
+        map1 = self.td3(lateral1 + nn.functional.interpolate(map2, scale_factor=2, mode="nearest"))
+        return F.pad(lateral0, pad1, "reflect"), map1, map2, map3, map4