Update app.py

app.py

@@ -1,67 +1,113 @@
+import gradio as gr
 import torch
 import torch.nn as nn
-import torchvision.transforms as transforms
-from PIL import Image
+from transformers import CLIPProcessor, CLIPModel, ViTFeatureExtractor, ViTForImageClassification
+from PIL import Image, ImageEnhance, ImageFilter, ImageOps
 import numpy as np
-import gradio as gr
+import cv2
 
-# -----------------------
-# Zero-DCE Model
-# -----------------------
-class ZeroDCE(nn.Module):
-    def __init__(self, num_layers=7):
-        super(ZeroDCE, self).__init__()
-        filters = 32
-        layers = []
-        layers += [nn.Conv2d(3, filters, 3, 1, 1), nn.ReLU(inplace=True)]
-        for _ in range(num_layers - 2):
-            layers += [nn.Conv2d(filters, filters, 3, 1, 1), nn.ReLU(inplace=True)]
-        layers += [nn.Conv2d(filters, 3, 3, 1, 1), nn.Tanh()]
-        self.net = nn.Sequential(*layers)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    def forward(self, x):
-        return x + self.net(x)
+# ----------------- Aesthetic Scoring (CLIP + regression head) -----------------
+class AestheticPredictor(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.clip = CLIPModel.from_pretrained("openai/clip-vit-base-patch16")
+        self.mlp = nn.Sequential(
+            nn.Linear(self.clip.config.projection_dim, 512),
+            nn.ReLU(),
+            nn.Linear(512, 1)
+        )
+    def forward(self, pixel_values):
+        outputs = self.clip(pixel_values=pixel_values).pooler_output
+        return self.mlp(outputs)
 
-# -----------------------
-# Setup
-# -----------------------
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-zero_dce = ZeroDCE().to(device)
+processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")
+ae_model = AestheticPredictor().to(device)
+ae_model.eval()
 
-# -----------------------
-# Image Enhancement
-# -----------------------
-def enhance_image(image, model):
-    # Always force RGB
-    image = image.convert("RGB")
+def aesthetic_score(img):
+    inputs = processor(images=img, return_tensors="pt").to(device)
+    with torch.no_grad():
+        s = ae_model(inputs['pixel_values'])
+    return float(s.item())
 
-    transform = transforms.Compose([
-        transforms.ToTensor(),
-    ])
-    img_tensor = transform(image).unsqueeze(0).to(device)
+# ----------------- Classifier -----------------
+cls_processor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
+cls_model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224").to(device)
+cls_model.eval()
 
+def classify(img):
+    inputs = cls_processor(images=img, return_tensors="pt").to(device)
     with torch.no_grad():
-        enhanced = model(img_tensor)
+        logits = cls_model(**inputs).logits
+    label = cls_model.config.id2label[logits.argmax(-1).item()]
+    return label.lower()
 
-    # Convert back to PIL
-    enhanced = enhanced.squeeze(0).cpu().permute(1, 2, 0).numpy()
-    enhanced = (enhanced * 255).clip(0, 255).astype(np.uint8)
+# ----------------- Cinematic Filters (hardcoded) -----------------
+CINEMATIC_STYLES = {
+    "portrait": [
+        dict(name="Warm Glow", gamma=0.9, contrast=1.3, saturation=1.2, color_tint=(240, 200, 180), vignette=0.4),
+        dict(name="Moody Teal-Orange", gamma=1.0, contrast=1.4, saturation=1.0, color_tint=(30, 210, 180), vignette=0.5),
+    ],
+    "landscape": [
+        dict(name="Epic Teal-Orange", gamma=1.0, contrast=1.5, saturation=1.3, color_tint=(25, 200, 160), vignette=0.3),
+        dict(name="Soft Film Look", gamma=0.95, contrast=1.2, saturation=1.1, color_tint=(220, 200, 180), vignette=0.2),
+    ],
+    "default": [
+        dict(name="Classic Cinema", gamma=1.0, contrast=1.4, saturation=1.1, color_tint=(220, 180, 160), vignette=0.3),
+    ],
+}
 
-    return Image.fromarray(enhanced)
+def apply_cinematic(img, style):
+    img = img.convert("RGB")
+    # gamma
+    arr = np.array(img).astype(np.float32) / 255.0
+    arr = arr ** style["gamma"]
+    # contrast & saturation
+    img = Image.fromarray((arr*255).astype(np.uint8))
+    img = ImageEnhance.Contrast(img).enhance(style["contrast"])
+    img = ImageEnhance.Color(img).enhance(style["saturation"])
+    # color tint via overlay
+    tint = Image.new("RGB", img.size, style["color_tint"])
+    img = Image.blend(img, tint, alpha=0.1)
+    # vignette
+    if style["vignette"]>0:
+        w, h = img.size
+        mask = Image.new("L", (w,h), 255)
+        draw = ImageDraw.Draw(mask)
+        draw.ellipse([(-w*style["vignette"], -h*style["vignette"]),
+                      (w*(1+style["vignette"]), h*(1+style["vignette"]))],
+                     fill=0)
+        img = Image.composite(img, ImageOps.colorize(mask, (0,0,0), (0,0,0)), mask)
+    return img
 
-# -----------------------
-# Gradio Interface
-# -----------------------
-def process_image(input_img):
-    return enhance_image(input_img, zero_dce)
+# ----------------- Pipeline -----------------
+def process(img):
+    label = classify(img)
+    key = "portrait" if "person" in label else ("landscape" if "landscape" in label else "default")
+    styles = CINEMATIC_STYLES.get(key, CINEMATIC_STYLES["default"])
+    candidates = []
+    for style in styles:
+        out = apply_cinematic(img, style)
+        score = aesthetic_score(out)
+        candidates.append((score, style["name"], out))
+    candidates.sort(reverse=True, key=lambda x: x[0])
+    # Output winner, plus gallery
+    gallery = [(o, f"{name}: {s:.2f}") for s,name,o in candidates]
+    winner_score, winner_name, winner_img = candidates[0]
+    return winner_img, f"Barely cinematic vibe: **{winner_name}** (score {winner_score:.2f})", gallery
 
+# ----------------- UI -----------------
 with gr.Blocks() as demo:
-    gr.Markdown("## 🌙 Low Light Image Enhancement (Zero-DCE)")
+    gr.Markdown("##  Cinematic AI — Instant Film-Grade Style")
     with gr.Row():
-        inp = gr.Image(type="pil", label="Upload Dark Image")
-        out = gr.Image(type="pil", label="Enhanced Image")
-    run_btn = gr.Button("Enhance")
-    run_btn.click(fn=process_image, inputs=inp, outputs=out)
+        inp = gr.Image(type="pil", label="Upload Your Photo")
+        out = gr.Image(type="pil", label="Cinematic Result")
+    btn = gr.Button("Make Cinematic")
+    info = gr.Markdown()
+    gallery = gr.Gallery(label="All Styles (ranked)", columns=2)
+    btn.click(process, [inp], [out, info, gallery])
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch()