Update app.py

app.py

@@ -7,73 +7,111 @@ import numpy as np
 from huggingface_hub import hf_hub_download
 import gradio as gr
 
-# Chuẩn hóa dự đoán
-def normPRED(d):
-    return (d - torch.min(d)) / (torch.max(d) - torch.min(d))
-
-# Hàm suy luận với U2NET
-def inference(net, input_img):
-    input_img = input_img / np.max(input_img)
-    tmpImg = np.zeros((input_img.shape[0], input_img.shape[1], 3))
-    tmpImg[:, :, 0] = (input_img[:, :, 2] - 0.406) / 0.225
-    tmpImg[:, :, 1] = (input_img[:, :, 1] - 0.456) / 0.224
-    tmpImg[:, :, 2] = (input_img[:, :, 0] - 0.485) / 0.229
-    tmpImg = torch.from_numpy(tmpImg.transpose((2, 0, 1))[np.newaxis, :, :, :]).type(torch.FloatTensor)
-    tmpImg = Variable(tmpImg.cuda() if torch.cuda.is_available() else tmpImg)
-    d1, _, _, _, _, _, _ = net(tmpImg)
-    pred = normPRED(1.0 - d1[:, 0, :, :])
-    return pred.cpu().data.numpy().squeeze()
-
-# Hàm chính để xử lý ảnh đầu vào và trả về ảnh chân dung
-def process_image(img, apply_bw, brightness, contrast, saturation, white_balance, hue, highlights_shadows, sharpness, noise_reduction,
-                  apply_adjustments):
-    # Chuyển ảnh sang đen trắng nếu cần
-    if apply_bw:
-        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
-    
-    # Áp dụng các điều chỉnh nếu được yêu cầu
-    if apply_adjustments:
-        # Độ sáng và Độ tương phản
+class PortraitGenerator:
+    def __init__(self):
+        self.u2net = self.load_u2net_model()
+
+    def normPRED(self, d):
+        return (d - torch.min(d)) / (torch.max(d) - torch.min(d))
+
+    def inference(self, input_img):
+        input_img = input_img / np.max(input_img)
+        tmpImg = np.zeros((input_img.shape[0], input_img.shape[1], 3))
+        tmpImg[:, :, 0] = (input_img[:, :, 2] - 0.406) / 0.225
+        tmpImg[:, :, 1] = (input_img[:, :, 1] - 0.456) / 0.224
+        tmpImg[:, :, 2] = (input_img[:, :, 0] - 0.485) / 0.229
+        tmpImg = torch.from_numpy(tmpImg.transpose((2, 0, 1))[np.newaxis, :, :, :]).type(torch.FloatTensor)
+        tmpImg = Variable(tmpImg.cuda() if torch.cuda.is_available() else tmpImg)
+        d1, _, _, _, _, _, _ = self.u2net(tmpImg)
+        pred = self.normPRED(1.0 - d1[:, 0, :, :])
+        return pred.cpu().data.numpy().squeeze()
+
+    def adjust_image(self, img, apply_bw, brightness, contrast, saturation, white_balance, hue, highlights_shadows, sharpness, noise_reduction):
+        # Convert to grayscale if needed
+        if apply_bw:
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+
+        # Adjust brightness and contrast
         img = cv2.convertScaleAbs(img, alpha=contrast / 50.0, beta=brightness - 50)
+
+        # Adjust saturation
+        if saturation != 50:
+            hsv_img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
+            hsv_img[:, :, 1] = np.clip(hsv_img[:, :, 1] * (saturation / 50.0), 0, 255)
+            img = cv2.cvtColor(hsv_img, cv2.COLOR_HSV2BGR)
+
+        # Adjust white balance
+        if white_balance != 50:
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
+            l, a, b = cv2.split(img)
+            a = a * (white_balance / 50.0)
+            b = b * (white_balance / 50.0)
+            img = cv2.merge((l, a, b))
+            img = cv2.cvtColor(img, cv2.COLOR_LAB2BGR)
+
+        # Adjust hue
+        if hue != 50:
+            hsv_img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
+            hsv_img[:, :, 0] = np.clip(hsv_img[:, :, 0] * (hue / 50.0), 0, 180)
+            img = cv2.cvtColor(hsv_img, cv2.COLOR_HSV2BGR)
+
+        # Adjust highlights and shadows
+        if highlights_shadows != 50:
+            img = cv2.convertScaleAbs(img, alpha=1.0, beta=(highlights_shadows - 50) * 5.1)
+
+        # Adjust sharpness
+        if sharpness != 50:
+            kernel = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]]) * (sharpness / 50.0)
+            img = cv2.filter2D(img, -1, kernel)
+
+        # Reduce noise
+        if noise_reduction != 50:
+            img = cv2.fastNlMeansDenoisingColored(img, None, noise_reduction / 50.0 * 10, noise_reduction / 50.0 * 10, 7, 21)
+
+        return img
+
+    def process_image(self, img, apply_bw, brightness, contrast, saturation, white_balance, hue, highlights_shadows, sharpness, noise_reduction, apply_adjustments, generate_final):
+        if not generate_final:
+            preview_img = self.adjust_image(img, apply_bw, brightness, contrast, saturation, white_balance, hue, highlights_shadows, sharpness, noise_reduction)
+            return preview_img
         
-        # Các điều chỉnh khác có thể được thêm vào đây (bão hòa, cân bằng trắng, v.v.)
-        # Placeholder cho các điều chỉnh chi tiết hơn
-    
-    # Chạy suy luận với U2NET
-    result = inference(u2net, img)
-    return (result * 255).astype(np.uint8)
-
-# Tải mô hình từ Hugging Face Hub
-def load_u2net_model():
-    model_path = hf_hub_download(repo_id="Arrcttacsrks/U2net", filename="u2net_portrait.pth", use_auth_token=os.getenv("HF_TOKEN"))
-    net = U2NET(3, 1)
-    net.load_state_dict(torch.load(model_path, map_location="cuda" if torch.cuda.is_available() else "cpu"))
-    net.eval()
-    return net
-
-# Khởi tạo mô hình U2NET
-u2net = load_u2net_model()
-
-# Tạo giao diện với Gradio
-iface = gr.Interface(
-    fn=process_image,
-    inputs=[
-        gr.Image(type="numpy", label="Upload your image"),
-        gr.Checkbox(label="Black & White Image"),
-        gr.Slider(0, 100, value=50, label="Brightness"),
-        gr.Slider(0, 100, value=50, label="Contrast"),
-        gr.Slider(0, 100, value=50, label="Saturation"),
-        gr.Slider(0, 100, value=50, label="White Balance"),
-        gr.Slider(0, 100, value=50, label="Hue"),
-        gr.Slider(0, 100, value=50, label="Highlights and Shadows"),
-        gr.Slider(0, 100, value=50, label="Sharpness"),
-        gr.Slider(0, 100, value=50, label="Noise Reduction"),
-        gr.Checkbox(label="Apply Adjustments")
-    ],
-    outputs=gr.Image(type="numpy", label="Portrait Result"),
-    title="Portrait Generation with U2NET",
-    description="Upload an image to generate its portrait with optional adjustments."
-)
-
-iface.launch()
+        adjusted_img = self.adjust_image(img, apply_bw, brightness, contrast, saturation, white_balance, hue, highlights_shadows, sharpness, noise_reduction)
+        result = self.inference(adjusted_img)
+        return (result * 255).astype(np.uint8)
+
+    def load_u2net_model(self):
+        model_path = hf_hub_download(repo_id="Arrcttacsrks/U2net", filename="u2net_portrait.pth", use_auth_token=os.getenv("HF_TOKEN"))
+        net = U2NET(3, 1)
+        net.load_state_dict(torch.load(model_path, map_location="cuda" if torch.cuda.is_available() else "cpu"))
+        net.eval()
+        return net
+
+def main():
+    portrait_generator = PortraitGenerator()
+
+    iface = gr.Interface(
+        fn=portrait_generator.process_image,
+        inputs=[
+            gr.Image(type="numpy", label="Upload your image"),
+            gr.Checkbox(label="Black & White Image"),
+            gr.Slider(0, 100, value=50, label="Brightness"),
+            gr.Slider(0, 100, value=50, label="Contrast"),
+            gr.Slider(0, 100, value=50, label="Saturation"),
+            gr.Slider(0, 100, value=50, label="White Balance"),
+            gr.Slider(0, 100, value=50, label="Hue"),
+            gr.Slider(0, 100, value=50, label="Highlights and Shadows"),
+            gr.Slider(0, 100, value=50, label="Sharpness"),
+            gr.Slider(0, 100, value=50, label="Noise Reduction"),
+            gr.Checkbox(label="Apply Adjustments"),
+            gr.Checkbox(label="Generate Final Portrait")
+        ],
+        outputs=gr.Image(type="numpy", label="Preview or Portrait Result"),
+        title="Portrait Generation with U2NET",
+        description="Upload an image to generate its portrait with optional adjustments. Enable 'Generate Final Portrait' for final output."
+    )
+
+    iface.launch()
+
+if __name__ == "__main__":
+    main()