Create dress_enhancer.py

model/dress_enhancer.py

@@ -0,0 +1,36 @@
+import torch
+from basicsr.archs.swinir_arch import SwinIR
+import cv2
+import numpy as np
+from PIL import Image
+from torchvision.transforms.functional import to_tensor, to_pil_image
+
+# Load pretrained SwinIR model
+def load_swinir_model():
+    model = SwinIR(
+        upscale=2,
+        in_chans=3,
+        img_size=64,
+        window_size=8,
+        img_range=1.0,
+        depths=[6, 6, 6, 6],
+        embed_dim=180,
+        num_heads=[6, 6, 6, 6],
+        mlp_ratio=2,
+        upsampler='nearest+conv',
+        resi_connection='1conv'
+    )
+    model_path = 'https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/swinir_sr_classical_patch64_x2_48.pth'
+    ckpt = torch.hub.load_state_dict_from_url(model_path, map_location='cpu')
+    model.load_state_dict(ckpt['params'], strict=True)
+    model.eval()
+    return model
+
+swinir = load_swinir_model()
+
+@torch.no_grad()
+def enhance_clothes(image_pil):
+    image = to_tensor(image_pil).unsqueeze(0)
+    output = swinir(image)
+    output_image = output.squeeze().clamp(0, 1)
+    return to_pil_image(output_image)