gaur3009 commited on
Commit
b17fae5
·
verified ·
1 Parent(s): 7bb1c5f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +91 -104
app.py CHANGED
@@ -4,132 +4,119 @@ import torch
4
  import cv2
5
  from PIL import Image
6
  from torchvision import transforms
7
- from cloth_segmentation.networks.u2net import U2NET
8
- import matplotlib.colors as mcolors
9
- import traceback
10
 
11
- # Load U²-Net
12
  model_path = "cloth_segmentation/networks/u2net.pth"
13
  model = U2NET(3, 1)
14
- state_dict = torch.load(model_path, map_location=torch.device("cpu"))
15
- state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
16
  model.load_state_dict(state_dict)
17
  model.eval()
18
 
19
- # Util to get BGR color from name
20
- def get_bgr_from_color_name(color_name):
21
- try:
22
- rgb = mcolors.to_rgb(color_name.lower())
23
- return tuple(int(255 * c) for c in rgb[::-1]) # Convert to BGR
24
- except:
25
- return (0, 0, 255) # Default to red
26
-
27
- # Mask refinement
28
- def refine_mask(mask):
29
- close_kernel = np.ones((5, 5), np.uint8)
30
- mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, close_kernel)
31
- erode_kernel = np.ones((3, 3), np.uint8)
32
- mask = cv2.erode(mask, erode_kernel, iterations=1)
33
- mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, close_kernel)
34
- return cv2.GaussianBlur(mask, (5, 5), 1.5)
35
-
36
- # U²-Net segmentation
37
  def segment_dress(image_np):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
38
  transform_pipeline = transforms.Compose([
39
  transforms.ToTensor(),
40
  transforms.Resize((320, 320))
41
  ])
 
42
  image = Image.fromarray(image_np).convert("RGB")
43
  input_tensor = transform_pipeline(image).unsqueeze(0)
44
 
45
  with torch.no_grad():
46
  output = model(input_tensor)[0][0].squeeze().cpu().numpy()
47
 
48
- output = (output - output.min()) / (output.max() - output.min() + 1e-8)
49
- adaptive_thresh = np.mean(output) + 0.2
50
- dress_mask = (output > adaptive_thresh).astype(np.uint8) * 255
51
- return refine_mask(cv2.resize(dress_mask, (image_np.shape[1], image_np.shape[0]), interpolation=cv2.INTER_NEAREST))
52
-
53
- # Optional GrabCut refinement
54
- def apply_grabcut(image_np, dress_mask):
55
- bgd_model = np.zeros((1, 65), np.float64)
56
- fgd_model = np.zeros((1, 65), np.float64)
57
- mask = np.where(dress_mask > 0, cv2.GC_PR_FGD, cv2.GC_BGD).astype('uint8')
58
- coords = cv2.findNonZero(dress_mask)
59
- if coords is not None:
60
- x, y, w, h = cv2.boundingRect(coords)
61
- rect = (x, y, w, h)
62
- cv2.grabCut(image_np, mask, rect, bgd_model, fgd_model, 3, cv2.GC_INIT_WITH_MASK)
63
- refined = np.where((mask == cv2.GC_FGD) | (mask == cv2.GC_PR_FGD), 255, 0).astype("uint8")
64
- return refine_mask(refined)
65
-
66
- # LAB color recoloring
67
- def recolor_dress(image_np, dress_mask, target_color):
68
- target_color_lab = cv2.cvtColor(np.uint8([[target_color]]), cv2.COLOR_BGR2LAB)[0][0]
 
 
69
  img_lab = cv2.cvtColor(image_np, cv2.COLOR_RGB2LAB)
70
-
71
- dress_pixels = img_lab[dress_mask > 0]
72
- if len(dress_pixels) == 0:
73
- return image_np
74
-
75
- mean_L, mean_A, mean_B = np.mean(dress_pixels, axis=0)
76
- a_shift = target_color_lab[1] - mean_A
77
- b_shift = target_color_lab[2] - mean_B
78
-
79
- img_lab[..., 1] = np.clip(img_lab[..., 1] + (dress_mask / 255.0) * a_shift, 0, 255)
80
- img_lab[..., 2] = np.clip(img_lab[..., 2] + (dress_mask / 255.0) * b_shift, 0, 255)
81
-
82
- img_recolored = cv2.cvtColor(img_lab.astype(np.uint8), cv2.COLOR_LAB2RGB)
83
- feathered_mask = cv2.GaussianBlur(dress_mask, (21, 21), 7)
84
- lightness_mask = (img_lab[..., 0] / 255.0) ** 0.7
85
- adaptive_feather = (feathered_mask * lightness_mask).astype(np.uint8)
86
-
87
- return (image_np * (1 - adaptive_feather[..., None] / 255) + img_recolored * (adaptive_feather[..., None] / 255)).astype(np.uint8)
88
-
89
- # Main function with enhanced error handling
90
- def change_dress_color(img, color_prompt):
91
- if img is None or not color_prompt:
92
- return img
93
-
94
- try:
95
- img_np = np.array(img)
96
- target_bgr = get_bgr_from_color_name(color_prompt)
97
-
98
- dress_mask = segment_dress(img_np)
99
- if np.sum(dress_mask) < 1000:
100
- return img
101
-
102
- dress_mask = apply_grabcut(img_np, dress_mask)
103
- img_recolored = recolor_dress(img_np, dress_mask, target_bgr)
104
- return Image.fromarray(img_recolored)
105
-
106
- except Exception as e:
107
- print(f"Error: {e}")
108
- traceback.print_exc()
109
- return img
110
-
111
- # Create a simple function that wraps the main functionality
112
- def process_image(image, color_name):
113
- if image is None:
114
  return None
115
- return change_dress_color(image, color_name)
116
 
117
- # Create Gradio interface with explicit input/output definitions
118
- iface = gr.Interface(
119
- fn=process_image,
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
120
  inputs=[
121
- gr.Image(type="pil", label="Upload Image"),
122
- gr.Textbox(label="Dress Color", placeholder="Enter color name (e.g. blue, lavender, crimson)")
123
  ],
124
- outputs=gr.Image(type="pil", label="Result"),
125
- title="🎨 AI Dress Recolorer",
126
- description="Upload an image of a person wearing a dress and enter a color name to recolor the dress",
127
- examples=[
128
- ["examples/dress1.jpg", "red"],
129
- ["examples/dress2.jpg", "blue"]
130
- ]
131
  )
132
 
133
- # Launch the application
134
  if __name__ == "__main__":
135
- iface.launch(server_name="0.0.0.0", server_port=7860)
 
4
  import cv2
5
  from PIL import Image
6
  from torchvision import transforms
7
+ from cloth_segmentation.networks.u2net import U2NET # Import U²-Net
 
 
8
 
9
+ # Load U²-Net model
10
  model_path = "cloth_segmentation/networks/u2net.pth"
11
  model = U2NET(3, 1)
12
+ state_dict = torch.load(model_path, map_location=torch.device('cpu'))
13
+ state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()} # Remove 'module.' prefix
14
  model.load_state_dict(state_dict)
15
  model.eval()
16
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
17
  def segment_dress(image_np):
18
+ """Segment the dress using U²-Net & refine with Lab color space."""
19
+
20
+ # Convert to Lab space
21
+ img_lab = cv2.cvtColor(image_np, cv2.COLOR_RGB2LAB)
22
+ L, A, B = cv2.split(img_lab)
23
+
24
+ # Use K-means clustering to detect dominant dress region
25
+ pixel_values = img_lab.reshape((-1, 3)).astype(np.float32)
26
+ k = 3 # Three clusters: background, skin, dress
27
+ _, labels, centers = cv2.kmeans(pixel_values, k, None, (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0), 10, cv2.KMEANS_RANDOM_CENTERS)
28
+ labels = labels.reshape(image_np.shape[:2])
29
+
30
+ # Assume dress is the largest non-background cluster
31
+ unique_labels, counts = np.unique(labels, return_counts=True)
32
+ dress_label = unique_labels[np.argmax(counts[1:]) + 1] # Avoid background
33
+
34
+ # Create dress mask
35
+ mask = (labels == dress_label).astype(np.uint8) * 255
36
+
37
+ # Use U²-Net prediction to refine segmentation
38
  transform_pipeline = transforms.Compose([
39
  transforms.ToTensor(),
40
  transforms.Resize((320, 320))
41
  ])
42
+
43
  image = Image.fromarray(image_np).convert("RGB")
44
  input_tensor = transform_pipeline(image).unsqueeze(0)
45
 
46
  with torch.no_grad():
47
  output = model(input_tensor)[0][0].squeeze().cpu().numpy()
48
 
49
+ u2net_mask = (output > 0.5).astype(np.uint8) * 255
50
+ u2net_mask = cv2.resize(u2net_mask, (image_np.shape[1], image_np.shape[0]), interpolation=cv2.INTER_NEAREST)
51
+
52
+ # Combine K-means and U²-Net masks
53
+ refined_mask = cv2.bitwise_and(mask, u2net_mask)
54
+
55
+ return refined_mask
56
+
57
+ def detect_design(image_np, dress_mask):
58
+ """Detect the design part of the dress and separate it from fabric."""
59
+ gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
60
+ edges = cv2.Canny(gray, 50, 150)
61
+
62
+ # Expand detected edges to mask the design area
63
+ kernel = np.ones((5, 5), np.uint8)
64
+ design_mask = cv2.dilate(edges, kernel, iterations=2)
65
+
66
+ # Keep only the design within the dress area
67
+ design_mask = cv2.bitwise_and(design_mask, dress_mask)
68
+ return design_mask
69
+
70
+ def recolor_dress(image_np, mask, design_mask, target_color):
71
+ """Change dress color while preserving texture and design."""
72
  img_lab = cv2.cvtColor(image_np, cv2.COLOR_RGB2LAB)
73
+ target_color_lab = cv2.cvtColor(np.uint8([[target_color]]), cv2.COLOR_BGR2LAB)[0][0]
74
+
75
+ # Preserve lightness (L) and change only chromatic channels (A & B)
76
+ blend_factor = 0.7
77
+ img_lab[..., 1] = np.where((mask > 128) & (design_mask == 0), img_lab[..., 1] * (1 - blend_factor) + target_color_lab[1] * blend_factor, img_lab[..., 1])
78
+ img_lab[..., 2] = np.where((mask > 128) & (design_mask == 0), img_lab[..., 2] * (1 - blend_factor) + target_color_lab[2] * blend_factor, img_lab[..., 2])
79
+
80
+ img_recolored = cv2.cvtColor(img_lab, cv2.COLOR_LAB2RGB)
81
+ return img_recolored
82
+
83
+ def change_dress_color(image_path, color):
84
+ """Change the dress color naturally while keeping textures and design."""
85
+ if image_path is None:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
86
  return None
 
87
 
88
+ img = Image.open(image_path).convert("RGB")
89
+ img_np = np.array(img)
90
+ dress_mask = segment_dress(img_np)
91
+ design_mask = detect_design(img_np, dress_mask)
92
+
93
+ if dress_mask is None:
94
+ return img # No dress detected
95
+
96
+ # Convert the selected color to BGR
97
+ color_map = {
98
+ "Red": (0, 0, 255), "Blue": (255, 0, 0), "Green": (0, 255, 0), "Yellow": (0, 255, 255),
99
+ "Purple": (128, 0, 128), "Orange": (0, 165, 255), "Cyan": (255, 255, 0), "Magenta": (255, 0, 255),
100
+ "White": (255, 255, 255), "Black": (0, 0, 0)
101
+ }
102
+ new_color_bgr = np.array(color_map.get(color, (0, 0, 255)), dtype=np.uint8) # Default to Red
103
+
104
+ # Recolor the dress naturally
105
+ img_recolored = recolor_dress(img_np, dress_mask, design_mask, new_color_bgr)
106
+
107
+ return Image.fromarray(img_recolored)
108
+
109
+ # Gradio Interface
110
+ demo = gr.Interface(
111
+ fn=change_dress_color,
112
  inputs=[
113
+ gr.Image(type="filepath", label="Upload Dress Image"),
114
+ gr.Radio(["Red", "Blue", "Green", "Yellow", "Purple", "Orange", "Cyan", "Magenta", "White", "Black"], label="Choose New Dress Color")
115
  ],
116
+ outputs=gr.Image(type="pil", label="Color Changed Dress"),
117
+ title="Dress Color Changer",
118
+ description="Upload an image of a dress and select a new color to change its appearance naturally while preserving the design."
 
 
 
 
119
  )
120
 
 
121
  if __name__ == "__main__":
122
+ demo.launch()