creek

The AI Image Editor of 2030

models/style.py (14382B)


  1# -*- coding: utf-8 -*-
  2"""
  3Style
  4Original file: https://colab.research.google.com/drive/1YFqRnpIrbZiN2fwOVcr3FUMdHS3WZdwf
  5"""
  6
  7!pip install git+https://github.com/openai/CLIP.git --q
  8!pip install torch torchvision ftfy regex tqdm --q
  9
 10import torch
 11import clip
 12from PIL import Image
 13import numpy as np
 14from google.colab import files
 15import io
 16
 17device = "cuda" if torch.cuda.is_available() else "cpu"
 18print(f"Using device: {device}")
 19
 20model, preprocess = clip.load("ViT-B/32", device=device)
 21
 22STYLE_CLASSES = [
 23    # Pop / Comic / Print
 24    "pop-art",
 25    "halftone-print",
 26    "comic-book",
 27    "graphic-novel",
 28    "newspaper-comic-strip",
 29    # Cartoon / Animation
 30    "cartoon-style",
 31    "minimal-flat-cartoon",
 32    "rubber-hose",
 33    "retro-cartoon",
 34    "anime",
 35    "chibi",
 36    "manga-inking",
 37    "ghibli-style",
 38    "superhero-comic",
 39    "disney-style",
 40    # Line-Based
 41    "line-art",
 42    "continuous-line",
 43    "blind-contour",
 44    "cross-hatching",
 45    "stippling",
 46    "ink-drawing",
 47    "technical-blueprint",
 48    "engraving",
 49    "monoline-icons",
 50    # Traditional Media
 51    "watercolor",
 52    "gouache",
 53    "oil-painting",
 54    "acrylic-paint",
 55    "pastel-chalk",
 56    "oil-pastel",
 57    "marker-render",
 58    "crayon",
 59    "ink-wash",
 60    "colored-pencil",
 61    "mixed-media",
 62    # Realism / Naturalism
 63    "realism",
 64    "naturalism",
 65    "photorealism",
 66    "hyperrealism",
 67    "portrait-realism",
 68    "botanical-illustration",
 69    "wildlife-realism",
 70    "scientific-illustration",
 71    "documentary-illustration",
 72    # Abstract / Modernist
 73    "abstract-expressionism",
 74    "minimal-abstract",
 75    "geometric-abstraction",
 76    "color-field",
 77    "cubism",
 78    "futurism",
 79    "bauhaus",
 80    "constructivism",
 81    "surreal",
 82    "magical-realism",
 83    "op-art",
 84    "escher-impossible",
 85    # Collage / Cut
 86    "paper-cutout",
 87    "analog-collage",
 88    "digital-collage",
 89    "photomontage",
 90    # Historical / Cultural
 91    "ukiyo-e",
 92    "indian-miniature",
 93    "persian-miniature",
 94    "art-nouveau",
 95    "art-deco",
 96    "medieval-illustration",
 97    "byzantine-icon",
 98    # Retro / Vintage
 99    "vintage",
100    "mid-century-modern",
101    "psychedelic-60s",
102    "groovy-70s",
103    "vaporwave",
104    "synthwave",
105    "retro-futurism",
106    "risograph",
107    "screenprint",
108    # Street / Subculture
109    "graffiti",
110    "street-art",
111    "lowbrow-pop",
112    "sticker-bomb",
113    "zine-grunge",
114    # Digital / 3D / Tech
115    "pixel-art",
116    "isometric-pixel",
117    "voxel-art",
118    "low-poly-3d",
119    "hyperreal-3d",
120    "clay-render",
121    "cel-shaded-3d",
122    "holographic-gradient",
123    "glitch-art",
124    "chromatic-aberration",
125    "datamosh",
126    "vhs-analog",
127]
128
129
130STYLE_DESCRIPTIONS = {
131    # Pop / Comic / Print
132    "pop-art": "pop art, comic book style, bold colors, halftone dots, thick outlines",
133    "halftone-print": "halftone dot shading, CMYK print texture, vintage comic printing",
134    "comic-book": "dynamic panels, bold inking, action lines, superhero comic style",
135    "graphic-novel": "moody shadows, dramatic black ink, narrative panels, serious tone",
136    "newspaper-comic-strip": "simple line work, flat shading, classic newspaper strip style",
137    # Cartoon / Animation
138    "cartoon-style": "bold outlines, flat colors, exaggerated characters and expressions",
139    "minimal-flat-cartoon": "simple vector shapes, flat color fills, minimal details",
140    "rubber-hose": "1930s animation, noodle limbs, vintage rubber hose characters",
141    "retro-cartoon": "old-school grain, hand-drawn frames, muted vintage palette",
142    "anime": "clean linework, cel-shading, large expressive eyes, Japanese animation",
143    "chibi": "super-deformed characters, big heads, tiny bodies, cute aesthetic",
144    "manga-inking": "black and white screentones, clean ink lines, manga style",
145    "ghibli-style": "soft painterly shading, warm palettes, whimsical nature scenes",
146    "superhero-comic": "muscular characters, dramatic shading, intense action poses",
147    "disney-style": "smooth polished shading, expressive eyes, classic animation aesthetic",
148    # Line-Based
149    "line-art": "clean linear outlines, no shading, monochrome strokes",
150    "continuous-line": "one unbroken flowing line forming the subject",
151    "blind-contour": "distorted outlines, drawn without looking, loose sketch style",
152    "cross-hatching": "intersecting lines for shading, dense ink texture",
153    "stippling": "dot-based shading, gradual tonal gradients from dots",
154    "ink-drawing": "strong black ink strokes, high contrast, brush or pen feel",
155    "technical-blueprint": "precise technical diagrams, blue background, white lines",
156    "engraving": "fine parallel carved lines, antique print look",
157    "monoline-icons": "uniform line weight, minimal shapes, iconographic design",
158    # Traditional Media
159    "watercolor": "transparent washes, soft gradients, bleeding pigment, paper texture",
160    "gouache": "opaque matte colors, flat coverage, velvety paint texture",
161    "oil-painting": "rich color blends, heavy brushstrokes, textured canvas feel",
162    "acrylic-paint": "opaque fast-drying strokes, strong colors, painterly texture",
163    "pastel-chalk": "soft dusty texture, delicate blending, chalky edges",
164    "oil-pastel": "waxy thick strokes, vibrant smears, layered pigments",
165    "marker-render": "streaky marker fills, smooth gradients, alcohol ink texture",
166    "crayon": "childlike wax texture, rough strokes, uneven coloring",
167    "ink-wash": "monochrome fluid washes, sumi-e brush style, soft gradients",
168    "colored-pencil": "fine pencil texture, layered shading, visible grain",
169    "mixed-media": "combined materials, layered textures, hybrid collage effects",
170    # Realism / Naturalism
171    "realism": "accurate representation, natural lighting, lifelike detail",
172    "naturalism": "organic colors, soft shading, true-to-life textures",
173    "photorealism": "photo-like accuracy, crisp details, polished lighting",
174    "hyperrealism": "ultra-detailed textures, exaggerated clarity, surreal realism",
175    "portrait-realism": "true facial anatomy, realistic skin tones, fine shading",
176    "botanical-illustration": "high-detail plants, scientific precision, fine lines",
177    "wildlife-realism": "accurate animal anatomy, natural fur/feather rendering",
178    "scientific-illustration": "precise labeled diagrams, technical accuracy",
179    "documentary-illustration": "real-life scenes, muted tones, reportage style",
180    # Abstract / Modernist
181    "abstract-expressionism": "gestural strokes, energetic splashes, emotional color",
182    "minimal-abstract": "simple geometric forms, large negative space, clarity",
183    "geometric-abstraction": "strict shapes, patterns, bold geometry",
184    "color-field": "large flat color areas, subtle tonal transitions",
185    "cubism": "fragmented shapes, multiple viewpoints, angular abstraction",
186    "futurism": "dynamic angles, motion streaks, mechanical aesthetic",
187    "bauhaus": "primary colors, clean geometry, functional composition",
188    "constructivism": "bold diagonals, industrial forms, propaganda flavor",
189    "surreal": "dreamlike imagery, strange juxtapositions, impossible forms",
190    "magical-realism": "ordinary scenes with subtle fantastical elements",
191    "op-art": "optical illusions, vibrating lines, high-contrast patterns",
192    "escher-impossible": "paradoxical architecture, infinite loops, tessellations",
193    # Collage / Cut
194    "paper-cutout": "flat colored shapes, crisp edges, layered paper look",
195    "analog-collage": "torn paper edges, magazine scraps, glue texture",
196    "digital-collage": "layered photography fragments, digital assembly",
197    "photomontage": "multiple photos combined into surreal compositions",
198    # Historical / Cultural
199    "ukiyo-e": "japanese woodblock style, flat colors, bold outlines, wave patterns",
200    "indian-miniature": "fine detailing, ornate patterns, traditional storytelling",
201    "persian-miniature": "vibrant colors, intricate decoration, ornate scenes",
202    "art-nouveau": "flowing curves, floral motifs, elegant decorative lines",
203    "art-deco": "geometric symmetry, metallic accents, luxury aesthetic",
204    "medieval-illustration": "flat perspective, gold leaf accents, illuminated style",
205    "byzantine-icon": "gold backgrounds, frontal figures, religious icon style",
206    # Retro / Vintage
207    "vintage": "aged tones, film grain, retro print texture, nostalgic aesthetic",
208    "mid-century-modern": "retro shapes, muted tones, playful geometric forms",
209    "psychedelic-60s": "vivid gradients, swirling patterns, trippy colors",
210    "groovy-70s": "warm earthy colors, wavy lines, bohemian retro vibe",
211    "vaporwave": "neon pinks and blues, retro grids, greek busts, glitchy nostalgia",
212    "synthwave": "neon purples, chrome shine, retrofuturistic sunset grids",
213    "retro-futurism": "futuristic style imagined in the past, old sci-fi aesthetic",
214    "risograph": "duotone neon inks, grainy overlays, misaligned color layers",
215    "screenprint": "flat inks, bold shapes, layered silkscreen texture",
216    # Street / Subculture
217    "graffiti": "spray paint texture, drips, tags, bold lettering",
218    "street-art": "urban wall murals, stencil work, mixed street media",
219    "lowbrow-pop": "quirky surrealism, outsider cartoon style, punk influence",
220    "sticker-bomb": "dense overlapping stickers, colorful chaotic layers",
221    "zine-grunge": "photocopy texture, rough cut edges, DIY punk style",
222    # Digital / 3D / Tech
223    "pixel-art": "pixelated blocks, low resolution, retro game style",
224    "isometric-pixel": "45-degree pixel perspective, 3D pixel illusion",
225    "voxel-art": "3D cubes, blocky shapes, volumetric pixel aesthetic",
226    "low-poly-3d": "simple polygon shapes, faceted surfaces, flat shading",
227    "hyperreal-3d": "lifelike 3D rendering, photoreal materials and lighting",
228    "clay-render": "soft clay surface, sculpted feel, neutral lighting",
229    "cel-shaded-3d": "toon outlines, flat shading, stylized animation look",
230    "holographic-gradient": "iridescent metallic rainbow gradients, glossy shine",
231    "glitch-art": "broken pixels, digital errors, RGB misalignment",
232    "chromatic-aberration": "red-blue fringing, lens distortion edges",
233    "datamosh": "compression smears, frame artifacts, melting visuals",
234    "vhs-analog": "scanlines, tape noise, retro analog glitch",
235}
236
237
238PROMPT_TEMPLATES = [
239    "a photo in the style of {}",
240    "artistic rendering of {}",
241    "a {} style image",
242    "a painting in the style of {}",
243    "a high quality example of {} art",
244    "visuals depicting {}",
245    "artwork created using {}",
246    "a poster with {}",
247    "a photograph looking like {}",
248    "an illustration inspired by {}",
249    "a digital artwork following the rules of {}",
250    "a creative piece showcasing {}",
251    "a visual composition designed in {} style",
252    "an artistic composition influenced by {}",
253    "a scene presented with {} aesthetics",
254    "an image heavily influenced by {} techniques",
255    "a graphical representation in {} aesthetics",
256    "a stylized depiction matching {}",
257    "the visual language of {} applied to an artwork",
258    "a concept image using {} characteristics",
259    "a stylized portrait in the manner of {}",
260    "a landscape interpreted through {} style",
261    "a figure drawing executed in {} technique",
262    "a creative poster influenced by {}",
263    "a concept sketch using {} visual rules",
264    "an editorial illustration drawn in {} style",
265    "a rendered scene with {} visual identity",
266    "an artistic study done with {} aesthetics",
267    "an experiment in {} visual style",
268    "a fully rendered image designed around {} style",
269]
270
271
272class StyleClassifier:
273    def __init__(self):
274        self.centroids = None
275        self.class_names = STYLE_CLASSES
276        self.build_text_features()
277
278    def build_text_features(self):
279        """
280        Pre-computes the embedding centroids for all styles.
281        """
282        print("Building text embeddings for styles...")
283        class_centroids = []
284
285        with torch.no_grad():
286            for style in self.class_names:
287                desc = STYLE_DESCRIPTIONS[style]
288
289                prompts = [tpl.format(desc) for tpl in PROMPT_TEMPLATES]
290
291                tokens = clip.tokenize(prompts).to(device)
292
293                text_features = model.encode_text(tokens)
294
295                text_features = text_features / text_features.norm(dim=-1, keepdim=True)
296
297                centroid = text_features.mean(dim=0)
298
299                centroid = centroid / centroid.norm()
300
301                class_centroids.append(centroid)
302
303        self.centroids = torch.stack(class_centroids).to(device)
304        print("Text embeddings built successfully.")
305
306    def predict(self, image_path, threshold=0.22):
307        try:
308            image = Image.open(image_path).convert("RGB")
309        except Exception as e:
310            print(f"Error loading image: {e}")
311            return "error", []
312
313        img_tensor = preprocess(image).unsqueeze(0).to(device)
314
315        with torch.no_grad():
316            img_feat = model.encode_image(img_tensor)
317            img_feat = img_feat / img_feat.norm(dim=-1, keepdim=True)
318
319        similarity = 100.0 * img_feat @ self.centroids.T
320
321        probs = similarity.softmax(dim=-1).cpu().numpy()[0]
322        raw_scores = similarity.cpu().numpy()[0] / 100.0
323
324        results = []
325        for i, style in enumerate(self.class_names):
326            results.append(
327                {
328                    "style": style,
329                    "score": float(raw_scores[i]),
330                    "confidence_pct": float(probs[i] * 100),
331                }
332            )
333
334        results.sort(key=lambda x: x["score"], reverse=True)
335
336        top_result = results[0]
337
338        if top_result["score"] < threshold:
339            prediction = "undefined/mixed"
340        else:
341            prediction = top_result["style"]
342
343        return prediction, results
344
345
346#  Runner
347
348classifier = StyleClassifier()
349
350print("\nPlease upload an image to analyze:")
351uploaded = files.upload()
352
353for filename in uploaded.keys():
354    print(f"\nAnalyzing: {filename}...")
355    prediction, details = classifier.predict(filename)
356
357    print(f"TOP PREDICTION: {prediction.upper()}")
358    print("-" * 30)
359    print("Full Breakdown:")
360    for res in details[:5]:
361        print(
362            f"{res['style']:<20} | Score: {res['score']:.4f} | Conf: {res['confidence_pct']:.2f}%"
363        )