Update app.py

app.py

@@ -10,37 +10,52 @@ from gtts import gTTS
 import uuid
 import tempfile
 
-# Device and model loading
+# Setup device and models
 device = "cuda" if torch.cuda.is_available() else "cpu"
 clip_model, clip_preprocess = clip.load("ViT-B/32", device=device)
 yolo_model = YOLO('yolov8n.pt').to(device)
 fashion_model = YOLO('best.pt').to(device)
 
-# Style prompts and templates
+# Style prompts and categories
 style_prompts = {
-    'drippy': [...],  # truncated for brevity
-    'mid': [...],
-    'not_drippy': [...]
+    'drippy': [
+        "avant-garde streetwear", "high-fashion designer outfit", "trendsetting urban attire",
+        "luxury sneakers and chic accessories", "cutting-edge, bold style"
+    ],
+    'mid': [
+        "casual everyday outfit", "modern minimalistic attire", "comfortable yet stylish look",
+        "simple, relaxed streetwear", "balanced, practical fashion"
+    ],
+    'not_drippy': [
+        "disheveled outfit", "poorly coordinated fashion", "unfashionable, outdated attire",
+        "tacky, mismatched ensemble", "sloppy, uninspired look"
+    ]
 }
 
-clothing_prompts = [...]
+clothing_prompts = [
+    "t-shirt", "dress shirt", "blouse", "hoodie", "jacket", "sweater", "coat", "dress", "skirt",
+    "pants", "jeans", "trousers", "shorts", "sneakers", "boots", "heels", "sandals", "cap", "hat",
+    "scarf", "gloves", "bag", "accessory", "tank-top", "haircut"
+]
 
 response_templates = {
-    'drippy': [...],
-    'mid': [...],
-    'not_drippy': [...]
+    'drippy': [
+        "You're Drippy, bruh – fire {item}!", "{item} goes crazy, on god!",
+        "Certified drippy with that {item}."
+    ],
+    'mid': [
+        "Drop the {item} and you might get a text back.", "It's alright, but I'd upgrade the {item}.",
+        "Mid fit alert. That {item} is holding you back."
+    ],
+    'not_drippy': [
+        "Bro thought that {item} was tuff!", "Oh hell nah! Burn that {item}!",
+        "Crimes against fashion, especially that {item}! Also… maybe get a haircut.",
+        "Never walk out the house again with that {item}."
+    ]
 }
 
-CATEGORY_LABEL_MAP = {
-    "drippy": "drippy",
-    "mid": "mid",
-    "not_drippy": "trash"
-}
-
-all_prompts = []
-for cat_prompts in style_prompts.values():
-    all_prompts.extend(cat_prompts)
-all_prompts.extend(clothing_prompts)
+CATEGORY_LABEL_MAP = {"drippy": "drippy", "mid": "mid", "not_drippy": "trash"}
+all_prompts = [prompt for cat in style_prompts.values() for prompt in cat] + clothing_prompts
 
 def get_top_clothing(probs, n=3):
     clothing_probs = probs[len(all_prompts) - len(clothing_prompts):]
@@ -61,20 +76,28 @@ def analyze_outfit(img: Image.Image):
         cropped_img = img.crop((x1, y1, x2, y2))
 
     image_tensor = clip_preprocess(cropped_img).unsqueeze(0).to(device)
-    text_tokens = clip.tokenize(all_prompts).to(device)
+    text_tokens = clip.tokenize([str(p) for p in all_prompts]).to(device)
     with torch.no_grad():
         logits, _ = clip_model(image_tensor, text_tokens)
         probs = logits.softmax(dim=-1).cpu().numpy()[0]
 
-    drip_score = np.mean(probs[:len(style_prompts['drippy'])])
-    mid_score = np.mean(probs[len(style_prompts['drippy']):len(style_prompts['drippy'])+len(style_prompts['mid'])])
-    not_score = np.mean(probs[len(style_prompts['drippy'])+len(style_prompts['mid']):])
-
-    category_key = max(['drippy', 'mid', 'not_drippy'], key=lambda k: np.mean(
-        probs[:len(style_prompts[k])] if k == 'drippy' else
-        probs[len(style_prompts['drippy']):len(style_prompts['drippy'])+len(style_prompts['mid'])] if k == 'mid' else
-        probs[len(style_prompts['drippy'])+len(style_prompts['mid']):]
-    ))
+    drip_len = len(style_prompts['drippy'])
+    mid_len = len(style_prompts['mid'])
+    not_len = len(style_prompts['not_drippy'])
+    
+    drip_score = np.mean(probs[:drip_len])
+    mid_score = np.mean(probs[drip_len:drip_len + mid_len])
+    not_score = np.mean(probs[drip_len + mid_len:drip_len + mid_len + not_len])
+
+    if drip_score > mid_score and drip_score > not_score:
+        category_key = 'drippy'
+        final_score = drip_score
+    elif mid_score > not_score:
+        category_key = 'mid'
+        final_score = mid_score
+    else:
+        category_key = 'not_drippy'
+        final_score = not_score
 
     category_label = CATEGORY_LABEL_MAP[category_key]
     clothing_items = get_top_clothing(probs)
@@ -83,60 +106,38 @@ def analyze_outfit(img: Image.Image):
 
     tts_path = os.path.join(tempfile.gettempdir(), f"drip_{uuid.uuid4().hex}.mp3")
     gTTS(response, lang="en").save(tts_path)
-    final_score_str = f"{max(drip_score, mid_score, not_score):.2f}"
 
+    final_score_str = f"{final_score:.2f}"
     category_html = f"""
-        <div style='text-align: center;'>
-            <h2 style='color: #1f04ff;'>Your fit is <b>{category_label.upper()}</b></h2>
-            <p style='font-size: 18px;'>Drip Score: <strong>{final_score_str}</strong></p>
+        <div style='padding:1rem; text-align:center;'>
+            <h2 style='margin-bottom:0.5rem;'>Your fit is <span style='color:#1f04ff'>{category_label}</span>!</h2>
+            <p style='font-size:1.1rem;'>Drip Score: <strong>{final_score_str}</strong></p>
         </div>
     """
 
     return category_html, tts_path, response
 
-# Gradio interface with cleaner styling
-custom_css = """
-.container {
-    max-width: 700px;
-    margin: 0 auto;
-    font-family: 'Arial', sans-serif;
-}
-button {
-    background-color: #1f04ff;
-    color: white;
-    border-radius: 6px;
-    padding: 10px 20px;
-    font-size: 16px;
-}
-button:hover {
-    background-color: #3c2fff;
-}
-.gradio-container {
-    background: #f9f9f9;
-    border-radius: 10px;
-    padding: 20px;
-    box-shadow: 0 4px 10px rgba(0,0,0,0.1);
-}
-"""
-
-with gr.Blocks(css=custom_css) as demo:
-    with gr.Column(elem_classes=["container"]):
-        gr.Markdown("""
-            # 👟 DripAI
-            Upload your outfit to get judged by the algorithm.
-            No bias. No mercy. Just drip.
-        """)
-        input_image = gr.Image(type='pil', label="Upload your outfit")
-        analyze_button = gr.Button("Analyze My Fit")
+# Gradio UI layout
+with gr.Blocks(css="""
+.container { max-width: 600px; margin: 0 auto; padding: 2rem; }
+button { background-color: #1f04ff; color: white; border: none; padding: 0.75rem 1.5rem; border-radius: 6px; cursor: pointer; font-weight: bold; }
+button:hover { background-color: #1500cc; }
+#resultbox { border: 1px solid #e3e3e3; border-radius: 10px; padding: 1rem; background: #fafafa; }
+""") as demo:
+
+    with gr.Group(elem_classes=["container"]):
+        input_image = gr.Image(type='pil', sources=['upload', 'webcam', 'clipboard'], label="Upload or Snap Your Fit")
+        analyze_button = gr.Button("🔥 Analyze My Fit")
 
         category_html = gr.HTML()
-        audio_output = gr.Audio(autoplay=True, label="AI Feedback")
-        response_box = gr.Textbox(lines=2, label="Generated Response")
+        audio_output = gr.Audio(autoplay=True, label="")
+        response_box = gr.Textbox(label="Response", lines=2, interactive=False)
 
         analyze_button.click(
             fn=analyze_outfit,
-            inputs=[input_image],
+            inputs=input_image,
             outputs=[category_html, audio_output, response_box],
         )
 
-demo.launch()
+if __name__ == '__main__':
+    demo.launch()