Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| from torch import nn | |
| from transformers import CLIPVisionModel, CLIPImageProcessor | |
| from PIL import Image, ImageEnhance, ImageOps | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # ----------------------------- | |
| # Aesthetic Scoring Model | |
| # ----------------------------- | |
| class AestheticScorer(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.vision_model = CLIPVisionModel.from_pretrained("openai/clip-vit-base-patch16") | |
| self.mlp = nn.Sequential( | |
| nn.Linear(self.vision_model.config.hidden_size, 512), | |
| nn.ReLU(), | |
| nn.Linear(512, 1) | |
| ) | |
| def forward(self, pixel_values): | |
| outputs = self.vision_model(pixel_values=pixel_values) | |
| pooled = outputs.pooler_output | |
| return self.mlp(pooled) | |
| processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch16") | |
| ae_model = AestheticScorer().to(device) | |
| ae_model.eval() | |
| # ----------------------------- | |
| # Cinematic Enhancement | |
| # ----------------------------- | |
| def apply_cinematic(image: Image.Image) -> Image.Image: | |
| img = image.convert("RGB") | |
| # slight letterboxing for cinematic aspect ratio feel | |
| w, h = img.size | |
| bar_height = int(h * 0.07) # 7% black bars | |
| new_img = Image.new("RGB", (w, h), (0, 0, 0)) | |
| new_img.paste(img, (0, 0)) | |
| # adjust colors for cinematic mood | |
| enhancer = ImageEnhance.Color(img) | |
| img = enhancer.enhance(1.2) # richer colors | |
| enhancer = ImageEnhance.Contrast(img) | |
| img = enhancer.enhance(1.15) # more contrast | |
| enhancer = ImageEnhance.Brightness(img) | |
| img = enhancer.enhance(0.95) # slightly darker | |
| # add subtle warm tint (teal & orange style) | |
| r, g, b = img.split() | |
| r = r.point(lambda i: i * 1.05) | |
| b = b.point(lambda i: i * 0.95) | |
| img = Image.merge("RGB", (r, g, b)) | |
| # add black bars | |
| new_img.paste(img, (0, 0)) | |
| draw = ImageOps.expand(img, border=(0, bar_height, 0, bar_height), fill="black") | |
| return draw | |
| # ----------------------------- | |
| # Scoring | |
| # ----------------------------- | |
| def aesthetic_score(image: Image.Image) -> float: | |
| inputs = processor(images=image, return_tensors="pt").to(device) | |
| with torch.no_grad(): | |
| score = ae_model(inputs['pixel_values']).cpu().item() | |
| return round(float(score), 2) | |
| # ----------------------------- | |
| # Pipeline | |
| # ----------------------------- | |
| def process(image): | |
| # score original | |
| original_score = aesthetic_score(image) | |
| # cinematic version | |
| cinematic_img = apply_cinematic(image) | |
| # score cinematic | |
| cinematic_score = aesthetic_score(cinematic_img) | |
| # ranking suggestion | |
| suggestions = sorted( | |
| [("Original", original_score), ("Cinematic", cinematic_score)], | |
| key=lambda x: x[1], | |
| reverse=True, | |
| ) | |
| result_text = "📊 Aesthetic Ranking:\n" | |
| for rank, (name, score) in enumerate(suggestions, 1): | |
| result_text += f"{rank}. {name} — Score: {score}\n" | |
| return cinematic_img, result_text | |
| # ----------------------------- | |
| # Gradio UI | |
| # ----------------------------- | |
| with gr.Blocks() as demo: | |
| gr.Markdown("## 🎬 Cinematic Photo Enhancer") | |
| with gr.Row(): | |
| inp = gr.Image(type="pil", label="Upload your photo") | |
| out_img = gr.Image(type="pil", label="Cinematic Output") | |
| out_txt = gr.Textbox(label="Suggestions & Scores") | |
| run_btn = gr.Button("Enhance Photo") | |
| run_btn.click(process, inputs=inp, outputs=[out_img, out_txt]) | |
| demo.launch() | |