Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from model import Trainer | |
| import torch | |
| import cv2 | |
| import tempfile | |
| import numpy as np | |
| def predict_beauty_score(img): | |
| trainer = Trainer() | |
| # Save the numpy array as an image temporarily | |
| with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp: | |
| # Convert RGB to BGR for cv2 | |
| img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) | |
| cv2.imwrite(tmp.name, img_bgr) | |
| # Use the temporary file path | |
| image_tensor = trainer.image_to_tensor(tmp.name) | |
| prediction = trainer.predict(image_tensor) | |
| score = prediction.item() * 100 | |
| # Decide which GIF to show based on the score | |
| if score < 20: | |
| gif_url = "https://i.pinimg.com/originals/9f/79/2a/9f792aed5881d409425de1a4361bc06b.gif" | |
| elif score < 40: | |
| gif_url = "https://i.pinimg.com/originals/ba/f5/c8/baf5c89c099b34decb7f4507b5144366.gif" | |
| elif score < 55: | |
| gif_url = "https://i.pinimg.com/originals/87/b6/dc/87b6dcfeec6f38a3836b1caf1d8fceab.gif" | |
| elif score < 60: | |
| gif_url = "https://i.pinimg.com/originals/3a/90/b8/3a90b87a337b79b9c8b7a3d9bf7250d7.gif" | |
| else: | |
| gif_url = "https://i.pinimg.com/originals/f6/02/01/f6020120d9e99f7b106c557cdc1edb1f.gif" | |
| # Create formatted HTML outpu | |
| html_output = f""" | |
| <div style='text-align: center; padding: 20px;'> | |
| <h2 style='color: #FFFFFF; margin-bottom: 10px;'>Score</h2> | |
| <div style='font-size: 48px; font-weight: bold; color: #1a73e8;'> | |
| {score:.3f} | |
| </div> | |
| <br/> | |
| <div style='display: flex; justify-content: center;'> | |
| <img src="{gif_url}" alt="GIF" width="200" height="200" /> | |
| </div> | |
| <p>Enjoy the fun! :)</p> | |
| </div> | |
| """ | |
| return html_output | |
| # Create Gradio interface | |
| demo = gr.Interface( | |
| fn=predict_beauty_score, | |
| inputs=gr.Image(), # Simple image input | |
| outputs=gr.HTML(), # Using HTML output for custom formatting | |
| title="Image Beauty Score Predictor", | |
| description=""" | |
| <ul> | |
| <li>Upload an image to get its beauty score prediction</li> | |
| <li>Please remember that this is just for fun and is not intended to downplay anybody and I, as the model creator believe everyone is beautiful the way it is</li> | |
| <li>Respect each other and most importantly have fun :)</li> | |
| </ul> | |
| """, | |
| examples=[ | |
| ["6082308423334085331.jpg"], | |
| ["a00e1c819a87fa56bb1e6058d9814bae.jpg"], | |
| ["Screenshot_2024-12-29_211520.png"] | |
| ], | |
| cache_examples=True | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |