Mavthunder's picture
Update app.py
0f173ed verified
raw
history blame
3.49 kB
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()