Spaces:
Running
on
Zero
Running
on
Zero
Upload 2 files
Browse files- app.py +543 -0
- style_transfer.py +752 -0
app.py
ADDED
@@ -0,0 +1,543 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import numpy as np
|
3 |
+
import torch
|
4 |
+
import torch.nn as nn
|
5 |
+
import gradio as gr
|
6 |
+
import time
|
7 |
+
import spaces
|
8 |
+
import timm
|
9 |
+
from torchvision.ops import nms, box_iou
|
10 |
+
import torch.nn.functional as F
|
11 |
+
from torchvision import transforms
|
12 |
+
from PIL import Image, ImageDraw, ImageFont, ImageFilter
|
13 |
+
from breed_health_info import breed_health_info
|
14 |
+
from breed_noise_info import breed_noise_info
|
15 |
+
from dog_database import get_dog_description
|
16 |
+
from scoring_calculation_system import UserPreferences, calculate_compatibility_score
|
17 |
+
from recommendation_html_format import format_recommendation_html, get_breed_recommendations
|
18 |
+
from history_manager import UserHistoryManager
|
19 |
+
from search_history import create_history_tab, create_history_component
|
20 |
+
from styles import get_css_styles
|
21 |
+
from breed_detection import create_detection_tab
|
22 |
+
from breed_comparison import create_comparison_tab
|
23 |
+
from breed_recommendation import create_recommendation_tab
|
24 |
+
from breed_visualization import create_visualization_tab
|
25 |
+
from style_transfer import DogStyleTransfer, create_style_transfer_tab
|
26 |
+
from html_templates import (
|
27 |
+
format_description_html,
|
28 |
+
format_single_dog_result,
|
29 |
+
format_multiple_breeds_result,
|
30 |
+
format_unknown_breed_message,
|
31 |
+
format_not_dog_message,
|
32 |
+
format_hint_html,
|
33 |
+
format_multi_dog_container,
|
34 |
+
format_breed_details_html,
|
35 |
+
get_color_scheme,
|
36 |
+
get_akc_breeds_link
|
37 |
+
)
|
38 |
+
from model_architecture import BaseModel, dog_breeds
|
39 |
+
from urllib.parse import quote
|
40 |
+
from ultralytics import YOLO
|
41 |
+
import asyncio
|
42 |
+
import traceback
|
43 |
+
|
44 |
+
history_manager = UserHistoryManager()
|
45 |
+
|
46 |
+
class ModelManager:
|
47 |
+
"""
|
48 |
+
Singleton class for managing model instances and device allocation
|
49 |
+
specifically designed for Hugging Face Spaces deployment.
|
50 |
+
"""
|
51 |
+
_instance = None
|
52 |
+
_initialized = False
|
53 |
+
_yolo_model = None
|
54 |
+
_breed_model = None
|
55 |
+
_device = None
|
56 |
+
|
57 |
+
def __new__(cls):
|
58 |
+
if cls._instance is None:
|
59 |
+
cls._instance = super().__new__(cls)
|
60 |
+
return cls._instance
|
61 |
+
|
62 |
+
def __init__(self):
|
63 |
+
if not ModelManager._initialized:
|
64 |
+
self._device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
65 |
+
ModelManager._initialized = True
|
66 |
+
|
67 |
+
@property
|
68 |
+
def device(self):
|
69 |
+
if self._device is None:
|
70 |
+
self._device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
71 |
+
return self._device
|
72 |
+
|
73 |
+
@property
|
74 |
+
def yolo_model(self):
|
75 |
+
if self._yolo_model is None:
|
76 |
+
self._yolo_model = YOLO('yolov8x.pt')
|
77 |
+
return self._yolo_model
|
78 |
+
|
79 |
+
@property
|
80 |
+
def breed_model(self):
|
81 |
+
if self._breed_model is None:
|
82 |
+
self._breed_model = BaseModel(
|
83 |
+
num_classes=len(dog_breeds),
|
84 |
+
device=self.device
|
85 |
+
).to(self.device)
|
86 |
+
|
87 |
+
checkpoint = torch.load(
|
88 |
+
'ConvNextV2Base_best_model.pth',
|
89 |
+
map_location=self.device
|
90 |
+
)
|
91 |
+
self._breed_model.load_state_dict(checkpoint['base_model'], strict=False)
|
92 |
+
self._breed_model.eval()
|
93 |
+
return self._breed_model
|
94 |
+
|
95 |
+
# Initialize model manager
|
96 |
+
model_manager = ModelManager()
|
97 |
+
|
98 |
+
def preprocess_image(image):
|
99 |
+
"""Preprocesses images for model input"""
|
100 |
+
if isinstance(image, np.ndarray):
|
101 |
+
image = Image.fromarray(image)
|
102 |
+
|
103 |
+
transform = transforms.Compose([
|
104 |
+
transforms.Resize((224, 224)),
|
105 |
+
transforms.ToTensor(),
|
106 |
+
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
107 |
+
])
|
108 |
+
|
109 |
+
return transform(image).unsqueeze(0)
|
110 |
+
|
111 |
+
@spaces.GPU
|
112 |
+
def predict_single_dog(image):
|
113 |
+
"""Predicts dog breed for a single image"""
|
114 |
+
image_tensor = preprocess_image(image).to(model_manager.device)
|
115 |
+
|
116 |
+
with torch.no_grad():
|
117 |
+
logits = model_manager.breed_model(image_tensor)[0]
|
118 |
+
probs = F.softmax(logits, dim=1)
|
119 |
+
|
120 |
+
top5_prob, top5_idx = torch.topk(probs, k=5)
|
121 |
+
breeds = [dog_breeds[idx.item()] for idx in top5_idx[0]]
|
122 |
+
probabilities = [prob.item() for prob in top5_prob[0]]
|
123 |
+
|
124 |
+
sum_probs = sum(probabilities[:3])
|
125 |
+
relative_probs = [f"{(prob/sum_probs * 100):.2f}%" for prob in probabilities[:3]]
|
126 |
+
|
127 |
+
return probabilities[0], breeds[:3], relative_probs
|
128 |
+
|
129 |
+
def enhanced_preprocess(image, is_standing=False, has_overlap=False):
|
130 |
+
"""
|
131 |
+
Enhanced image preprocessing function with special handling for different poses
|
132 |
+
and overlapping cases.
|
133 |
+
"""
|
134 |
+
target_size = 224
|
135 |
+
w, h = image.size
|
136 |
+
|
137 |
+
if is_standing:
|
138 |
+
if h > w * 1.5:
|
139 |
+
new_h = target_size
|
140 |
+
new_w = min(target_size, int(w * (target_size / h)))
|
141 |
+
new_w = max(new_w, int(target_size * 0.6))
|
142 |
+
elif has_overlap:
|
143 |
+
scale = min(target_size/w, target_size/h) * 0.95
|
144 |
+
new_w = int(w * scale)
|
145 |
+
new_h = int(h * scale)
|
146 |
+
else:
|
147 |
+
scale = min(target_size/w, target_size/h)
|
148 |
+
new_w = int(w * scale)
|
149 |
+
new_h = int(h * scale)
|
150 |
+
|
151 |
+
resized = image.resize((new_w, new_h), Image.Resampling.LANCZOS)
|
152 |
+
final_image = Image.new('RGB', (target_size, target_size), (240, 240, 240))
|
153 |
+
paste_x = (target_size - new_w) // 2
|
154 |
+
paste_y = (target_size - new_h) // 2
|
155 |
+
final_image.paste(resized, (paste_x, paste_y))
|
156 |
+
|
157 |
+
return final_image
|
158 |
+
|
159 |
+
@spaces.GPU
|
160 |
+
def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.3):
|
161 |
+
"""
|
162 |
+
Enhanced multiple dog detection with improved bounding box handling and
|
163 |
+
intelligent boundary adjustments.
|
164 |
+
"""
|
165 |
+
results = model_manager.yolo_model(image, conf=conf_threshold, iou=iou_threshold)[0]
|
166 |
+
img_width, img_height = image.size
|
167 |
+
detected_boxes = []
|
168 |
+
|
169 |
+
# Phase 1: Initial detection and processing
|
170 |
+
for box in results.boxes:
|
171 |
+
if box.cls.item() == 16: # Dog class
|
172 |
+
xyxy = box.xyxy[0].tolist()
|
173 |
+
confidence = box.conf.item()
|
174 |
+
x1, y1, x2, y2 = map(int, xyxy)
|
175 |
+
w = x2 - x1
|
176 |
+
h = y2 - y1
|
177 |
+
|
178 |
+
detected_boxes.append({
|
179 |
+
'coords': [x1, y1, x2, y2],
|
180 |
+
'width': w,
|
181 |
+
'height': h,
|
182 |
+
'center_x': (x1 + x2) / 2,
|
183 |
+
'center_y': (y1 + y2) / 2,
|
184 |
+
'area': w * h,
|
185 |
+
'confidence': confidence,
|
186 |
+
'aspect_ratio': w / h if h != 0 else 1
|
187 |
+
})
|
188 |
+
|
189 |
+
if not detected_boxes:
|
190 |
+
return [(image, 1.0, [0, 0, img_width, img_height], False)]
|
191 |
+
|
192 |
+
# Phase 2: Analysis of detection relationships
|
193 |
+
avg_height = sum(box['height'] for box in detected_boxes) / len(detected_boxes)
|
194 |
+
avg_width = sum(box['width'] for box in detected_boxes) / len(detected_boxes)
|
195 |
+
avg_area = sum(box['area'] for box in detected_boxes) / len(detected_boxes)
|
196 |
+
|
197 |
+
def calculate_iou(box1, box2):
|
198 |
+
x1 = max(box1['coords'][0], box2['coords'][0])
|
199 |
+
y1 = max(box1['coords'][1], box2['coords'][1])
|
200 |
+
x2 = min(box1['coords'][2], box2['coords'][2])
|
201 |
+
y2 = min(box1['coords'][3], box2['coords'][3])
|
202 |
+
|
203 |
+
if x2 <= x1 or y2 <= y1:
|
204 |
+
return 0.0
|
205 |
+
|
206 |
+
intersection = (x2 - x1) * (y2 - y1)
|
207 |
+
area1 = box1['area']
|
208 |
+
area2 = box2['area']
|
209 |
+
return intersection / (area1 + area2 - intersection)
|
210 |
+
|
211 |
+
# Phase 3: Processing each detection
|
212 |
+
processed_boxes = []
|
213 |
+
overlap_threshold = 0.2
|
214 |
+
|
215 |
+
for i, box_info in enumerate(detected_boxes):
|
216 |
+
x1, y1, x2, y2 = box_info['coords']
|
217 |
+
w = box_info['width']
|
218 |
+
h = box_info['height']
|
219 |
+
center_x = box_info['center_x']
|
220 |
+
center_y = box_info['center_y']
|
221 |
+
|
222 |
+
# Check for overlaps
|
223 |
+
has_overlap = False
|
224 |
+
for j, other_box in enumerate(detected_boxes):
|
225 |
+
if i != j and calculate_iou(box_info, other_box) > overlap_threshold:
|
226 |
+
has_overlap = True
|
227 |
+
break
|
228 |
+
|
229 |
+
# Adjust expansion strategy
|
230 |
+
base_expansion = 0.03
|
231 |
+
max_expansion = 0.05
|
232 |
+
|
233 |
+
is_standing = h > 1.5 * w
|
234 |
+
is_sitting = 0.8 <= h/w <= 1.2
|
235 |
+
is_abnormal_size = (h * w) > (avg_area * 1.5) or (h * w) < (avg_area * 0.5)
|
236 |
+
|
237 |
+
if has_overlap:
|
238 |
+
h_expansion = w_expansion = base_expansion * 0.8
|
239 |
+
else:
|
240 |
+
if is_standing:
|
241 |
+
h_expansion = min(base_expansion * 1.2, max_expansion)
|
242 |
+
w_expansion = base_expansion
|
243 |
+
elif is_sitting:
|
244 |
+
h_expansion = w_expansion = base_expansion
|
245 |
+
else:
|
246 |
+
h_expansion = w_expansion = base_expansion * 0.9
|
247 |
+
|
248 |
+
# Position compensation
|
249 |
+
if center_x < img_width * 0.2 or center_x > img_width * 0.8:
|
250 |
+
w_expansion *= 0.9
|
251 |
+
|
252 |
+
if is_abnormal_size:
|
253 |
+
h_expansion *= 0.8
|
254 |
+
w_expansion *= 0.8
|
255 |
+
|
256 |
+
# Calculate final bounding box
|
257 |
+
expansion_w = w * w_expansion
|
258 |
+
expansion_h = h * h_expansion
|
259 |
+
|
260 |
+
new_x1 = max(0, center_x - (w + expansion_w)/2)
|
261 |
+
new_y1 = max(0, center_y - (h + expansion_h)/2)
|
262 |
+
new_x2 = min(img_width, center_x + (w + expansion_w)/2)
|
263 |
+
new_y2 = min(img_height, center_y + (h + expansion_h)/2)
|
264 |
+
|
265 |
+
# Crop and process image
|
266 |
+
cropped_image = image.crop((int(new_x1), int(new_y1),
|
267 |
+
int(new_x2), int(new_y2)))
|
268 |
+
|
269 |
+
processed_image = enhanced_preprocess(
|
270 |
+
cropped_image,
|
271 |
+
is_standing=is_standing,
|
272 |
+
has_overlap=has_overlap
|
273 |
+
)
|
274 |
+
|
275 |
+
processed_boxes.append((
|
276 |
+
processed_image,
|
277 |
+
box_info['confidence'],
|
278 |
+
[new_x1, new_y1, new_x2, new_y2],
|
279 |
+
True
|
280 |
+
))
|
281 |
+
|
282 |
+
return processed_boxes
|
283 |
+
|
284 |
+
@spaces.GPU
|
285 |
+
def predict(image):
|
286 |
+
"""
|
287 |
+
Main prediction function that handles both single and multiple dog detection.
|
288 |
+
Args:
|
289 |
+
image: PIL Image or numpy array
|
290 |
+
Returns:
|
291 |
+
tuple: (html_output, annotated_image, initial_state)
|
292 |
+
"""
|
293 |
+
if image is None:
|
294 |
+
return format_hint_html("Please upload an image to start."), None, None
|
295 |
+
|
296 |
+
try:
|
297 |
+
if isinstance(image, np.ndarray):
|
298 |
+
image = Image.fromarray(image)
|
299 |
+
|
300 |
+
# 檢測圖片中的物體
|
301 |
+
dogs = detect_multiple_dogs(image)
|
302 |
+
color_scheme = get_color_scheme(len(dogs) == 1)
|
303 |
+
|
304 |
+
# 準備標註
|
305 |
+
annotated_image = image.copy()
|
306 |
+
draw = ImageDraw.Draw(annotated_image)
|
307 |
+
|
308 |
+
try:
|
309 |
+
font = ImageFont.truetype("arial.ttf", 24)
|
310 |
+
except:
|
311 |
+
font = ImageFont.load_default()
|
312 |
+
|
313 |
+
dogs_info = ""
|
314 |
+
|
315 |
+
# 處理每個檢測到的物體
|
316 |
+
for i, (cropped_image, detection_confidence, box, is_dog) in enumerate(dogs):
|
317 |
+
print(f"Predict processing - Object {i+1}:")
|
318 |
+
print(f" Is dog: {is_dog}")
|
319 |
+
print(f" Detection confidence: {detection_confidence:.4f}")
|
320 |
+
|
321 |
+
# 如果是狗且進行品種預測,在這裡也加入打印語句
|
322 |
+
if is_dog:
|
323 |
+
top1_prob, topk_breeds, relative_probs = predict_single_dog(cropped_image)
|
324 |
+
print(f" Breed prediction - Top probability: {top1_prob:.4f}")
|
325 |
+
print(f" Top breeds: {topk_breeds[:3]}")
|
326 |
+
color = color_scheme if len(dogs) == 1 else color_scheme[i % len(color_scheme)]
|
327 |
+
|
328 |
+
# 繪製框和標籤
|
329 |
+
draw.rectangle(box, outline=color, width=4)
|
330 |
+
label = f"Dog {i+1}" if is_dog else f"Object {i+1}"
|
331 |
+
label_bbox = draw.textbbox((0, 0), label, font=font)
|
332 |
+
label_width = label_bbox[2] - label_bbox[0]
|
333 |
+
label_height = label_bbox[3] - label_bbox[1]
|
334 |
+
|
335 |
+
# 繪製標籤背景和文字
|
336 |
+
label_x = box[0] + 5
|
337 |
+
label_y = box[1] + 5
|
338 |
+
draw.rectangle(
|
339 |
+
[label_x - 2, label_y - 2, label_x + label_width + 4, label_y + label_height + 4],
|
340 |
+
fill='white',
|
341 |
+
outline=color,
|
342 |
+
width=2
|
343 |
+
)
|
344 |
+
draw.text((label_x, label_y), label, fill=color, font=font)
|
345 |
+
|
346 |
+
try:
|
347 |
+
# 首先檢查是否為狗
|
348 |
+
if not is_dog:
|
349 |
+
dogs_info += format_not_dog_message(color, i+1)
|
350 |
+
continue
|
351 |
+
|
352 |
+
# 如果是狗,進行品種預測
|
353 |
+
top1_prob, topk_breeds, relative_probs = predict_single_dog(cropped_image)
|
354 |
+
combined_confidence = detection_confidence * top1_prob
|
355 |
+
|
356 |
+
# 根據信心度決定輸出格式
|
357 |
+
if combined_confidence < 0.15:
|
358 |
+
dogs_info += format_unknown_breed_message(color, i+1)
|
359 |
+
elif top1_prob >= 0.4:
|
360 |
+
breed = topk_breeds[0]
|
361 |
+
description = get_dog_description(breed)
|
362 |
+
if description is None:
|
363 |
+
description = {
|
364 |
+
"Name": breed,
|
365 |
+
"Size": "Unknown",
|
366 |
+
"Exercise Needs": "Unknown",
|
367 |
+
"Grooming Needs": "Unknown",
|
368 |
+
"Care Level": "Unknown",
|
369 |
+
"Good with Children": "Unknown",
|
370 |
+
"Description": f"Identified as {breed.replace('_', ' ')}"
|
371 |
+
}
|
372 |
+
dogs_info += format_single_dog_result(breed, description, color)
|
373 |
+
else:
|
374 |
+
dogs_info += format_multiple_breeds_result(
|
375 |
+
topk_breeds,
|
376 |
+
relative_probs,
|
377 |
+
color,
|
378 |
+
i+1,
|
379 |
+
lambda breed: get_dog_description(breed) or {
|
380 |
+
"Name": breed,
|
381 |
+
"Size": "Unknown",
|
382 |
+
"Exercise Needs": "Unknown",
|
383 |
+
"Grooming Needs": "Unknown",
|
384 |
+
"Care Level": "Unknown",
|
385 |
+
"Good with Children": "Unknown",
|
386 |
+
"Description": f"Identified as {breed.replace('_', ' ')}"
|
387 |
+
}
|
388 |
+
)
|
389 |
+
except Exception as e:
|
390 |
+
print(f"Error formatting results for dog {i+1}: {str(e)}")
|
391 |
+
dogs_info += format_unknown_breed_message(color, i+1)
|
392 |
+
|
393 |
+
# 包裝最終的HTML輸出
|
394 |
+
html_output = format_multi_dog_container(dogs_info)
|
395 |
+
|
396 |
+
# 準備初始狀態
|
397 |
+
initial_state = {
|
398 |
+
"dogs_info": dogs_info,
|
399 |
+
"image": annotated_image,
|
400 |
+
"is_multi_dog": len(dogs) > 1,
|
401 |
+
"html_output": html_output
|
402 |
+
}
|
403 |
+
|
404 |
+
return html_output, annotated_image, initial_state
|
405 |
+
|
406 |
+
except Exception as e:
|
407 |
+
error_msg = f"An error occurred: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
408 |
+
print(error_msg)
|
409 |
+
return format_hint_html(error_msg), None, None
|
410 |
+
|
411 |
+
|
412 |
+
def show_details_html(choice, previous_output, initial_state):
|
413 |
+
"""
|
414 |
+
Generate detailed HTML view for a selected breed.
|
415 |
+
|
416 |
+
Args:
|
417 |
+
choice: str, Selected breed option
|
418 |
+
previous_output: str, Previous HTML output
|
419 |
+
initial_state: dict, Current state information
|
420 |
+
|
421 |
+
Returns:
|
422 |
+
tuple: (html_output, gradio_update, updated_state)
|
423 |
+
"""
|
424 |
+
if not choice:
|
425 |
+
return previous_output, gr.update(visible=True), initial_state
|
426 |
+
|
427 |
+
try:
|
428 |
+
breed = choice.split("More about ")[-1]
|
429 |
+
description = get_dog_description(breed)
|
430 |
+
html_output = format_breed_details_html(description, breed)
|
431 |
+
|
432 |
+
# Update state
|
433 |
+
initial_state["current_description"] = html_output
|
434 |
+
initial_state["original_buttons"] = initial_state.get("buttons", [])
|
435 |
+
|
436 |
+
return html_output, gr.update(visible=True), initial_state
|
437 |
+
|
438 |
+
except Exception as e:
|
439 |
+
error_msg = f"An error occurred while showing details: {e}"
|
440 |
+
print(error_msg)
|
441 |
+
return format_hint_html(error_msg), gr.update(visible=True), initial_state
|
442 |
+
|
443 |
+
def main():
|
444 |
+
with gr.Blocks(css=get_css_styles()) as iface:
|
445 |
+
# Header HTML
|
446 |
+
|
447 |
+
gr.HTML("""
|
448 |
+
<header style='text-align: center; padding: 20px; margin-bottom: 20px;'>
|
449 |
+
<h1 style='font-size: 2.5em; margin-bottom: 10px; color: #2D3748;'>
|
450 |
+
🐾 PawMatch AI
|
451 |
+
</h1>
|
452 |
+
<h2 style='font-size: 1.2em; font-weight: normal; color: #4A5568; margin-top: 5px;'>
|
453 |
+
Your Smart Dog Breed Guide
|
454 |
+
</h2>
|
455 |
+
<div style='width: 50px; height: 3px; background: linear-gradient(90deg, #4299e1, #48bb78); margin: 15px auto;'></div>
|
456 |
+
<p style='color: #718096; font-size: 0.9em;'>
|
457 |
+
Powered by AI • Breed Recognition • Smart Matching • Companion Guide
|
458 |
+
</p>
|
459 |
+
</header>
|
460 |
+
""")
|
461 |
+
|
462 |
+
# 先創建歷史組件實例(但不創建標籤頁)
|
463 |
+
history_component = create_history_component()
|
464 |
+
|
465 |
+
# Initialize style transfor
|
466 |
+
dog_style_transfer = DogStyleTransfer()
|
467 |
+
|
468 |
+
with gr.Tabs():
|
469 |
+
# 1. breed detection
|
470 |
+
example_images = [
|
471 |
+
'Border_Collie.jpg',
|
472 |
+
'Golden_Retriever.jpeg',
|
473 |
+
'Saint_Bernard.jpeg',
|
474 |
+
'Samoyed.jpeg',
|
475 |
+
'French_Bulldog.jpeg'
|
476 |
+
]
|
477 |
+
detection_components = create_detection_tab(predict, example_images)
|
478 |
+
|
479 |
+
# 2. breed comparison
|
480 |
+
comparison_components = create_comparison_tab(
|
481 |
+
dog_breeds=dog_breeds,
|
482 |
+
get_dog_description=get_dog_description,
|
483 |
+
breed_health_info=breed_health_info,
|
484 |
+
breed_noise_info=breed_noise_info
|
485 |
+
)
|
486 |
+
|
487 |
+
# 3. breed recommendation
|
488 |
+
recommendation_components = create_recommendation_tab(
|
489 |
+
UserPreferences=UserPreferences,
|
490 |
+
get_breed_recommendations=get_breed_recommendations,
|
491 |
+
format_recommendation_html=format_recommendation_html,
|
492 |
+
history_component=history_component
|
493 |
+
)
|
494 |
+
|
495 |
+
# 4. Visualization Analysis
|
496 |
+
with gr.Tab("Visualization Analysis"):
|
497 |
+
create_visualization_tab(
|
498 |
+
dog_breeds=dog_breeds,
|
499 |
+
get_dog_description=get_dog_description,
|
500 |
+
calculate_compatibility_score=calculate_compatibility_score,
|
501 |
+
UserPreferences=UserPreferences
|
502 |
+
)
|
503 |
+
|
504 |
+
# 5. Style Transfer tab
|
505 |
+
with gr.Tab("Style Transfer"):
|
506 |
+
style_transfer_components = create_style_transfer_tab(dog_style_transfer)
|
507 |
+
|
508 |
+
|
509 |
+
# 6. History Search
|
510 |
+
create_history_tab(history_component)
|
511 |
+
|
512 |
+
# Footer
|
513 |
+
gr.HTML('''
|
514 |
+
<div style="
|
515 |
+
display: flex;
|
516 |
+
align-items: center;
|
517 |
+
justify-content: center;
|
518 |
+
gap: 20px;
|
519 |
+
padding: 20px 0;
|
520 |
+
">
|
521 |
+
<p style="
|
522 |
+
font-family: 'Arial', sans-serif;
|
523 |
+
font-size: 14px;
|
524 |
+
font-weight: 500;
|
525 |
+
letter-spacing: 2px;
|
526 |
+
background: linear-gradient(90deg, #555, #007ACC);
|
527 |
+
-webkit-background-clip: text;
|
528 |
+
-webkit-text-fill-color: transparent;
|
529 |
+
margin: 0;
|
530 |
+
text-transform: uppercase;
|
531 |
+
display: inline-block;
|
532 |
+
">EXPLORE THE CODE →</p>
|
533 |
+
<a href="https://github.com/Eric-Chung-0511/Learning-Record/tree/main/Data%20Science%20Projects/PawMatchAI" style="text-decoration: none;">
|
534 |
+
<img src="https://img.shields.io/badge/GitHub-PawMatch_AI-007ACC?logo=github&style=for-the-badge">
|
535 |
+
</a>
|
536 |
+
</div>
|
537 |
+
''')
|
538 |
+
|
539 |
+
return iface
|
540 |
+
|
541 |
+
if __name__ == "__main__":
|
542 |
+
iface = main()
|
543 |
+
iface.launch()
|
style_transfer.py
ADDED
@@ -0,0 +1,752 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
from PIL import Image, ImageEnhance
|
3 |
+
import numpy as np
|
4 |
+
import gradio as gr
|
5 |
+
from diffusers import StableDiffusionImg2ImgPipeline
|
6 |
+
import time
|
7 |
+
import os
|
8 |
+
import base64
|
9 |
+
from io import BytesIO
|
10 |
+
|
11 |
+
class DogStyleTransfer:
|
12 |
+
"""
|
13 |
+
Class for handling dog image style transfer using Stable Diffusion.
|
14 |
+
This class manages model loading, image preprocessing, and style transfer operations.
|
15 |
+
"""
|
16 |
+
def __init__(self):
|
17 |
+
self.models = {}
|
18 |
+
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
19 |
+
|
20 |
+
# Check xformers availability
|
21 |
+
self.xformers_available = False
|
22 |
+
try:
|
23 |
+
import xformers
|
24 |
+
self.xformers_available = True
|
25 |
+
print(f"xformers {xformers.__version__} is available and will be used for memory-efficient attention")
|
26 |
+
except ImportError:
|
27 |
+
print("xformers not found - will use default attention mechanism")
|
28 |
+
except Exception as e:
|
29 |
+
print(f"Error checking xformers: {str(e)} - will use default attention mechanism")
|
30 |
+
|
31 |
+
# Define style to model mapping based on availability
|
32 |
+
if self.device == "cuda":
|
33 |
+
self.style_model_mapping = {
|
34 |
+
"Japanese Anime Style": "Linaqruf/anything-v3.0", # Specialized for anime style
|
35 |
+
"Classic Cartoon": "nitrosocke/mo-di-diffusion", # Specialized for Disney style
|
36 |
+
"Oil Painting": "runwayml/stable-diffusion-v1-5", # Original model good for painting styles
|
37 |
+
"Watercolor": "dreamlike-art/dreamlike-photoreal-2.0", # Photorealistic art style model
|
38 |
+
"Cyberpunk": "dreamlike-art/dreamlike-diffusion-1.0" # Dreamlike style model
|
39 |
+
}
|
40 |
+
else:
|
41 |
+
# Lightweight models for CPU mode
|
42 |
+
self.style_model_mapping = {
|
43 |
+
"Japanese Anime Style": "runwayml/stable-diffusion-v1-5",
|
44 |
+
"Classic Cartoon": "runwayml/stable-diffusion-v1-5",
|
45 |
+
"Oil Painting": "runwayml/stable-diffusion-v1-5",
|
46 |
+
"Watercolor": "runwayml/stable-diffusion-v1-5",
|
47 |
+
"Cyberpunk": "runwayml/stable-diffusion-v1-5"
|
48 |
+
}
|
49 |
+
|
50 |
+
# style prompts with each feature
|
51 |
+
self.style_prompts = {
|
52 |
+
"Japanese Anime Style": "masterpiece, highest quality, genuine anime style illustration of a (dog:1.5), (bold anime aesthetics:1.5), (vibrant saturated colors:1.3), clean distinct lineart, stylized simplified features, expressive anime eyes, (preserve exact animal species:1.8), (maintain original animal breed:1.7), distinctive animal characteristics, (iconic anime art style:1.4), dramatic shading, flat color areas with highlight accents, simplified background elements, characteristic anime proportions, retain animal identity while stylizing, professional anime production quality, no watermarks, no signatures, (do not change animal species:1.8)",
|
53 |
+
"Classic Cartoon": "masterpiece, highest quality classic cartoon illustration of a dog, (golden age animation style:1.3), hand-drawn cel animation quality, bold clean outlines, (vibrant solid color fills:1.2), exaggerated expressive features, playful animated poses, classic Disney/Pixar influenced design, professional animation studio quality, simplified but expressive details, perfect smooth linework, rounded stylized forms, cheerful color palette, dynamic motion lines, classic cartoon physics, expressive oversized eyes, joyful personality captured, squash and stretch principles applied, classic cartoon proportions, professional character design, perfect animation keyframe quality, appealing character expression, masterful use of simple shapes, iconic cartoon aesthetic, no watermarks, no signatures",
|
54 |
+
"Oil Painting": "masterpiece, museum quality oil painting of a dog, (impasto technique:1.3), visible textured brushstrokes, layered oil pigments, rich depth of color, classical composition, (dramatic chiaroscuro lighting:1.2), Renaissance painting technique, glazing layers, sophisticated color harmony, warm and cool tones balance, expert painterly details, canvas texture visible, traditional realistic portrait style, fine art quality, gallery exhibition standard, rich shadows and highlights, volumetric form definition, atmospheric perspective, professional oil painting techniques, traditional varnished finish, color complexity with subtle undertones, expertly captured fur textures, strong compositional focus, emotional depth, timeless artistic quality, no watermarks, no signatures",
|
55 |
+
"Watercolor": "masterpiece, highest quality watercolor painting of a dog, (wet-on-wet technique:1.3), flowing color blends, translucent paint layers, visible paper texture, (controlled paint blooms:1.2), delicate color washes, spontaneous paint flow, preserved white spaces, soft color bleeding effects, subtle granulation textures, feathered edges, luminous transparency, loose expressive brushwork, artistic color pooling, gradient color transitions, minimalist background, playful splatter accents, artistic negative space usage, light-filled composition, watercolor paper grain visible, atmospheric color diffusion, professional traditional watercolor techniques, delicate brush details combined with flowing textures, no watermarks, no signatures",
|
56 |
+
"Cyberpunk": "masterpiece, highest quality, hyper-detailed cyberpunk digital art of a dog, (advanced technological integration:1.4), holographic collar interface, bionic limb enhancements, neural implant visuals, data visualization overlay, augmented reality HUD elements, (neon light reflections:1.3), wet street reflections, volumetric fog effects, urban dystopian background, megacity skyline, glowing circuitry details, optical fiber accents, synthetic materials, dramatic neon-lit contrast, cybernetic enhancements, high tech visors, digital distortion effects, information flow visualization, glitchy textures, metallic surfaces with advanced patina, dark atmospheric tone with vibrant neon accents, electrical energy effects, retro-futuristic design elements, near-future technology aesthetic, no watermarks, no signatures"
|
57 |
+
}
|
58 |
+
|
59 |
+
# Feature preservation prompts with weighted emphasis
|
60 |
+
self.feature_preservation = {
|
61 |
+
"common": "faithful representation of original animal species:(1.6), preserve original animal face structure:(1.5), maintain exact species characteristics:(1.4), accurate distinctive features:(1.3), consistent anatomical structure:(1.2), recognizable animal identity",
|
62 |
+
"Japanese Anime Style": "anime style dog with preserved realistic proportions, distinctive dog breed characteristics maintained, dog facial features clearly recognizable",
|
63 |
+
"Classic Cartoon": "cartoon style with accurate dog proportions, characteristic breed features preserved, recognizable dog expressions",
|
64 |
+
"Oil Painting": "oil painting technique while maintaining anatomical accuracy, realistic dog proportions, distinctive breed characteristics",
|
65 |
+
"Watercolor": "watercolor aesthetic with precise breed representation, accurate dog anatomy, distinctive dog features preserved",
|
66 |
+
"Cyberpunk": "cyberpunk elements while maintaining accurate dog proportions, recognizable breed features, true-to-life dog expression"
|
67 |
+
}
|
68 |
+
|
69 |
+
# Negative prompts
|
70 |
+
self.negative_prompts = {
|
71 |
+
"common": "deformed, distorted, disfigured, poorly drawn, bad anatomy, wrong anatomy, extra limbs, missing limbs, floating limbs, disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, watermark, signature, text, change of species, wrong animal species, incorrect animal type, different animal, human features",
|
72 |
+
"dog_specific": "human face, human features, anthropomorphic, humanoid, human-like features, cartoon eyes, unrealistic eyes",
|
73 |
+
"Japanese Anime Style": "photorealistic, 3d render, western cartoon style, pixar style, realistic textured skin",
|
74 |
+
"Classic Cartoon": "anime style, manga, realistic, detailed skin texture, painterly, sketch, watercolor style",
|
75 |
+
"Oil Painting": "flat colors, digital art, cartoon, cell shaded, smooth texture, anime style",
|
76 |
+
"Watercolor": "digital art, 3d render, vector art, perfect linework, hard edges, bold lines",
|
77 |
+
"Cyberpunk": "watercolor paint, oil painting, natural scene, traditional art, vintage style, soft colors",
|
78 |
+
"species_preservation": "species transformation, change of animal type, incorrect animal features, wrong animal proportions, mixed animal characteristics"
|
79 |
+
}
|
80 |
+
|
81 |
+
# Style descriptions for UI display
|
82 |
+
self.style_descriptions = {
|
83 |
+
"Japanese Anime Style": "Characterized by vibrant colors, large expressive eyes, and stylized features common in Japanese animation.",
|
84 |
+
"Classic Cartoon": "Friendly, rounded features with bold outlines and bright colors typical of classic animated films.",
|
85 |
+
"Oil Painting": "Rich textures and depth created through visible brushstrokes and layered color application.",
|
86 |
+
"Watercolor": "Soft, transparent washes of color with flowing transitions and subtle color blending.",
|
87 |
+
"Cyberpunk": "Futuristic sci-fi aesthetic with neon colors, high contrast, and technological elements."
|
88 |
+
}
|
89 |
+
|
90 |
+
# Set model cache path
|
91 |
+
self.model_cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "dog_style_transfer")
|
92 |
+
os.makedirs(self.model_cache_dir, exist_ok=True)
|
93 |
+
|
94 |
+
# Display system info for debugging
|
95 |
+
self._print_system_info()
|
96 |
+
|
97 |
+
def _print_system_info(self):
|
98 |
+
"""Print system information for debugging purposes"""
|
99 |
+
print("\n===== System Information =====")
|
100 |
+
print(f"Device: {self.device}")
|
101 |
+
print(f"PyTorch version: {torch.__version__}")
|
102 |
+
|
103 |
+
if self.device == "cuda":
|
104 |
+
print(f"CUDA available: {torch.cuda.is_available()}")
|
105 |
+
print(f"CUDA version: {torch.version.cuda if hasattr(torch.version, 'cuda') else 'Unknown'}")
|
106 |
+
print(f"GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'Not available'}")
|
107 |
+
print(f"GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB" if torch.cuda.is_available() else "Not available")
|
108 |
+
|
109 |
+
print(f"xformers available: {self.xformers_available}")
|
110 |
+
print("============================\n")
|
111 |
+
|
112 |
+
def load_model(self, style_name):
|
113 |
+
"""Load the appropriate model based on style, handling xformers compatibility"""
|
114 |
+
# Get model ID for the style
|
115 |
+
model_id = self.style_model_mapping.get(style_name, "runwayml/stable-diffusion-v1-5")
|
116 |
+
|
117 |
+
# Check if model is already loaded
|
118 |
+
if model_id not in self.models:
|
119 |
+
print(f"Loading model {model_id} for {style_name} style...")
|
120 |
+
|
121 |
+
try:
|
122 |
+
# Load model with cache directory
|
123 |
+
model = StableDiffusionImg2ImgPipeline.from_pretrained(
|
124 |
+
model_id,
|
125 |
+
cache_dir=self.model_cache_dir,
|
126 |
+
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
|
127 |
+
safety_checker=None # Remove safety checker to improve speed
|
128 |
+
)
|
129 |
+
|
130 |
+
if self.device == "cuda":
|
131 |
+
model = model.to("cuda")
|
132 |
+
# Enable memory optimization
|
133 |
+
model.enable_attention_slicing()
|
134 |
+
|
135 |
+
# Try to enable xformers
|
136 |
+
try:
|
137 |
+
if hasattr(model, 'enable_xformers_memory_efficient_attention'):
|
138 |
+
print("Attempting to enable xformers memory efficient attention...")
|
139 |
+
model.enable_xformers_memory_efficient_attention()
|
140 |
+
print("xformers memory efficient attention enabled successfully!")
|
141 |
+
except Exception as e:
|
142 |
+
print(f"Warning: Could not enable xformers memory efficient attention: {e}")
|
143 |
+
print("Proceeding without xformers optimization - this may use more memory but should still work.")
|
144 |
+
|
145 |
+
# Store model
|
146 |
+
self.models[model_id] = model
|
147 |
+
print(f"Model {model_id} loaded successfully!")
|
148 |
+
except Exception as e:
|
149 |
+
print(f"Error loading model {model_id}: {str(e)}")
|
150 |
+
# Fall back to basic model if specific model fails
|
151 |
+
if model_id != "runwayml/stable-diffusion-v1-5":
|
152 |
+
print("Falling back to default model...")
|
153 |
+
return self.load_model("Oil Painting") # Use generic model as fallback
|
154 |
+
raise
|
155 |
+
|
156 |
+
return self.models[model_id]
|
157 |
+
|
158 |
+
def preprocess_image(self, image, animal_type='dog'):
|
159 |
+
"""Enhanced preprocessing for dog images before style transfer"""
|
160 |
+
# Convert to PIL image if needed
|
161 |
+
if isinstance(image, np.ndarray):
|
162 |
+
# Handle RGBA images by converting to RGB
|
163 |
+
if image.shape[2] == 4:
|
164 |
+
image = image[:, :, :3]
|
165 |
+
image = Image.fromarray(np.uint8(image))
|
166 |
+
|
167 |
+
# Resize while maintaining aspect ratio
|
168 |
+
width, height = image.size
|
169 |
+
max_size = 512 # SD models typically use 512x512 input
|
170 |
+
scaling_factor = min(max_size / width, max_size / height)
|
171 |
+
new_width = int(width * scaling_factor)
|
172 |
+
new_height = int(height * scaling_factor)
|
173 |
+
image = image.resize((new_width, new_height), Image.LANCZOS)
|
174 |
+
|
175 |
+
# Enhance contrast to emphasize dog features
|
176 |
+
enhancer = ImageEnhance.Contrast(image)
|
177 |
+
image = enhancer.enhance(1.2) # Slightly enhance contrast
|
178 |
+
|
179 |
+
# Sharpen to improve detail
|
180 |
+
enhancer = ImageEnhance.Sharpness(image)
|
181 |
+
image = enhancer.enhance(1.3) # Enhance sharpness
|
182 |
+
|
183 |
+
# Pad if not 512x512, instead of cropping
|
184 |
+
if new_width != 512 or new_height != 512:
|
185 |
+
new_img = Image.new("RGB", (512, 512), (255, 255, 255))
|
186 |
+
# Center the resized image
|
187 |
+
offset = ((512 - new_width) // 2, (512 - new_height) // 2)
|
188 |
+
new_img.paste(image, offset)
|
189 |
+
image = new_img
|
190 |
+
|
191 |
+
if animal_type != 'dog':
|
192 |
+
self.feature_preservation['common'] = 'strict preservation of original animal species:(1.8),' + self.feature_preservation["common"]
|
193 |
+
|
194 |
+
return image
|
195 |
+
|
196 |
+
def transform_style(self, image, style_name, strength=0.75, guidance_scale=7.5):
|
197 |
+
"""
|
198 |
+
Transform image to selected style with improved prompts and parameters
|
199 |
+
|
200 |
+
Args:
|
201 |
+
image: Input image
|
202 |
+
style_name: Name of the style to apply
|
203 |
+
strength: Style transformation strength (0-1)
|
204 |
+
guidance_scale: Guidance scale for stable diffusion
|
205 |
+
|
206 |
+
Returns:
|
207 |
+
tuple: (transformed_image, error_message)
|
208 |
+
"""
|
209 |
+
try:
|
210 |
+
if image is None:
|
211 |
+
return None, "Please upload a dog image first!"
|
212 |
+
|
213 |
+
start_time = time.time()
|
214 |
+
print(f"Starting style transfer: {style_name}")
|
215 |
+
|
216 |
+
# Adjust parameters based on style
|
217 |
+
if style_name == "Japanese Anime Style":
|
218 |
+
guidance_scale = 9.0 # Higher guidance for anime style
|
219 |
+
strength = 0.8
|
220 |
+
num_steps = 50
|
221 |
+
elif style_name == "Classic Cartoon":
|
222 |
+
guidance_scale = 8.0
|
223 |
+
strength = 0.75
|
224 |
+
num_steps = 40
|
225 |
+
elif style_name == "Oil Painting" or style_name == "Watercolor":
|
226 |
+
guidance_scale = 8.0 # Medium guidance for art styles
|
227 |
+
strength = 0.85
|
228 |
+
num_steps = 50
|
229 |
+
elif style_name == "Cyberpunk":
|
230 |
+
guidance_scale = 10.0 # Very high guidance for cyberpunk
|
231 |
+
strength = 0.85
|
232 |
+
num_steps = 50
|
233 |
+
else:
|
234 |
+
num_steps = 40
|
235 |
+
|
236 |
+
# Load model for style
|
237 |
+
try:
|
238 |
+
pipe = self.load_model(style_name)
|
239 |
+
except Exception as e:
|
240 |
+
print(f"Failed to load specific model for {style_name}: {str(e)}")
|
241 |
+
# Fall back to default model
|
242 |
+
pipe = self.load_model("Oil Painting")
|
243 |
+
|
244 |
+
# Enhanced image preprocessing
|
245 |
+
pil_image = self.preprocess_image(image)
|
246 |
+
|
247 |
+
# Get style prompt and add feature preservation
|
248 |
+
base_prompt = self.style_prompts.get(style_name, "digital art style, a dog")
|
249 |
+
|
250 |
+
# Feature preservation prompts - combining common and style-specific
|
251 |
+
feature_preservation = f"{self.feature_preservation['common']}, {self.feature_preservation.get(style_name, '')}"
|
252 |
+
|
253 |
+
# Enhanced positive prompt with feature preservation
|
254 |
+
prompt = f"{base_prompt}, {feature_preservation}, (high quality, detailed, sharp focus, professional photography):(1.2)"
|
255 |
+
|
256 |
+
# Use negative prompt - combining common and style-specific
|
257 |
+
negative_prompt = f"{self.negative_prompts['common']}, {self.negative_prompts['dog_specific']}, {self.negative_prompts.get(style_name, '')}"
|
258 |
+
|
259 |
+
print(f"Using prompt: {prompt}")
|
260 |
+
print(f"Using negative prompt: {negative_prompt}")
|
261 |
+
print(f"Transformation parameters - Strength: {strength}, Guidance Scale: {guidance_scale}, Steps: {num_steps}")
|
262 |
+
|
263 |
+
# Limit steps if too large to avoid memory issues
|
264 |
+
if num_steps > 60 and self.device == "cuda":
|
265 |
+
print("Reducing inference steps to save memory")
|
266 |
+
num_steps = 60
|
267 |
+
|
268 |
+
try:
|
269 |
+
# Generate transformed image
|
270 |
+
result = pipe(
|
271 |
+
prompt=prompt,
|
272 |
+
negative_prompt=negative_prompt,
|
273 |
+
image=pil_image,
|
274 |
+
strength=strength,
|
275 |
+
guidance_scale=guidance_scale,
|
276 |
+
num_inference_steps=num_steps
|
277 |
+
).images[0]
|
278 |
+
|
279 |
+
except RuntimeError as e:
|
280 |
+
# Handle CUDA out of memory errors
|
281 |
+
if "CUDA out of memory" in str(e):
|
282 |
+
print("CUDA out of memory error, trying with reduced parameters")
|
283 |
+
# Retry with lower settings
|
284 |
+
return self._retry_with_lower_settings(pipe, prompt, negative_prompt, pil_image, strength, guidance_scale)
|
285 |
+
else:
|
286 |
+
# Try without negative prompt
|
287 |
+
print(f"Error with negative prompt, retrying without it: {str(e)}")
|
288 |
+
try:
|
289 |
+
result = pipe(
|
290 |
+
prompt=prompt,
|
291 |
+
image=pil_image,
|
292 |
+
strength=strength,
|
293 |
+
guidance_scale=guidance_scale,
|
294 |
+
num_inference_steps=30 # Reduce steps
|
295 |
+
).images[0]
|
296 |
+
except Exception as retry_error:
|
297 |
+
print(f"Retry also failed: {str(retry_error)}")
|
298 |
+
raise
|
299 |
+
|
300 |
+
proc_time = time.time() - start_time
|
301 |
+
print(f"Style transfer completed in {proc_time:.2f} seconds")
|
302 |
+
|
303 |
+
return np.array(result), None
|
304 |
+
|
305 |
+
except Exception as e:
|
306 |
+
error_message = str(e)
|
307 |
+
# Provide user-friendly error messages
|
308 |
+
if "xformers" in error_message.lower():
|
309 |
+
print(f"xformers related error: {error_message}")
|
310 |
+
return None, "Style transfer error: xformers optimization unavailable, but functionality not affected. Please click 'Transform Style' button again to continue."
|
311 |
+
elif "CUDA out of memory" in error_message:
|
312 |
+
print(f"CUDA memory error: {error_message}")
|
313 |
+
return None, "GPU memory insufficient. Try reducing parameters or using a smaller image."
|
314 |
+
else:
|
315 |
+
print(f"Error during style transfer: {error_message}")
|
316 |
+
return None, f"Style transfer error: {error_message}"
|
317 |
+
|
318 |
+
|
319 |
+
def _retry_with_lower_settings(self, pipe, prompt, negative_prompt, image, strength, guidance_scale):
|
320 |
+
"""Retry with lower settings when memory is insufficient"""
|
321 |
+
try:
|
322 |
+
# First attempt: Reduce inference steps
|
323 |
+
print("Attempting with lower settings (steps=20)...")
|
324 |
+
result = pipe(
|
325 |
+
prompt=prompt,
|
326 |
+
negative_prompt=negative_prompt,
|
327 |
+
image=image,
|
328 |
+
strength=strength,
|
329 |
+
guidance_scale=guidance_scale,
|
330 |
+
num_inference_steps=20 # Significantly reduce steps
|
331 |
+
).images[0]
|
332 |
+
return np.array(result), None
|
333 |
+
|
334 |
+
except Exception as first_error:
|
335 |
+
# Log first failure
|
336 |
+
print(f"First retry attempt failed: {str(first_error)}")
|
337 |
+
|
338 |
+
# Second attempt: Minimum settings
|
339 |
+
try:
|
340 |
+
print("Attempting with minimum settings (steps=15, strength=0.6)...")
|
341 |
+
result = pipe(
|
342 |
+
prompt=prompt,
|
343 |
+
image=image,
|
344 |
+
strength=0.6, # Lower strength
|
345 |
+
guidance_scale=7.0, # Use standard setting
|
346 |
+
num_inference_steps=15 # Minimum steps
|
347 |
+
).images[0]
|
348 |
+
return np.array(result), None
|
349 |
+
|
350 |
+
except Exception as second_error:
|
351 |
+
# Log all failures
|
352 |
+
print(f"Second retry attempt also failed: {str(second_error)}")
|
353 |
+
print("All retry attempts failed")
|
354 |
+
|
355 |
+
# Return clear error message
|
356 |
+
error_msg = f"Unable to complete style transfer, even with minimal settings: {str(second_error)}"
|
357 |
+
return None, error_msg
|
358 |
+
|
359 |
+
def get_available_styles(self):
|
360 |
+
"""Get all available style options"""
|
361 |
+
return list(self.style_prompts.keys())
|
362 |
+
|
363 |
+
def get_style_description(self, style_name):
|
364 |
+
"""Get description for a specific style"""
|
365 |
+
return self.style_descriptions.get(style_name, "")
|
366 |
+
|
367 |
+
def get_model_info(self, style_name):
|
368 |
+
"""Get the model information for a specific style"""
|
369 |
+
model_id = self.style_model_mapping.get(style_name, "runwayml/stable-diffusion-v1-5")
|
370 |
+
return f"Powered by: {model_id}"
|
371 |
+
|
372 |
+
def get_image_download_link(self, image):
|
373 |
+
"""
|
374 |
+
Generate a data URL for downloading the image
|
375 |
+
|
376 |
+
Args:
|
377 |
+
image: PIL Image or numpy array
|
378 |
+
|
379 |
+
Returns:
|
380 |
+
str: Base64 encoded data URL
|
381 |
+
"""
|
382 |
+
if image is None:
|
383 |
+
return None
|
384 |
+
|
385 |
+
# Convert numpy array to PIL Image if needed
|
386 |
+
if isinstance(image, np.ndarray):
|
387 |
+
image = Image.fromarray(np.uint8(image))
|
388 |
+
|
389 |
+
# Save image to bytes buffer
|
390 |
+
buffer = BytesIO()
|
391 |
+
image.save(buffer, format="PNG")
|
392 |
+
img_str = base64.b64encode(buffer.getvalue()).decode()
|
393 |
+
|
394 |
+
return f"data:image/png;base64,{img_str}"
|
395 |
+
|
396 |
+
def create_style_transfer_tab(dog_style_transfer):
|
397 |
+
"""Create style transfer tab with UI components"""
|
398 |
+
|
399 |
+
with gr.Column():
|
400 |
+
gr.Markdown("""
|
401 |
+
# 🎨 Dog Style Transformation
|
402 |
+
|
403 |
+
Transform your dog photos into different artistic styles! Upload a dog picture, choose your preferred style, and create unique artwork.
|
404 |
+
""")
|
405 |
+
|
406 |
+
# Add model info and style description display
|
407 |
+
gr.Markdown("""
|
408 |
+
<div style="background-color: #f0f8ff; padding: 16px; border-radius: 8px; margin: 16px 0; border: 1px solid #cce5ff; box-shadow: 0 3px 10px rgba(0,0,123,0.1);">
|
409 |
+
<h3>🐶 Upload a dog photo and select an artistic style</h3>
|
410 |
+
<p>After uploading your dog photo, the system will transform it into your chosen artistic style. Try different styles to create stunning effects!</p>
|
411 |
+
<p>Our system uses specialized models for each style to ensure the best results.</p>
|
412 |
+
<p style="margin-top: 10px; padding: 8px; background-color: #fff9e6; border-left: 4px solid #ffd966; border-radius: 4px;"><b>⏱️ Patience is a virtue!</b> While AI is working its magic, your dog might have time to learn a new trick or two. The transformation can take up to 30 seconds, depending on how photogenic your furry friend is! 🐾</p>
|
413 |
+
<p style="margin-top: 10px; padding: 8px; background-color: #e6f9ff; border-left: 4px solid #66c2ff; border-radius: 4px;"><b>🤫 A Little Secret:</b> Although we designed this tool for dogs, it can actually transform any photo! Portraits, landscapes, even your favorite teddy bear — feel free to try them all! Just don’t tell the other dogs… they might get jealous! 😉</p>
|
414 |
+
<p style="margin-top: 10px; padding: 8px; background-color: #e6f9e6; border-left: 4px solid #66cc77; border-radius: 4px;"><b>✨ Unlimited Creativity!</b> Sometimes, AI might surprise you with unexpected creative interpretations, adding unique colors or features to your image. ✨</p>
|
415 |
+
</div>
|
416 |
+
""")
|
417 |
+
|
418 |
+
with gr.Row():
|
419 |
+
with gr.Column(scale=1):
|
420 |
+
# Upload image component
|
421 |
+
input_image = gr.Image(
|
422 |
+
label="Upload Dog Photo",
|
423 |
+
type="numpy"
|
424 |
+
)
|
425 |
+
|
426 |
+
style_dropdown = gr.Dropdown(
|
427 |
+
choices=dog_style_transfer.get_available_styles(),
|
428 |
+
value=dog_style_transfer.get_available_styles()[0],
|
429 |
+
label="Select Artistic Style"
|
430 |
+
)
|
431 |
+
|
432 |
+
# Display style description
|
433 |
+
style_description = gr.Markdown(
|
434 |
+
dog_style_transfer.get_style_description(dog_style_transfer.get_available_styles()[0])
|
435 |
+
)
|
436 |
+
|
437 |
+
# Display model info
|
438 |
+
model_info = gr.Markdown(
|
439 |
+
dog_style_transfer.get_model_info(dog_style_transfer.get_available_styles()[0])
|
440 |
+
)
|
441 |
+
|
442 |
+
with gr.Row():
|
443 |
+
strength_slider = gr.Slider(
|
444 |
+
minimum=0.3,
|
445 |
+
maximum=0.9,
|
446 |
+
value=0.75,
|
447 |
+
step=0.05,
|
448 |
+
label="Style Intensity (lower values preserve more original details)"
|
449 |
+
)
|
450 |
+
|
451 |
+
# customize Transform style buttom
|
452 |
+
style_button = gr.Button("Transform Style", variant="primary")
|
453 |
+
|
454 |
+
gr.Markdown("""
|
455 |
+
<style>
|
456 |
+
button.primary {
|
457 |
+
background: linear-gradient(90deg, #ff6b6b, #ffa36b, #ffd56b) !important;
|
458 |
+
color: white !important;
|
459 |
+
font-weight: 600 !important;
|
460 |
+
text-shadow: 0 1px 1px rgba(0,0,0,0.2) !important;
|
461 |
+
border: none !important;
|
462 |
+
}
|
463 |
+
|
464 |
+
button.primary:hover {
|
465 |
+
background: linear-gradient(90deg, #ff5b5b, #ff936b, #ffcf6b) !important;
|
466 |
+
box-shadow: 0 4px 8px rgba(0,0,0,0.3) !important;
|
467 |
+
}
|
468 |
+
</style>
|
469 |
+
""")
|
470 |
+
|
471 |
+
# Progress indicator
|
472 |
+
status_indicator = gr.Textbox(
|
473 |
+
label="Status",
|
474 |
+
value="Upload an image and press 'Transform Style' to begin",
|
475 |
+
interactive=False
|
476 |
+
)
|
477 |
+
|
478 |
+
error_output = gr.Textbox(
|
479 |
+
visible=False,
|
480 |
+
label="Error Message"
|
481 |
+
)
|
482 |
+
|
483 |
+
with gr.Column(scale=1):
|
484 |
+
# Output image component
|
485 |
+
output_image = gr.Image(
|
486 |
+
label="Style Transformation Result"
|
487 |
+
)
|
488 |
+
|
489 |
+
# Hidden component to store the download link
|
490 |
+
download_link = gr.HTML(visible=False)
|
491 |
+
|
492 |
+
# HTML component for actual download
|
493 |
+
download_html = gr.HTML(visible=False)
|
494 |
+
|
495 |
+
gr.Markdown("""
|
496 |
+
### Tips for Best Results
|
497 |
+
- Use images with clear dog features and good lighting
|
498 |
+
- For best results, use images where the dog is the main subject
|
499 |
+
- Different styles work better with different dog breeds
|
500 |
+
- Lower the style intensity to preserve more original details
|
501 |
+
""")
|
502 |
+
|
503 |
+
gr.HTML("""
|
504 |
+
<style>
|
505 |
+
.style-box {
|
506 |
+
background: linear-gradient(145deg, #ffffff, #f5f7fa);
|
507 |
+
border-radius: 12px;
|
508 |
+
box-shadow: 0 4px 20px rgba(0,0,0,0.08);
|
509 |
+
padding: 25px 30px;
|
510 |
+
margin: 30px 0;
|
511 |
+
border: 1px solid rgba(0,0,0,0.05);
|
512 |
+
position: relative;
|
513 |
+
}
|
514 |
+
|
515 |
+
.style-box::before {
|
516 |
+
content: '';
|
517 |
+
position: absolute;
|
518 |
+
top: 0;
|
519 |
+
left: 0;
|
520 |
+
width: 6px;
|
521 |
+
height: 100%;
|
522 |
+
background: linear-gradient(to bottom, #ff6b6b, #ffa36b, #ffd56b);
|
523 |
+
border-radius: 6px 0 0 6px;
|
524 |
+
}
|
525 |
+
|
526 |
+
.style-box h2 {
|
527 |
+
color: #333;
|
528 |
+
font-size: 24px;
|
529 |
+
margin-bottom: 20px;
|
530 |
+
padding-bottom: 10px;
|
531 |
+
border-bottom: 2px solid #f0f0f0;
|
532 |
+
}
|
533 |
+
|
534 |
+
.style-name {
|
535 |
+
font-weight: bold;
|
536 |
+
color: #333;
|
537 |
+
}
|
538 |
+
|
539 |
+
.style-desc {
|
540 |
+
margin-bottom: 15px;
|
541 |
+
padding-bottom: 15px;
|
542 |
+
border-bottom: 1px solid #f0f0f0;
|
543 |
+
}
|
544 |
+
|
545 |
+
.style-desc:last-child {
|
546 |
+
margin-bottom: 0;
|
547 |
+
padding-bottom: 0;
|
548 |
+
border-bottom: none;
|
549 |
+
}
|
550 |
+
</style>
|
551 |
+
|
552 |
+
<div class="style-box">
|
553 |
+
<h2>Style Effect Descriptions</h2>
|
554 |
+
<p>Each style transforms your dog photo in a unique way:</p>
|
555 |
+
|
556 |
+
<div class="style-desc">
|
557 |
+
<p><span class="style-name">Japanese Anime Style:</span> Vibrant artwork with fluid animation qualities, expressive features, and dramatic lighting effects. Features soft color gradients, detailed line work, and emotional depth.</p>
|
558 |
+
</div>
|
559 |
+
|
560 |
+
<div class="style-desc">
|
561 |
+
<p><span class="style-name">Classic Cartoon:</span> Traditional animation style with bold outlines, solid color fills, and playful character design. Displays exaggerated expressions, simplified forms, and dynamic poses.</p>
|
562 |
+
</div>
|
563 |
+
|
564 |
+
<div class="style-desc">
|
565 |
+
<p><span class="style-name">Oil Painting:</span> Classical art technique with visible textured brushstrokes and layered color application. Shows rich depth, dramatic lighting contrast, and sophisticated color harmony.</p>
|
566 |
+
</div>
|
567 |
+
|
568 |
+
<div class="style-desc">
|
569 |
+
<p><span class="style-name">Watercolor:</span> Delicate painting style with flowing color blends and translucent layers. Features soft edges, color bleeding effects, and visible paper texture elements.</p>
|
570 |
+
</div>
|
571 |
+
|
572 |
+
<div class="style-desc">
|
573 |
+
<p><span class="style-name">Cyberpunk:</span> High-tech futuristic aesthetic with advanced technological elements and neon accents. Incorporates holographic interfaces, digital effects, and urban dystopian elements.</p>
|
574 |
+
</div>
|
575 |
+
</div>
|
576 |
+
""")
|
577 |
+
|
578 |
+
# Setup event triggers
|
579 |
+
def update_progress(value, desc):
|
580 |
+
"""Update progress bar and description"""
|
581 |
+
return gr.update(value=value), gr.update(value=desc)
|
582 |
+
|
583 |
+
def process_style_transfer(image, style, strength):
|
584 |
+
"""Process style transfer and prepare download options"""
|
585 |
+
if image is None:
|
586 |
+
return (
|
587 |
+
None,
|
588 |
+
gr.update(visible=True, value="Please upload a dog image first!"),
|
589 |
+
gr.update(visible=False),
|
590 |
+
gr.update(visible=False),
|
591 |
+
gr.update(visible=False),
|
592 |
+
gr.update(value="Upload an image and press 'Transform Style' to begin")
|
593 |
+
)
|
594 |
+
|
595 |
+
# Display processing status
|
596 |
+
status_message = "Processing your image... This may take a moment."
|
597 |
+
|
598 |
+
# Perform style transfer
|
599 |
+
result, error = dog_style_transfer.transform_style(
|
600 |
+
image,
|
601 |
+
style,
|
602 |
+
strength
|
603 |
+
)
|
604 |
+
|
605 |
+
if error:
|
606 |
+
return (
|
607 |
+
None,
|
608 |
+
gr.update(visible=True, value=error),
|
609 |
+
gr.update(visible=False),
|
610 |
+
gr.update(visible=False),
|
611 |
+
gr.update(visible=False),
|
612 |
+
gr.update(value="Error occurred. Please try again.")
|
613 |
+
)
|
614 |
+
|
615 |
+
# Generate download link for the image
|
616 |
+
if result is not None:
|
617 |
+
pil_image = Image.fromarray(result)
|
618 |
+
download_data = dog_style_transfer.get_image_download_link(pil_image)
|
619 |
+
download_html_content = f"""
|
620 |
+
<style>
|
621 |
+
.download-btn {{
|
622 |
+
display: inline-block;
|
623 |
+
background: linear-gradient(90deg, #3498db, #2ecc71);
|
624 |
+
color: white !important; /* 確保文字為白色 */
|
625 |
+
text-shadow: 0 1px 2px rgba(0,0,0,0.3) !important; /* 增強文字陰影使白色更突出 */
|
626 |
+
font-weight: 700 !important; /* 加粗字體 */
|
627 |
+
padding: 12px 24px;
|
628 |
+
text-align: center;
|
629 |
+
text-decoration: none;
|
630 |
+
font-size: 16px;
|
631 |
+
border-radius: 25px;
|
632 |
+
cursor: pointer;
|
633 |
+
transition: all 0.3s ease;
|
634 |
+
border: none;
|
635 |
+
box-shadow: 0 3px 6px rgba(0,0,0,0.16);
|
636 |
+
letter-spacing: 0.5px; /* 字母間距,提高可讀性 */
|
637 |
+
}}
|
638 |
+
|
639 |
+
.download-btn:hover {{
|
640 |
+
transform: translateY(-2px);
|
641 |
+
box-shadow: 0 5px 10px rgba(0,0,0,0.25);
|
642 |
+
background: linear-gradient(90deg, #2980b9, #27ae60);
|
643 |
+
}}
|
644 |
+
|
645 |
+
.download-btn:active {{
|
646 |
+
transform: translateY(0);
|
647 |
+
box-shadow: 0 2px 4px rgba(0,0,0,0.1);
|
648 |
+
}}
|
649 |
+
</style>
|
650 |
+
|
651 |
+
<a href="{download_data}" download="dog_{style.replace(' ', '_')}.png" class="download-btn">
|
652 |
+
Download Transformed Image
|
653 |
+
</a>
|
654 |
+
"""
|
655 |
+
|
656 |
+
# Store download data in a hidden element
|
657 |
+
# We'll make this invisible to avoid showing base64 encoded data
|
658 |
+
hidden_download_data = download_data
|
659 |
+
|
660 |
+
return (
|
661 |
+
result,
|
662 |
+
gr.update(visible=False),
|
663 |
+
gr.update(visible=False, value=hidden_download_data), # Keep the data hidden
|
664 |
+
gr.update(visible=True, value=download_html_content), # Show the HTML button
|
665 |
+
gr.update(value="Transform Completed! You can download the image")
|
666 |
+
)
|
667 |
+
else:
|
668 |
+
# Handle the case where result is None but no error was returned
|
669 |
+
return (
|
670 |
+
None,
|
671 |
+
gr.update(visible=True, value="Style transfer failed with no specific error. Please try again."),
|
672 |
+
gr.update(visible=False),
|
673 |
+
gr.update(visible=False),
|
674 |
+
gr.update(visible=False),
|
675 |
+
gr.update(value="Something went wrong. Please try again.")
|
676 |
+
)
|
677 |
+
|
678 |
+
# Update style description and model info
|
679 |
+
def update_style_info(style):
|
680 |
+
return dog_style_transfer.get_style_description(style), dog_style_transfer.get_model_info(style)
|
681 |
+
|
682 |
+
style_button.click(
|
683 |
+
fn=process_style_transfer,
|
684 |
+
inputs=[input_image, style_dropdown, strength_slider],
|
685 |
+
outputs=[
|
686 |
+
output_image,
|
687 |
+
error_output,
|
688 |
+
download_link,
|
689 |
+
download_html,
|
690 |
+
]
|
691 |
+
)
|
692 |
+
|
693 |
+
style_dropdown.change(
|
694 |
+
fn=update_style_info,
|
695 |
+
inputs=[style_dropdown],
|
696 |
+
outputs=[style_description, model_info]
|
697 |
+
)
|
698 |
+
|
699 |
+
|
700 |
+
# Add example images
|
701 |
+
example_dogs = [
|
702 |
+
["Border_Collie.jpg", "Japanese Anime Style"],
|
703 |
+
["Golden_Retriever.jpeg", "Classic Cartoon"],
|
704 |
+
["Saint_Bernard.jpeg", "Oil Painting"],
|
705 |
+
["Samoyed.jpeg", "Watercolor"],
|
706 |
+
["French_Bulldog.jpeg", "Cyberpunk"]
|
707 |
+
]
|
708 |
+
|
709 |
+
# Check if Examples feature is available
|
710 |
+
try:
|
711 |
+
gr.Examples(
|
712 |
+
examples=example_dogs,
|
713 |
+
inputs=[input_image, style_dropdown]
|
714 |
+
)
|
715 |
+
except Exception as e:
|
716 |
+
print(f"Note: Examples feature not available in your Gradio version: {e}")
|
717 |
+
|
718 |
+
gr.HTML("""
|
719 |
+
<style>
|
720 |
+
.attribution-box {
|
721 |
+
font-size: 0.85em;
|
722 |
+
color: #666;
|
723 |
+
margin-top: 20px;
|
724 |
+
padding: 18px;
|
725 |
+
border-radius: 8px;
|
726 |
+
background-color: #f8f9fa;
|
727 |
+
border: 1px solid #e9ecef;
|
728 |
+
font-style: italic;
|
729 |
+
}
|
730 |
+
|
731 |
+
.attribution-box h4 {
|
732 |
+
margin-top: 0;
|
733 |
+
color: #495057;
|
734 |
+
font-style: normal;
|
735 |
+
font-weight: 600;
|
736 |
+
margin-bottom: 12px;
|
737 |
+
}
|
738 |
+
|
739 |
+
.attribution-box p {
|
740 |
+
margin: 8px 0;
|
741 |
+
line-height: 1.5;
|
742 |
+
}
|
743 |
+
</style>
|
744 |
+
|
745 |
+
<div class="attribution-box">
|
746 |
+
<h4>Attribution</h4>
|
747 |
+
<p>This application uses pre-trained diffusion models from Hugging Face for image style transfer. All models are used according to their respective open source licenses for educational and non-commercial purposes.</p>
|
748 |
+
<p>Powered by the open source Diffusers library from Hugging Face.</p>
|
749 |
+
</div>
|
750 |
+
""")
|
751 |
+
|
752 |
+
return input_image, style_dropdown, style_button, output_image
|