Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import clip | |
| import numpy as np | |
| import random | |
| import os | |
| from PIL import Image | |
| from ultralytics import YOLO # Needed for both person and fashion detection | |
| from gtts import gTTS | |
| import uuid | |
| import time | |
| import tempfile | |
| def analyze_outfit(input_img, yolo_person_model, yolo_fashion_model, clip_model, clip_preprocess, | |
| all_prompts, style_prompts_end_index, FASHION_CLASSES, CATEGORY_LABEL_MAP, | |
| response_templates, YOLO_PERSON_CONF_THRESHOLD, YOLO_FASHION_CONF_THRESHOLD, | |
| YOLO_FASHION_HIGH_CONF_THRESHOLD, DEVICE): | |
| # Handle both file paths and PIL Images | |
| if isinstance(input_img, str): | |
| try: | |
| input_img = Image.open(input_img) | |
| except Exception as e: | |
| return (f"<p style='color: #FF5555;'>Error loading image: {str(e)}</p>", | |
| None, "Image loading error") | |
| # Existing code continues... | |
| if input_img is None: | |
| return ("<p style='color: #FF5555; text-align: center;'>Please upload an image.</p>", | |
| None, "Error: No image provided.") | |
| img = input_img.convert("RGB").copy() | |
| #def analyze_outfit(image): | |
| #if image is None: | |
| #return ("<p style='color: #FF5555; text-align: center;'>Please upload an image.</p>", None, "Error: No image provided.") | |
| #image = image.convert("RGB").copy() | |
| #print(f"[DEBUG] image_path type: {type(image_path)} | value: {image_path}") | |
| # 1) YOLO Person Detection | |
| person_results = yolo_person_model(img, verbose=False, conf=YOLO_PERSON_CONF_THRESHOLD) | |
| boxes = person_results[0].boxes.xyxy.cpu().numpy() | |
| classes = person_results[0].boxes.cls.cpu().numpy() | |
| confidences = person_results[0].boxes.conf.cpu().numpy() | |
| # Filter for persons (class 0 in standard YOLOv8) | |
| person_indices = np.where(classes == 0)[0] | |
| cropped_img = img # Default to full image if no person found | |
| person_detected = False | |
| if len(person_indices) > 0: | |
| # Find the person detection with the highest confidence | |
| max_conf_person_idx = person_indices[np.argmax(confidences[person_indices])] | |
| x1, y1, x2, y2 = map(int, boxes[max_conf_person_idx]) | |
| # Ensure coordinates are valid and within image bounds | |
| x1, y1 = max(0, x1), max(0, y1) | |
| x2, y2 = min(img.width, x2), min(img.height, y2) | |
| if x1 < x2 and y1 < y2: # Check if the box has valid dimensions | |
| cropped_img = img.crop((x1, y1, x2, y2)) | |
| print(f"Person detected and cropped: Box {x1, y1, x2, y2}") | |
| person_detected = True | |
| else: | |
| print("Warning: Invalid person bounding box after clipping. Using full image.") | |
| cropped_img = img | |
| else: | |
| print("No person detected by yolo_person_model. Analyzing full image.") | |
| # 2) YOLO Fashion Model Detection (run on the cropped image if person was found) | |
| detected_fashion_item_name = None | |
| detected_fashion_item_conf = 0.0 | |
| if person_detected or True: # Or always run on the (potentially full) image? Let's always run for now. | |
| try: | |
| fashion_results = yolo_fashion_model(cropped_img, verbose=False, conf=YOLO_FASHION_CONF_THRESHOLD) | |
| fashion_boxes = fashion_results[0].boxes.xyxy.cpu().numpy() | |
| fashion_classes = fashion_results[0].boxes.cls.cpu().numpy().astype(int) | |
| fashion_confidences = fashion_results[0].boxes.conf.cpu().numpy() | |
| if len(fashion_classes) > 0: | |
| # Find the detection with the highest confidence | |
| best_fashion_idx = np.argmax(fashion_confidences) | |
| detected_class_id = fashion_classes[best_fashion_idx] | |
| detected_fashion_item_conf = fashion_confidences[best_fashion_idx] | |
| if detected_class_id in FASHION_CLASSES: | |
| detected_fashion_item_name = FASHION_CLASSES[detected_class_id] | |
| print(f"Fashion model detected: '{detected_fashion_item_name}' " | |
| f"with confidence {detected_fashion_item_conf:.2f}") | |
| else: | |
| print(f"Warning: Detected fashion class ID {detected_class_id} not in FASHION_CLASSES map.") | |
| else: | |
| print("No fashion items detected above threshold by yolo_fashion_model.") | |
| except Exception as e: | |
| print(f"Error during YOLO fashion model analysis: {e}") | |
| # Continue without fashion model input | |
| # 3) CLIP Analysis (always run on the cropped/full image) | |
| clip_detected_item = "look" # Default fallback item name | |
| clip_detected_item_prob = 0.0 | |
| category_key = 'mid' # Default category | |
| final_score_str = "N/A" | |
| try: | |
| image_tensor = clip_preprocess(cropped_img).unsqueeze(0).to(DEVICE) | |
| text_tokens = clip.tokenize(all_prompts).to(DEVICE) | |
| with torch.no_grad(): | |
| logits, _ = clip_model(image_tensor, text_tokens) | |
| all_probs = logits.softmax(dim=-1).cpu().numpy()[0] | |
| # Calculate style scores | |
| drip_len = len(style_prompts['drippy']) | |
| mid_len = len(style_prompts['mid']) | |
| drip_score = np.mean(all_probs[0 : drip_len]) | |
| mid_score = np.mean(all_probs[drip_len : drip_len + mid_len]) | |
| not_score = np.mean(all_probs[drip_len + mid_len : style_prompts_end_index]) | |
| # Determine overall style category AND DEFINE score_label | |
| score_label = "Style Score" # Initialize with a default/fallback | |
| if drip_score > 0.41 and drip_score > mid_score and drip_score > not_score: | |
| category_key = 'drippy' | |
| final_score = drip_score | |
| score_label = "Drip Score" # <<< DEFINE score_label | |
| elif mid_score > not_score: # Check mid_score > not_score explicitly | |
| category_key = 'mid' | |
| final_score = mid_score | |
| score_label = "Mid Score" # <<< DEFINE score_label | |
| else: | |
| category_key = 'not_drippy' | |
| final_score = not_score | |
| score_label = "Trash Score" # <<< DEFINE score_label # Or maybe "Rating Score" | |
| category_label = CATEGORY_LABEL_MAP[category_key] | |
| # final_score_str = f"{final_score:.2f}" # You might not need this raw score string anymore | |
| percentage_score = max(0, final_score * 100) | |
| percentage_score_str = f"{percentage_score:.0f}%" # Formats as integer (e.g., "3%", "15%", "0%") | |
| # Now score_label is defined before being used here | |
| print(f"Style analysis: Category={category_label}, Score = {score_label}={percentage_score_str} (Raw Score: {final_score:.4f})") | |
| # Get top clothing item from CLIP | |
| top_3_clip_items = get_top_clip_clothing(all_probs, n=3) # <<< Ask for top 3 items | |
| if top_3_clip_items: | |
| # Print the top 3 detected items | |
| detected_items_str = ", ".join([f"{item[0]} ({item[1]*100:.1f}%)" for item in top_3_clip_items]) # Show item and probability | |
| print(f"I think I detected: {detected_items_str}") | |
| # Still use the single *most* probable item for response generation logic later | |
| clip_detected_item, clip_detected_item_prob = top_3_clip_items[0] | |
| # Optional: You can keep or remove the print for the single top item below if the top-3 print is sufficient | |
| # print(f"Top clothing item identified by CLIP (for response): '{clip_detected_item}' " | |
| # f"with probability {clip_detected_item_prob:.2f}") | |
| else: | |
| print("I couldn't confidently identify specific clothing items via CLIP.") | |
| clip_detected_item = "piece" # Use a different fallback if CLIP fails | |
| clip_detected_item_prob = 0.0 # Ensure prob is defined | |
| except Exception as e: | |
| print(f"Error during CLIP analysis: {e}") | |
| # Use defaults, maybe return error message? | |
| return ("<p style='color: #FF5555;'>Error during CLIP analysis.</p>", | |
| None, f"Analysis Error: {e}") | |
| # 4) Determine the Final Item to Mention in Response | |
| final_clothing_item = "style" # Ultimate fallback generic term | |
| generic_response_needed = False | |
| if detected_fashion_item_name and detected_fashion_item_conf >= YOLO_FASHION_HIGH_CONF_THRESHOLD: | |
| # Priority 1: High-confidence fashion model detection | |
| final_clothing_item = detected_fashion_item_name | |
| print(f"Using highly confident fashion model item: '{final_clothing_item}'") | |
| elif detected_fashion_item_name and detected_fashion_item_conf >= YOLO_FASHION_CONF_THRESHOLD: | |
| # Priority 2: Medium-confidence fashion model detection (still prefer over CLIP) | |
| final_clothing_item = detected_fashion_item_name | |
| print(f"Using medium confidence fashion model item: '{final_clothing_item}'") | |
| elif clip_detected_item and clip_detected_item_prob > 0.05: # Check if CLIP prob is somewhat reasonable | |
| # Priority 3: CLIP detection (if fashion model didn't provide a strong candidate) | |
| final_clothing_item = clip_detected_item | |
| print(f"Using CLIP detected item: '{final_clothing_item}'") | |
| else: | |
| # Priority 4: Generic response needed (no confident detection from either model) | |
| final_clothing_item = random.choice(["fit", "look", "style", "vibe"]) # Randomize generic term | |
| generic_response_needed = True | |
| print(f"Using generic fallback item: '{final_clothing_item}'") | |
| # 5) Generate Response and TTS | |
| try: | |
| response_pool = response_templates[category_key] | |
| # Choose a random template from the entire response pool | |
| chosen_template = random.choice(response_pool) | |
| # Format the response, substituting the item name if needed | |
| response_text = chosen_template.format(item=final_clothing_item) if '{item}' in chosen_template else chosen_template | |
| tts_path = os.path.join(tempfile.gettempdir(), f"drip_{uuid.uuid4().hex}.mp3") | |
| tts = gTTS(text=response_text, lang='en', tld='com', slow=False) | |
| tts.save(tts_path) | |
| print(f"Generated TTS response: '{response_text}' saved to {tts_path}") | |
| # --- Updated HTML Output --- | |
| category_html = f""" | |
| <div class='results-container'> | |
| <h2 class='result-category'>RATING: {category_label.upper()}</h2> | |
| <p class='result-score'>{score_label}: {percentage_score_str}</p> | |
| </div> | |
| """ | |
| return category_html, tts_path, response_text | |
| except Exception as e: | |
| print(f"Error during response/TTS generation: {e}") | |
| percentage_score = max(0, final_score * 100) | |
| percentage_score_str = f"{percentage_score:.0f}%" | |
| category_html = f""" | |
| <div class='results-container'> | |
| <h2 class='result-category'>Result: {category_label.upper()}</h2> | |
| <p class='result-score'>{score_label}: {percentage_score_str}</p> | |
| <p class='result-error' style='color: #FFAAAA; font-size: 0.9em;'>Error generating audio/full response.</p> | |
| </div> | |
| """ | |
| # Still provide category info, but indicate TTS/response error | |
| return category_html, None, f"Analysis complete ({category_label}), but error generating audio/response." |