Update app.py

"cpu로 전환"

app.py

@@ -1,4 +1,3 @@
-import requests
 import streamlit as st
 import cv2
 import numpy as np
@@ -7,7 +6,35 @@ import torch
 from diffusers import StableDiffusionPipeline
 import io
 
-# Function to apply deepfake-like transformation using Stable Diffusion
+def change_hair_to_blonde(image):
+    # Convert to OpenCV format
+    image = np.array(image)
+    # Convert the image to HSV color space
+    hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
+
+    # Define the range for hair color (dark colors)
+    lower_hair = np.array([0, 0, 0])
+    upper_hair = np.array([180, 255, 30])
+
+    # Create a mask for hair
+    mask = cv2.inRange(hsv, lower_hair, upper_hair)
+
+    # Change hair color to blonde (light yellow)
+    hsv[mask > 0] = (30, 255, 200)
+
+    # Convert back to RGB color space
+    image_blonde = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
+    return image_blonde
+
+def add_noise(image):
+    # Convert to OpenCV format
+    image_np = np.array(image)
+    # Generate random noise
+    noise = np.random.normal(0, 25, image_np.shape).astype(np.uint8)
+    # Add noise to the image
+    noisy_image = cv2.add(image_np, noise)
+    return noisy_image
+
 def apply_deepfake(image):
     # Convert PIL image to bytes
     image_bytes = io.BytesIO()
@@ -16,8 +43,8 @@ def apply_deepfake(image):
     
     # Initialize the Stable Diffusion pipeline
     model_id = "CompVis/stable-diffusion-v1-4"
-    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-    pipe = pipe.to("cuda")
+    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
+    pipe = pipe.to("cpu")
 
     prompt = "A person with a different face but same pose and background"
     
@@ -41,9 +68,7 @@ if uploaded_file is not None:
         st.image(image, caption='Original Image.', use_column_width=True)
     elif action == 'B':
         # Add noise to the original image
-        image_np = np.array(image)
-        noise = np.random.normal(0, 25, image_np.shape).astype(np.uint8)
-        noisy_image = cv2.add(image_np, noise)
+        noisy_image = add_noise(image)
         st.image(noisy_image, caption='Image with Noise.', use_column_width=True)
     elif action == 'Deepfake':
         # Apply deepfake transformation