Update app.py

app.py

@@ -10,123 +10,129 @@ from gtts import gTTS
 import uuid
 import tempfile
 
-# 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)
+yolo_model = YOLO("yolov8n.pt").to(device)
+fashion_model = YOLO("best.pt").to(device)  # Your trained model
 
 style_prompts = {
-    '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"]
+    "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 = [
-    "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"
+    "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': ["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}."]
+    "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 = [prompt for cat in style_prompts.values() for prompt in cat] + clothing_prompts
+all_prompts = [p for cat in style_prompts.values() for p in cat] + clothing_prompts
 
 def get_top_clothing(probs, n=3):
     clothing_probs = probs[len(all_prompts) - len(clothing_prompts):]
     top_indices = np.argsort(clothing_probs)[-n:]
     return [clothing_prompts[i] for i in reversed(top_indices)]
 
-def analyze_outfit(img: Image.Image):
+def analyze_outfit(img):
     results = yolo_model(img)
     boxes = results[0].boxes.xyxy.cpu().numpy()
     classes = results[0].boxes.cls.cpu().numpy()
     confidences = results[0].boxes.conf.cpu().numpy()
 
     person_indices = np.where(classes == 0)[0]
-    cropped_img = img
+    cropped = img
     if len(person_indices) > 0:
-        max_conf_idx = np.argmax(confidences[person_indices])
-        x1, y1, x2, y2 = map(int, boxes[person_indices][max_conf_idx])
-        cropped_img = img.crop((x1, y1, x2, y2))
+        idx = np.argmax(confidences[person_indices])
+        x1, y1, x2, y2 = map(int, boxes[person_indices][idx])
+        cropped = img.crop((x1, y1, x2, y2))
 
-    image_tensor = clip_preprocess(cropped_img).unsqueeze(0).to(device)
+    # Run fashion model to get top class label
+    fashion_results = fashion_model(cropped, verbose=False)
+    top_item_idx = fashion_results[0].probs.top1
+    top_item_name = fashion_results[0].names[int(top_item_idx)]
+
+    # CLIP classification
+    image_tensor = clip_preprocess(cropped).unsqueeze(0).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_len = len(style_prompts['drippy'])
-    mid_len = len(style_prompts['mid'])
-    not_len = len(style_prompts['not_drippy'])
+    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])
+    not_score = np.mean(probs[drip_len + mid_len:])
 
     if drip_score > mid_score and drip_score > not_score:
-        category_key = 'drippy'
+        cat = "drippy"
         final_score = drip_score
     elif mid_score > not_score:
-        category_key = 'mid'
+        cat = "mid"
         final_score = mid_score
     else:
-        category_key = 'not_drippy'
+        cat = "not_drippy"
         final_score = not_score
 
-    category_label = CATEGORY_LABEL_MAP[category_key]
-    clothing_items = get_top_clothing(probs)
-    clothing_item = clothing_items[0]
-    response = random.choice(response_templates[category_key]).format(item=clothing_item)
-
+    label = CATEGORY_LABEL_MAP[cat]
+    response = random.choice(response_templates[cat]).format(item=top_item_name)
     tts_path = os.path.join(tempfile.gettempdir(), f"drip_{uuid.uuid4().hex}.mp3")
     gTTS(response, lang="en").save(tts_path)
 
-    final_score_str = f"{final_score:.2f}"
-    category_html = f"""
+    html = f"""
         <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>
+            <h2>Your fit is <span style='color:#1f04ff'>{label}</span>!</h2>
+            <p>Drip Score: <strong>{final_score:.2f}</strong></p>
         </div>
     """
+    return html, tts_path, response
 
-    return category_html, tts_path, response
-
-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:
-
+# Gradio UI
+with gr.Blocks(css=".container { max-width: 600px; margin: auto; padding: 2rem; }") as demo:
     with gr.Group(elem_classes=["container"]):
-        gr.Markdown("### Choose an input method:")
-
-        with gr.Row():
-            webcam_input = gr.Image(label="📸 Take a Photo", image_mode="RGB", type="pil", show_label=False, show_download_button=False, sources=["webcam"])
-            upload_input = gr.Image(label="🖼️ Upload Photo", image_mode="RGB", type="pil", show_label=False, show_download_button=False, sources=["upload"])
-
-        analyze_button = gr.Button("🔥 Analyze My Fit")
-
-        # Placeholder outputs
-        category_html = gr.HTML()
-        audio_output = gr.Audio(autoplay=True, label="")
-        response_box = gr.Textbox(label="Response", lines=2, interactive=False)
-
-        # Merge the two image inputs into one workflow
-        def merged_input(image1, image2):
-            return image1 if image1 is not None else image2
-
-        merged = gr.State()
-
-        webcam_input.change(fn=merged_input, inputs=[webcam_input, upload_input], outputs=merged)
-        upload_input.change(fn=merged_input, inputs=[webcam_input, upload_input], outputs=merged)
-
-        analyze_button.click(fn=analyze_outfit, inputs=[merged], outputs=[category_html, audio_output, response_box])
-
-if __name__ == '__main__':
+        method = gr.Radio(["Upload Image", "Use Webcam"], label="Choose how to submit your fit", value="Upload Image")
+        upload = gr.Image(type="pil", label="Upload Image", visible=True)
+        webcam = gr.Image(type="pil", label="Take Photo", visible=False, sources=["webcam"])
+        analyze = gr.Button("🔥 Analyze My Fit")
+        html = gr.HTML()
+        audio = gr.Audio(autoplay=True, label="")
+        textbox = gr.Textbox(label="Response", interactive=False, lines=2)
+
+        def toggle_inputs(method):
+            return gr.update(visible=method == "Upload Image"), gr.update(visible=method == "Use Webcam")
+
+        method.change(toggle_inputs, inputs=method, outputs=[upload, webcam])
+        analyze.click(fn=analyze_outfit, inputs=upload, outputs=[html, audio, textbox])
+        analyze.click(fn=analyze_outfit, inputs=webcam, outputs=[html, audio, textbox])
+
+if __name__ == "__main__":
     demo.launch()
-