creek

The AI Image Editor of 2030
commit e649603ddb86fc5bce757347b31da2a3ee99f28c
parent b04800d0a5ef4cf944514446fdbae47869f09506
Author: maydayv7 <maydayv7@gmail.com>
Date:   Thu,  4 Dec 2025 22:35:26 +0530

Add `models` and update docs

Diffstat:

D.gitattributes | 2--
M.gitignore | 13+++++++------
AINSTALL.md | 66++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
MREADME.md | 226+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++----------------
Dassets/clip_image_encoder.onnx | 3---
Dassets/dinov2.onnx | 3---
Dassets/dinov2.onnx.data | 3---
Dassets/fannet.onnx | 3---
Dflask/.gitattributes | 3---
Aflask/.gitignore | 6++++++
Dflask/BiRefNet/.gitignore | 3---
Dflask/BiRefNet/LICENSE | 3---
Dflask/BiRefNet/README.md | 3---
Dflask/BiRefNet/birefnet_fp16.pt | 3---
Dflask/BiRefNet/config.py | 3---
Dflask/BiRefNet/dataset.py | 3---
Dflask/BiRefNet/eval_existingOnes.py | 3---
Dflask/BiRefNet/evaluation/metrics.py | 3---
Dflask/BiRefNet/gen_best_ep.py | 3---
Dflask/BiRefNet/image_proc.py | 3---
Dflask/BiRefNet/inference.py | 3---
Dflask/BiRefNet/loss.py | 3---
Dflask/BiRefNet/make_a_copy.sh | 3---
Dflask/BiRefNet/models/backbones/build_backbone.py | 3---
Dflask/BiRefNet/models/backbones/dino_v3.py | 3---
Dflask/BiRefNet/models/backbones/pvt_v2.py | 3---
Dflask/BiRefNet/models/backbones/swin_v1.py | 3---
Dflask/BiRefNet/models/birefnet.py | 3---
Dflask/BiRefNet/models/modules/aspp.py | 3---
Dflask/BiRefNet/models/modules/decoder_blocks.py | 3---
Dflask/BiRefNet/models/modules/deform_conv.py | 3---
Dflask/BiRefNet/models/modules/lateral_blocks.py | 3---
Dflask/BiRefNet/models/modules/utils.py | 3---
Dflask/BiRefNet/requirements.txt | 3---
Dflask/BiRefNet/rm_cache.sh | 3---
Dflask/BiRefNet/sub.sh | 3---
Dflask/BiRefNet/test.sh | 3---
Dflask/BiRefNet/train.py | 3---
Dflask/BiRefNet/train.sh | 3---
Dflask/BiRefNet/train_test.sh | 3---
Dflask/BiRefNet/tutorials/BiRefNet_inference.ipynb | 3---
Dflask/BiRefNet/tutorials/BiRefNet_inference_video.ipynb | 3---
Dflask/BiRefNet/tutorials/BiRefNet_pth2onnx.ipynb | 3---
Dflask/BiRefNet/utils.py | 3---
Dflask/Florence-2-4bit-Quantized/added_tokens.json | 3---
Dflask/Florence-2-4bit-Quantized/config.json | 3---
Dflask/Florence-2-4bit-Quantized/configuration_florence2.py | 3---
Dflask/Florence-2-4bit-Quantized/generation_config.json | 3---
Dflask/Florence-2-4bit-Quantized/merges.txt | 3---
Dflask/Florence-2-4bit-Quantized/model.safetensors | 3---
Dflask/Florence-2-4bit-Quantized/modeling_florence2.py | 3---
Dflask/Florence-2-4bit-Quantized/preprocessor_config.json | 3---
Dflask/Florence-2-4bit-Quantized/processing_florence2.py | 3---
Dflask/Florence-2-4bit-Quantized/processor_config.json | 3---
Dflask/Florence-2-4bit-Quantized/special_tokens_map.json | 3---
Dflask/Florence-2-4bit-Quantized/tokenizer.json | 3---
Dflask/Florence-2-4bit-Quantized/tokenizer_config.json | 3---
Dflask/Florence-2-4bit-Quantized/vocab.json | 3---
Mflask/README.md | 6++++++
Dflask/local_inpainting_model/feature_extractor/preprocessor_config.json | 3---
Dflask/local_inpainting_model/model_index.json | 3---
Dflask/local_inpainting_model/safety_checker/config.json | 3---
Dflask/local_inpainting_model/safety_checker/model.safetensors | 3---
Dflask/local_inpainting_model/scheduler/scheduler_config.json | 3---
Dflask/local_inpainting_model/text_encoder/config.json | 3---
Dflask/local_inpainting_model/text_encoder/model.safetensors | 3---
Dflask/local_inpainting_model/tokenizer/merges.txt | 3---
Dflask/local_inpainting_model/tokenizer/special_tokens_map.json | 3---
Dflask/local_inpainting_model/tokenizer/tokenizer_config.json | 3---
Dflask/local_inpainting_model/tokenizer/vocab.json | 3---
Dflask/local_inpainting_model/unet/config.json | 3---
Dflask/local_inpainting_model/unet/diffusion_pytorch_model.safetensors | 3---
Dflask/local_inpainting_model/vae/config.json | 3---
Dflask/local_inpainting_model/vae/diffusion_pytorch_model.safetensors | 3---
Amodels/color.py | 113+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
Amodels/composition.py | 442+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
Amodels/emotions.py | 165+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
Amodels/era_cult_ref.py | 140+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
Amodels/inpainting_api&local.py | 275+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
Amodels/lighting.py | 319+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
Amodels/sketch2img_local_server.py | 62++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
Amodels/style.py | 383+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
Amodels/stylesheet_gen&prompt_gen.py | 672+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
Amodels/subject.py | 181+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
Amodels/testing_data_gen_template.py | 153+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
Amodels/text&mat/dinov2_vits14_fp16.pt | 0
Amodels/text&mat/texture&mat.py | 441+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
Amodels/text&mat/texture_centroids_s14.pt | 0
Amodels/typo.py | 78++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
89 files changed, 3690 insertions(+), 257 deletions(-)

diff --git a/.gitattributes b/.gitattributes @@ -1,2 +0,0 @@ -*.onnx filter=lfs diff=lfs merge=lfs -text -*.onnx.data filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore @@ -1,3 +1,5 @@ +.env + # Miscellaneous ~* *.class @@ -40,14 +42,13 @@ app.*.map.json /android/app/profile /android/app/release +# VS Code +.vscode + # Python venv python_env/ __pycache__ -# VS Code -.vscode/ - -# .env Files -.env -flask/.env +# Models +*onnx* diff --git a/INSTALL.md b/INSTALL.md @@ -0,0 +1,66 @@ +# Binary + +The pre-built APK file can be downloaded from [here](https://drive.google.com/file/d/1JnLvXiOlkQzK2VtBG40WsiTq1eKGJF5F/view) to experience the app on your Android phone without any hassle (uses a deployed [Modal.com](https://modal.com) backend). +To build the app from source, read the rest of this document. + +# Prerequisites + +- Flutter SDK **>=3.7** +- Python **3.8** +- Gradle **8.9.1** +- Java/JDK **17** & Kotlin +- Android SDK & NDK + - Platform Level **>=34** + - Build Tools **>=34.0.0** + - NDK Version **27.0.12077973** + - CMake **3.22.1** + +Ensure that all these programs with the correct versions are installed on your system before using this repository + +# Setup + +### 1. Download Models + +Download all the models from [here](https://drive.google.com/drive/folders/1b2cEsZao5miX7WxYmp160XkKd_lfBmPN) and place them inside the `assets` directory + +### 2. Setup Python + +The Android app compiles Python code using [Chaquopy](https://chaquo.com/chaquopy/), which requires the following setup, mandatorily using Python Version 3.8: + +#### A) Install Python 3.8 + +Download from [python.org](https://www.python.org/downloads/release/python-3810/) + +#### B) Set up Virtual Environment + +```bash +python -m venv python_env +``` + +### 3. Create .env + +Create the `.env` file in the root directory containing the following: + +```dotenv +URL_ASSET=<Background Removal Backend Endpoint> +URL_DESCRIBE=<Florence-v2 Description Backend Endpoint> +URL_GENERATE=<Stable Diffusion Image Generation Backend Endpoint> +URL_INPAINTING=<Stable Diffusion Inpainting Backend Endpoint> +URL_INPAINTING_API=<Fal.ai Inpainting Backend Endpoint> +URL_SKETCH_API=<Sketch to Image Backend Endpoint> +SHARED_SECRET_KEY=<Base64 Security Key> +``` + +### 4. Build the application APK + +```bash +flutter clean +flutter pub get +flutter build apk +``` + +If you have `adb` properly installed, you can use `flutter run` + +# Backend + +Read [flask/README.md](flask/README.md) to run the server diff --git a/README.md b/README.md @@ -1,66 +1,202 @@ -# CreekUI +# CreekUI — Re-Imagining Photoshop (The AI Editor of 2030) -A Flutter project with AI-based image analysis +CreekUI is a lightweight, intent-driven mobile image editor designed to shift the creative workflow from manual pixel manipulation to semantic, AI-assisted design. It separates creativity into structured phases that prioritize **semantic understanding**, **energy efficiency**, and **selective use of heavy compute**. -## Prerequisites +This repository contains the full analysis pipeline, style synthesis logic, and generation workflow that together simulate how a future (circa 2030) mobile NPU might function in creative applications. -- `git` & `git-lfs` -- Flutter SDK **>=3.7** -- Python **3.8** -- Gradle **8.9.1** -- Java/JDK **17** & Kotlin -- Android SDK & NDK - - Platform Level **>=34** - - Build Tools **>=34.0.0** - - NDK Version **27.0.12077973** - - CMake **3.22.1** +> [!NOTE] +> For installation and setup instructions, refer to [INSTALL.md](./INSTALL.md). +> For model details, check out the [models](./models) directory. -Ensure that all these programs with the correct versions are installed on your system before using this repository +--- -## Setup +<details> +<summary><b>DATASETS</b></summary> -### 1. Clone Repository +- DIS5K-TR Dataset +- FANnet Dataset +- EasyOCR Dataset +- LAION-5B Dataset (_Stable Diffusion v1.5_) +- LAION-AESTHETICS V2 4.5 (_Stable Diffusion v1.5 inpainting_) -Clone using `git lfs clone` to ensure all models are correctly pulled +</details> -### 2. Setup Python +--- -The Android app compiles Python code using [Chaquopy](https://chaquo.com/chaquopy/), which requires the following setup, mandatorily using Python Version 3.8: +> [!IMPORTANT] +> This architecture makes use of models provided by [Fal.ai](https://fal.ai/) API. -#### A) Install Python 3.8 +--- -Download from [python.org](https://www.python.org/downloads/release/python-3810/) +# 1. Project Overview -#### B) Set up Virtual Environment +CreekUI decomposes the creative workflow into three major stages: -```bash -python -m venv python_env -``` +1. **Image Analysis → Semantic representations** +2. **Stylesheet Generation → Parametric design language** +3. **Asset Synthesis → Inpainting / Sketch-to-Image / Final outputs** -### 3. Create .env +Only when the user finalizes intent does the system touch heavy pixel-generation models; everything before that operates on lightweight metadata. -Create the `.env` file in the root directory containing the following: +The runtime spans: -```dotenv -URL_ASSET=<Background Removal Backend Endpoint> -URL_DESCRIBE=<Florence-v2 Description Backend Endpoint> -URL_GENERATE=<Stable Diffusion Image Generation Backend Endpoint> -URL_INPAINTING=<Stable Diffusion Inpainting Backend Endpoint> -URL_INPAINTING_API=<Fal.ai Inpainting Backend Endpoint> -URL_SKETCH_API=<Sketch to Image Backend Endpoint> -SHARED_SECRET_KEY=<Base64 Security Key> -``` +- **Flutter UI** (front-end, canvas, orchestration) +- **Kotlin Native Layer** (method channel, concurrency, device coordination) +- **Chaquopy Python Runtime** (OpenCV, clustering, ONNX inference) +- **Flask ML Middleware** (simulated NPU backend for heavy models) -### 4. Build the APK +--- -```bash -flutter clean -flutter pub get -flutter build apk -``` +# 2. High-Level Architecture -If you have `adb` properly installed, you can use `flutter run` +## A. Presentation Layer (Flutter) -## Backend +The UI provides a responsive canvas and orchestrates analysis tasks. Key features: -Read [flask/README.md](flask/README.md) to run the server +- **`AnalysisQueueManager`** + Background queue that handles image jobs without blocking UI. + Batches note (cropped-region) jobs to reduce redundant image decodes. + +- **Canvas Tools** + + - Magic Draw (sketch-based generation) + - Layer history and reversible state + - Stylesheet asset import + +- **Prompt Builder Hooks** + Integrates semantic cues, stylesheets, and user notes into a unified request. + +## B. Native Intelligence Layer (Kotlin + Chaquopy) + +A hybrid Kotlin/Python pipeline handles local (non-generative) ML work: + +- Starts the Chaquopy Python runtime with controlled synchronization (`CountDownLatch`). +- Executes OpenCV operations (saliency, geometry, edge detection). +- Runs ONNX models for segmentation / embeddings where feasible. +- Uses Kotlin coroutines to keep analysis asynchronous and avoid blocking the UI thread. + +This layer ensures low-latency feedback and offline capability for most non-generative tasks. + +## C. Simulated NPU Backend (Flask Middleware) + +Since current mobile hardware can’t run multi-billion-parameter generative models at interactive speed, a local Flask server stands in as a simulation of a “2030 mobile NPU.” + +- Runs heavy models: Stable Diffusion, Flux-based models, Florence-2. +- Models are quantized to approximate mobile constraints. +- Provides GPU-accelerated endpoints while mimicking the API boundaries of on-device inference. + +**Security:** +All payloads are wrapped with **AES-256-GCM** using a fresh nonce per request to emulate the privacy guarantees of true on-device processing. + +--- + +# 3. Detailed Pipelines + +## A. Image Analysis Pipeline (Moodboard) + +Each image passes through several independent analyzers executed in parallel (thread-per-tag approach): + +1. **Composition Analysis** + Lightweight OpenCV functions compute saliency, centroids, edge maps, and evaluate rules such as thirds, symmetry, and negative space. + +2. **Segmentation (BiRefNet)** + Provides high-quality masks at low overhead after 16-bit quantization. + +3. **Texture Embeddings (DINOv2)** + Extracts patch-level embeddings to classify textures against prototype centroids. + +4. **Style / Era / Emotion (CLIP)** + Performs zero-shot similarity against curated descriptors. + +5. **Typography Recognition (EasyOCR + FANnet)** + OCR tokens → FANnet embedding → nearest-neighbor font family lookup. + +Each analyzer contributes structured JSON fragments. At the end of processing, the system merges these fragments into a comprehensive semantic summary for the image. + +## B. Stylesheet Synthesis (Parametric Design Layer) + +Multiple image summaries (from those uploaded to the moodboard) are fused into a project-level style profile. + +Steps include: + +- **Rank–decay weighted aggregation** to emphasize consistent patterns across the moodboard. +- **Palette extraction** via clustering. +- **Texture & layout consensus** using pooled embeddings and geometric heuristics. +- **Font aggregation** through repeated FANnet similarity scores. + +The output is a compact `stylesheet.json` containing the project’s: + +- Composition +- Typography +- Background/Texture +- Colours +- Material Look +- Lighting +- Style +- Era +- Emotion + +This stylesheet avoids repeated inference during user exploration: everything from brush behavior to prompt generation uses this parametric design layer. + +## C. Generation Pipeline (Sketch-to-Image / Inpainting) + +When the user provides intent (sketches, strokes, notes), the system constructs a unified prompt: + +1. **Caption** from a quantized Florence-2 model +2. **Style cues** extracted from the stylesheet +3. **User directives** or overrides (object placement, prompt, etc.) + +Depending on context: + +- **Sketch-to-Image** uses lightweight or high-quality external backends (Nano Banana, Flux). +- **Inpainting** performs localized edits using Stable Diffusion with a mask derived from canvas stroke differences, or FLUX LoRa Fill. + +## Local generation is prioritized for privacy-sensitive edits; API generation is used for high fidelity or speed when allowed. + +# 4. Compute Profile + +Approximate characteristics of each model in its quantized / optimized configuration: + +| Algorithm / Model | Inference Time | Number of Parameters | VRAM Usage | RAM Usage | +| ----------------------------------------- | -------------- | -------------------- | ---------- | --------- | +| Composition Detection Algo. | 30 ms | - | - | - | +| BiRefNet | 502 ms | 221 million | 3.5 GB | 2 GB | +| Typography Detection | 3 seconds | 238k | 50 MB | 70 MB | +| Background / Texture / Material Detection | 23 ms | 22 million | 350 MB | 2.46 GB | +| Colour & Colour Palette Detection | 10 ms | - | - | - | +| Lighting, Style, Era & Emotion Detection | 9 ms | 150 million | 240 MB | 280 MB | +| Florence 2 Base | 1 second | 230 million | 750 MB | 380 MB | +| Nano Banana API | 6 seconds | - | - | - | +| Flux Dev API | 5 seconds | - | - | - | +| Stable Diffusion v1.5 | 6 seconds | 983 million | 3.6 GB | 2.5 GB | +| Flux LoRa Fill API | 7 seconds | - | - | - | +| Stable Diffusion v1.5 Inpainting | 7 seconds | 860 million | 4 GB | 2.7 GB | + +--- + +# 5. Deployment & Runtime Decisions + +- **Quantization** provided the best balance of fidelity and memory footprint. +- **Distillation** was tested but discarded due to poor stability and hallucination issues. +- **Lazy model loading** and reuse of ONNX sessions reduce repeated overhead. +- **Thread-capped parallelism** ensures segmentation and embedding steps don’t overload memory. + +### On-Device Layer + +- Offload only lightweight analyzers (OpenCV, CLIP, DINOv2 depending on budget). +- Use asynchronous queues for all analysis tasks. +- Cache embeddings per image to avoid re-computation. + +### Backend Layer + +- Use mixed-precision and selective quantization for generative workloads. +- Keep session warm to avoid cold-start latency spikes. +- Choose between local-backend vs external APIs based on privacy and performance flags. + +--- + +# 6. Future Roadmap + +- **Global Style Profile**: Inferring a user's long-term visual taste across multiple projects to offer personalized defaults. +- **Composition-Aware Camera**: A real-time AR overlay that guides users to take photos that align with their current moodboard's composition scores. +- **Multi-Device Synchronization**: A seamless ecosystem where compute tasks are dynamically distributed across a user's phone, tablet, and desktop based on available power and thermal headroom. diff --git a/assets/clip_image_encoder.onnx b/assets/clip_image_encoder.onnx @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:65bf51cdaba76c0cb3f77a922da4a43ebb13d3a5674e9eeb922456d210ad2a4d -size 351826328 diff --git a/assets/dinov2.onnx b/assets/dinov2.onnx @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d191d5e0dba661920919126e790cf097b753c0dcfedfd9332e78fccbc6376821 -size 753403 diff --git a/assets/dinov2.onnx.data b/assets/dinov2.onnx.data @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:4fed072682368c03e6c0ed0c970af25411ceeff324359320c65fce068e6dd089 -size 88224768 diff --git a/assets/fannet.onnx b/assets/fannet.onnx @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:62dfd008648d6eb54daa1aabe25d9784bda65e185cd06c4873b6329d7821c5c3 -size 8403630 diff --git a/flask/.gitattributes b/flask/.gitattributes @@ -1,3 +0,0 @@ -local_inpainting_model/** filter=lfs diff=lfs merge=lfs -text -Florence-2-4bit-Quantized/** filter=lfs diff=lfs merge=lfs -text -BiRefNet/** filter=lfs diff=lfs merge=lfs -text diff --git a/flask/.gitignore b/flask/.gitignore @@ -0,0 +1,6 @@ +.env + +# Models +BiRefNet +Florence-2-4bit-Quantized +local_inpainting_model diff --git a/flask/BiRefNet/.gitignore b/flask/BiRefNet/.gitignore @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7939c19c49e8257b5311fea54e45cb79aa970fa0909629151f6fb3abaeb49b9d -size 2026 diff --git a/flask/BiRefNet/LICENSE b/flask/BiRefNet/LICENSE @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:92a7089e0915fc32bc40067560b398f1e6a7a5958abd7d04eda393629a5acefb -size 1066 diff --git a/flask/BiRefNet/README.md b/flask/BiRefNet/README.md @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b142fd5846651860b86ff640849309f549ad238080501c0220a6cc9da637609d -size 40449 diff --git a/flask/BiRefNet/birefnet_fp16.pt b/flask/BiRefNet/birefnet_fp16.pt @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a0e0be5d62afd3d1413a863999a06e67cf5084cbe634ecbd81c91448775a0b69 -size 444669483 diff --git a/flask/BiRefNet/config.py b/flask/BiRefNet/config.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:6580aacdd661a16575a700561f91eb341d6d2df1440591056f462a1a1b1044aa -size 11365 diff --git a/flask/BiRefNet/dataset.py b/flask/BiRefNet/dataset.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e684883e7ba07c95f0b98f8ebe7d9cb33749583ba043b0fba70c59a29677862d -size 10316 diff --git a/flask/BiRefNet/eval_existingOnes.py b/flask/BiRefNet/eval_existingOnes.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:774393da63c66aa2b45e7d2d0830ae256142a951b92b424f2a4bcdacde55dc11 -size 4354 diff --git a/flask/BiRefNet/evaluation/metrics.py b/flask/BiRefNet/evaluation/metrics.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:494d096de5e62c09cc19ffed903ada93b094605c432df9278a6eb40875ba7926 -size 40411 diff --git a/flask/BiRefNet/gen_best_ep.py b/flask/BiRefNet/gen_best_ep.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:5af9a4efe44afecd2a0cc8c2a9ebef8fd9545366ec0a4a22d265d203e86eeb85 -size 4197 diff --git a/flask/BiRefNet/image_proc.py b/flask/BiRefNet/image_proc.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:86fe45da1e49e10801b70a2e336f4222c5f6c0482fe613227cc5893c7553f045 -size 6836 diff --git a/flask/BiRefNet/inference.py b/flask/BiRefNet/inference.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:451d6ee07edf7d05f550fd3bb879696b9af8860cd6cbe7a26d58ea8a920593ec -size 4741 diff --git a/flask/BiRefNet/loss.py b/flask/BiRefNet/loss.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:db830cfa50b0c10acf6f24fb18ae11f1b3a6b0839df8883a7734c18ed93199df -size 9269 diff --git a/flask/BiRefNet/make_a_copy.sh b/flask/BiRefNet/make_a_copy.sh @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:5a620fcb6cef8d2e5932dd58fbea7d39a595894d2a3f62dfdca49dace3963557 -size 428 diff --git a/flask/BiRefNet/models/backbones/build_backbone.py b/flask/BiRefNet/models/backbones/build_backbone.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:f262062aea3f841f6e687b8544e98197e9a07d96f5edb67f705942e94104bc9f -size 2611 diff --git a/flask/BiRefNet/models/backbones/dino_v3.py b/flask/BiRefNet/models/backbones/dino_v3.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:ce92743fe56795c25cc4b0954ddb05ba3a396e25efa51cb4bc6dd7c782803bb8 -size 2337 diff --git a/flask/BiRefNet/models/backbones/pvt_v2.py b/flask/BiRefNet/models/backbones/pvt_v2.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:31550a9b1d9e1388a4d7d759bfbc2e2065a97e524a8036e618efcfd8941b7d14 -size 17064 diff --git a/flask/BiRefNet/models/backbones/swin_v1.py b/flask/BiRefNet/models/backbones/swin_v1.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a222afabd3f967d3bd95e9a88d902d57a2a0403f042df6b6f28e5d273ff484f0 -size 24710 diff --git a/flask/BiRefNet/models/birefnet.py b/flask/BiRefNet/models/birefnet.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7b3ccc1db3826e9cb350c5cd76de148d515826f9c19b342d9db617781c494871 -size 16702 diff --git a/flask/BiRefNet/models/modules/aspp.py b/flask/BiRefNet/models/modules/aspp.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:8fa5b8214294cb2a94eeb4cf470ad2a62db65655201152f8e421699493704e32 -size 5058 diff --git a/flask/BiRefNet/models/modules/decoder_blocks.py b/flask/BiRefNet/models/modules/decoder_blocks.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a133e43a24cbeb7e592181c136963eec3af52a660c747dd0bb3cd1f9d5d4c4d5 -size 2525 diff --git a/flask/BiRefNet/models/modules/deform_conv.py b/flask/BiRefNet/models/modules/deform_conv.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e16b32456eccf80b6ecd4d5ab64ef7f452b4f686cb9222e09b2de19abfe2bb42 -size 2439 diff --git a/flask/BiRefNet/models/modules/lateral_blocks.py b/flask/BiRefNet/models/modules/lateral_blocks.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:3ebc38a49438ddd4c705ccc6fa0e8bf903bd1f4d39cab8a4a2c04b7639d09809 -size 354 diff --git a/flask/BiRefNet/models/modules/utils.py b/flask/BiRefNet/models/modules/utils.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:0cd9561952e14985fb90d9785bba6e68039f8cf84a64b6e6da4d83e22561e8f2 -size 1481 diff --git a/flask/BiRefNet/requirements.txt b/flask/BiRefNet/requirements.txt @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:0d40d5a37d5427d70ea54aac210c972d54a3c05d99c9be5c024139804761b6b2 -size 155 diff --git a/flask/BiRefNet/rm_cache.sh b/flask/BiRefNet/rm_cache.sh @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:8fb2ff17f559976e2eb0d4901e68b95d096b2eee26487409fd2fc272205a8de2 -size 290 diff --git a/flask/BiRefNet/sub.sh b/flask/BiRefNet/sub.sh @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:10cc3576b8110a6d824abbdcdd14b727226bc67cc659bb0aafd1ba0b95daf905 -size 602 diff --git a/flask/BiRefNet/test.sh b/flask/BiRefNet/test.sh @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7e024bb1c604d8b717b9fa34cf0f2a0d1a17d24e5c90261d93bdb8c62297ad44 -size 710 diff --git a/flask/BiRefNet/train.py b/flask/BiRefNet/train.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:1b75f7d76ca74966587f676115d786f4944545703e5e4ad0c5f4c988e6d1eeff -size 11050 diff --git a/flask/BiRefNet/train.sh b/flask/BiRefNet/train.sh @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:3741a620a69afedbd5893fa3bedf1e5a38a5c463ed18c102eeeea34f637f17e8 -size 1463 diff --git a/flask/BiRefNet/train_test.sh b/flask/BiRefNet/train_test.sh @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:8c40d976a612b93cdb7f713bf52c88bb13dcd878fd33c8ec52eb0193eca9cc5e -size 244 diff --git a/flask/BiRefNet/tutorials/BiRefNet_inference.ipynb b/flask/BiRefNet/tutorials/BiRefNet_inference.ipynb @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:92d9d18ba6fde5d80fbf0b05375f3a2500ddb91c82991c62c3e0b3778f28f5b2 -size 58955 diff --git a/flask/BiRefNet/tutorials/BiRefNet_inference_video.ipynb b/flask/BiRefNet/tutorials/BiRefNet_inference_video.ipynb @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:2b9fb19d779a214eab60c4b100b85e97026e540ce6b22251d8fad8590dd4f8ec -size 61043 diff --git a/flask/BiRefNet/tutorials/BiRefNet_pth2onnx.ipynb b/flask/BiRefNet/tutorials/BiRefNet_pth2onnx.ipynb @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:129388fdf66b99d4cc7d193a63b61721bba631c368d4fcc668a7120056b7e90d -size 10453 diff --git a/flask/BiRefNet/utils.py b/flask/BiRefNet/utils.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7fbdcd4d5dedca7c251e1b5d6848494bb8512899b21f280dbe7a2625fa7043a4 -size 2908 diff --git a/flask/Florence-2-4bit-Quantized/added_tokens.json b/flask/Florence-2-4bit-Quantized/added_tokens.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:ba67457c577895525996e7a42d7da60a05856506ea4fd25ff9a023a64662f74b -size 22410 diff --git a/flask/Florence-2-4bit-Quantized/config.json b/flask/Florence-2-4bit-Quantized/config.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:161b960decdd1e205c551328cfab21aba2f53a352adc0076b9c1bfb6ac9b23e8 -size 6015 diff --git a/flask/Florence-2-4bit-Quantized/configuration_florence2.py b/flask/Florence-2-4bit-Quantized/configuration_florence2.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:de2e45a975b3582de05d2f4d963a3e9f9a3d20dccf78d28e0052932a0be93bdf -size 15119 diff --git a/flask/Florence-2-4bit-Quantized/generation_config.json b/flask/Florence-2-4bit-Quantized/generation_config.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e5bc064fef2956e8c40be5efa98bc45889bbe2b8015911045ccf7c0e42e1d8d8 -size 292 diff --git a/flask/Florence-2-4bit-Quantized/merges.txt b/flask/Florence-2-4bit-Quantized/merges.txt @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:1ce1664773c50f3e0cc8842619a93edc4624525b728b188a9e0be33b7726adc5 -size 456318 diff --git a/flask/Florence-2-4bit-Quantized/model.safetensors b/flask/Florence-2-4bit-Quantized/model.safetensors @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:ae4a6b8ee9e02a4dce366cc56084f9a97c0e0bf07a0da13f4f53e3f971e6d275 -size 200489948 diff --git a/flask/Florence-2-4bit-Quantized/modeling_florence2.py b/flask/Florence-2-4bit-Quantized/modeling_florence2.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:5162bf465e61b6e29cc113a467630ec3cb56ed8e4d46eb6207157f10fb9b8a24 -size 127455 diff --git a/flask/Florence-2-4bit-Quantized/preprocessor_config.json b/flask/Florence-2-4bit-Quantized/preprocessor_config.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:1396ec5a0a7adfe1c04fb777b09e8ba753be6dbb5868212ab3c3ef39d91fe031 -size 603 diff --git a/flask/Florence-2-4bit-Quantized/processing_florence2.py b/flask/Florence-2-4bit-Quantized/processing_florence2.py @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:f146023a507c009f425a49ee39aa037f4f25c64e14336e3e4f3f1d7377a68e98 -size 48676 diff --git a/flask/Florence-2-4bit-Quantized/processor_config.json b/flask/Florence-2-4bit-Quantized/processor_config.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:6eb1a2e487fa8b7521691a3382b5afd92a78e8ea34b385f8824fa07563e92245 -size 130 diff --git a/flask/Florence-2-4bit-Quantized/special_tokens_map.json b/flask/Florence-2-4bit-Quantized/special_tokens_map.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:72ff172dc769bc1551b1b4211628ce3271643bc60379e4da45d85a9be9332c39 -size 146627 diff --git a/flask/Florence-2-4bit-Quantized/tokenizer.json b/flask/Florence-2-4bit-Quantized/tokenizer.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:9b466914d9e7f9a39936c9bbe2ac28c86cd8b35c83a905699b01f83178f27c51 -size 3747961 diff --git a/flask/Florence-2-4bit-Quantized/tokenizer_config.json b/flask/Florence-2-4bit-Quantized/tokenizer_config.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:29290b47f21706f6a987cba4abb61e529fd71a607331860ace7c88d47041960a -size 197773 diff --git a/flask/Florence-2-4bit-Quantized/vocab.json b/flask/Florence-2-4bit-Quantized/vocab.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:ed19656ea1707df69134c4af35c8ceda2cc9860bf2c3495026153a133670ab5e -size 798293 diff --git a/flask/README.md b/flask/README.md @@ -31,6 +31,12 @@ modal setup modal deploy modal_app.py ``` +# Setup + +## Models + +Download all the models from [here](https://drive.google.com/drive/folders/143A296pgkTiUGqffE4xG8ADPPEGOY-W2) and place them inside this directory + ## Environment Variables Ensure that a `.env` file is present in this directory containing the following contents: diff --git a/flask/local_inpainting_model/feature_extractor/preprocessor_config.json b/flask/local_inpainting_model/feature_extractor/preprocessor_config.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:205d45875859846dc1ca211b9a19f5db72f47e953b396954ee34d48d8e95c3c6 -size 466 diff --git a/flask/local_inpainting_model/model_index.json b/flask/local_inpainting_model/model_index.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:68c908f4a7a55be3b21efddd2694f20c07056d6df55aabd0bf1d601c0a2b7163 -size 688 diff --git a/flask/local_inpainting_model/safety_checker/config.json b/flask/local_inpainting_model/safety_checker/config.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:43897735497f5b63df27d78e78af429d90efcf5803706398007e1ef5e2c53d5a -size 1167 diff --git a/flask/local_inpainting_model/safety_checker/model.safetensors b/flask/local_inpainting_model/safety_checker/model.safetensors @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:57ecdfa243b170f9b4cb3eefaf0f64552ef78fc0bf0eb1c5b9675308447184f6 -size 608016280 diff --git a/flask/local_inpainting_model/scheduler/scheduler_config.json b/flask/local_inpainting_model/scheduler/scheduler_config.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:f5105cbb737c7567c8743f0242340c6c3c34dce5c6652610a76e8bd44a2b6bac -size 528 diff --git a/flask/local_inpainting_model/text_encoder/config.json b/flask/local_inpainting_model/text_encoder/config.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d1cb5894d3c18c0a7f422186062e9f5a160f750fdafe53316a71358006b6f0b5 -size 554 diff --git a/flask/local_inpainting_model/text_encoder/model.safetensors b/flask/local_inpainting_model/text_encoder/model.safetensors @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:660c6f5b1abae9dc498ac2d21e1347d2abdb0cf6c0c0c8576cd796491d9a6cdd -size 246144152 diff --git a/flask/local_inpainting_model/tokenizer/merges.txt b/flask/local_inpainting_model/tokenizer/merges.txt @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:9fd691f7c8039210e0fced15865466c65820d09b63988b0174bfe25de299051a -size 524619 diff --git a/flask/local_inpainting_model/tokenizer/special_tokens_map.json b/flask/local_inpainting_model/tokenizer/special_tokens_map.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c4864a9376a8401918425bed71fc14fc0e81f9b59ec45c1cf96cccb2df508eac -size 472 diff --git a/flask/local_inpainting_model/tokenizer/tokenizer_config.json b/flask/local_inpainting_model/tokenizer/tokenizer_config.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c19f835a5e93d299b21dfee61c8f6a7cd06f5d83c8e828c013c25d2f36b47ccf -size 735 diff --git a/flask/local_inpainting_model/tokenizer/vocab.json b/flask/local_inpainting_model/tokenizer/vocab.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e089ad92ba36837a0d31433e555c8f45fe601ab5c221d4f607ded32d9f7a4349 -size 1059962 diff --git a/flask/local_inpainting_model/unet/config.json b/flask/local_inpainting_model/unet/config.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:038a57614b4bf87b2c20ffdf6ec39309cd40adf965f79436ea2d1297f7a3c47c -size 1932 diff --git a/flask/local_inpainting_model/unet/diffusion_pytorch_model.safetensors b/flask/local_inpainting_model/unet/diffusion_pytorch_model.safetensors @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:24b788b4a777748377cc20364eea4ae113c8c42f4468c16bc8c02fdae5492af9 -size 1719154104 diff --git a/flask/local_inpainting_model/vae/config.json b/flask/local_inpainting_model/vae/config.json @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d5bd71bcd70ade38b0367af3087fa3c4a6eb0e090fc70eca16a7282eaea5080d -size 916 diff --git a/flask/local_inpainting_model/vae/diffusion_pytorch_model.safetensors b/flask/local_inpainting_model/vae/diffusion_pytorch_model.safetensors @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:4fbcf0ebe55a0984f5a5e00d8c4521d52359af7229bb4d81890039d2aa16dd7c -size 167335342 diff --git a/models/color.py b/models/color.py @@ -0,0 +1,112 @@ +# -*- coding: utf-8 -*- +"""color_1.ipynb + +Automatically generated by Colab. + +Original file is located at + https://colab.research.google.com/drive/1o7cMSAesuuaxFQsteQL4ViwL_qNeTWbm +""" + +import cv2 +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.colors as mcolors +from sklearn.cluster import KMeans +from google.colab import files +from PIL import Image +import io + + +def get_dominant_colors(image, k=5): + """ + Extracts dominant colors using KMeans. + Expects a PIL Image or Numpy array (RGB). + """ + if isinstance(image, Image.Image): + image = np.array(image) + + img_small = cv2.resize(image, (150, 150), interpolation=cv2.INTER_AREA) + + pixels = img_small.reshape((-1, 3)) + + kmeans = KMeans(n_clusters=k, n_init='auto', random_state=42) + kmeans.fit(pixels) + + colors = kmeans.cluster_centers_.astype(int) + + return sorted(colors.tolist(), key=lambda x: sum(x)) + +def classify_mood(rgb_colors): + """ + Classifies mood based on HSV values using Matplotlib colors. + """ + norm_colors = np.array(rgb_colors) / 255.0 + + hsv_stats = mcolors.rgb_to_hsv(norm_colors) + + avg_sat = np.mean(hsv_stats[:, 1]) + avg_val = np.mean(hsv_stats[:, 2]) + + if avg_sat < 0.15 and avg_val > 0.65: return "Minimalist" + if avg_val < 0.35: return "Dark/Moody" + if avg_sat < 0.45 and avg_val > 0.75: return "Pastel" + if avg_sat > 0.65 and avg_val > 0.5: return "Neon" + + earthy_votes = sum(1 for h, s, v in hsv_stats if (0.02 <= h <= 0.42) and s < 0.8) + if earthy_votes >= 3: return "Earthy" + + warm_votes = sum(1 for h, s, v in hsv_stats if h < 0.17 or h > 0.83) + return "Warm" if warm_votes >= 3 else "Cool" + +def rgb_to_hex(rgb): + return '#{:02x}{:02x}{:02x}'.format(rgb[0], rgb[1], rgb[2]) + +# MAIN EXECUTION + +def analyze_uploaded_image(): + print("Please upload an image file...") + uploaded = files.upload() + + for fn in uploaded.keys(): + print(f'\nProcessing {fn}...') + + image_data = uploaded[fn] + image = Image.open(io.BytesIO(image_data)).convert('RGB') + + colors = get_dominant_colors(image) + if colors is None: + print("Error processing image.") + continue + + mood = classify_mood(colors) + hex_colors = [rgb_to_hex(c) for c in colors] + + print(f"Detected Mood: {mood}") + print("Extracted Palette:") + print("-" * 30) + for i, h in enumerate(hex_colors): + print(f"Color {i+1}: {h}") + print("-" * 30) + + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4)) + + ax1.imshow(image) + ax1.axis('off') + ax1.set_title(f"Mood: {mood}", fontweight='bold') + + y_pos = np.arange(len(colors)) + ax2.barh(y_pos, [1]*len(colors), color=hex_colors) + ax2.invert_yaxis() + ax2.axis('off') + + for i, h in enumerate(hex_colors): + brightness = sum(colors[i]) / 3 + text_col = 'black' if brightness > 128 else 'white' + ax2.text(0.5, i, f"{h}", ha='center', va='center', + color=text_col, fontweight='bold', fontsize=12) + + plt.tight_layout() + plt.show() + +if __name__ == "__main__": + analyze_uploaded_image() +\ No newline at end of file diff --git a/models/composition.py b/models/composition.py @@ -0,0 +1,441 @@ +# -*- coding: utf-8 -*- +"""layout + +Automatically generated by Colab. + +Original file is located at + https://colab.research.google.com/drive/1NzX6rAQSJkPTpdWWkOowDFlBpypoKl2U +""" + +pip install opencv-python numpy + +import cv2 +import numpy as np +import os +from math import hypot + +def resize_for_fast_processing(img, max_side=640): + h, w = img.shape[:2] + scale = max_side / max(h, w) if max(h, w) > max_side else 1.0 + if scale != 1.0: + img = cv2.resize(img, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA) + return img, scale + +def compute_saliency_gray(img): + sal = None + try: + salience = cv2.saliency.StaticSaliencySpectralResidual_create() + success, sal = salience.computeSaliency(img) + if not success: sal = None + except Exception: + sal = None + if sal is None: + gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) + blur1 = cv2.GaussianBlur(gray, (3, 3), 0) + blur2 = cv2.GaussianBlur(gray, (21, 21), 0) + sal = cv2.absdiff(blur1, blur2) + sal = sal.astype(np.float32) / (sal.max() + 1e-9) + else: + sal = sal.astype(np.float32) + if sal.max() > 0: sal /= sal.max() + return sal + +def saliency_centroid(sal): + h, w = sal.shape + M = sal.sum() + if M <= 1e-6: return (w / 2, h / 2) + ys, xs = np.indices(sal.shape) + cx = (xs * sal).sum() / M + cy = (ys * sal).sum() / M + return (float(cx), float(cy)) + +def detect_lines(gray, canny_thresh1=50, canny_thresh2=150, hough_thresh=50): + edges = cv2.Canny(gray, canny_thresh1, canny_thresh2) + lines = cv2.HoughLinesP(edges, rho=1, theta=np.pi / 180, threshold=hough_thresh, + minLineLength=min(gray.shape) // 10, maxLineGap=10) + if lines is None: return edges, [] + lines = [tuple(l[0]) for l in lines] + return edges, lines + +def project_sal_along_axis(sal, axis=0, smooth_k=15): + proj = sal.sum(axis=1-axis) if axis == 0 else sal.sum(axis=0) + k = max(3, smooth_k) + kernel = np.ones(k) / k + proj_s = np.convolve(proj, kernel, mode='same') + if proj_s.max() > 0: + proj_s = proj_s / proj_s.max() + return proj_s + +def count_saliency_peaks_along_y(sal, min_sep_fraction=0.1, threshold=0.2): + h, w = sal.shape + proj = project_sal_along_axis(sal, axis=0, smooth_k=max(5, int(h*0.03))) + peaks = [] + for i in range(1, len(proj)-1): + if proj[i] > proj[i-1] and proj[i] > proj[i+1] and proj[i] > threshold: + peaks.append(i) + min_sep = int(min_sep_fraction * h) + filtered = [] + for p in peaks: + if not filtered or p - filtered[-1] >= min_sep: + filtered.append(p) + return len(filtered), filtered + +def band_saliency_fractions(sal, bands=3): + h, w = sal.shape + sums = [] + total = sal.sum() + 1e-9 + for i in range(bands): + band = sal[:, i*w//bands:(i+1)*w//bands] + sums.append(band.sum() / total) + return sums + +def sal_bbox_and_margins(sal, thresh=0.3): + mask = sal > thresh + if mask.sum() == 0: + return (0,0,0,0), None + ys, xs = np.where(mask) + minx, maxx = int(xs.min()), int(xs.max()) + miny, maxy = int(ys.min()), int(ys.max()) + return (minx, miny, maxx, maxy), mask + +def component_bboxes_from_mask(mask): + mask_u8 = (mask.astype(np.uint8) * 255) + contours, _ = cv2.findContours(mask_u8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + bboxes = [] + for c in contours: + x,y,w,h = cv2.boundingRect(c) + bboxes.append((x,y,x+w-1,y+h-1)) + return bboxes + +def bbox_laplacian_variance(gray, bbox): + x1,y1,x2,y2 = bbox + roi = gray[y1:y2+1, x1:x2+1] + if roi.size == 0: return 0.0 + lap = cv2.Laplacian(roi, cv2.CV_64F) + return float(np.var(lap)) + + + +def _proj_and_peaks(sal, axis=0, smooth_k=15): + proj = sal.sum(axis=1-axis) if axis==0 else sal.sum(axis=0) + k = max(3, smooth_k) + kernel = np.ones(k) / k + proj_s = np.convolve(proj, kernel, mode='same') + if proj_s.max() > 0: proj_s = proj_s / proj_s.max() + return proj_s + +def vertical_stack_score(sal, lines, min_peaks=2): + h,w = sal.shape + proj = _proj_and_peaks(sal, axis=0, smooth_k=max(5, int(h*0.03))) + peaks = [i for i in range(1,len(proj)-1) if proj[i]>proj[i-1] and proj[i]>proj[i+1] and proj[i]>0.18] + if len(peaks)==0: + return 0.0, {'peaks':0} + avg_peak = float(np.mean([proj[p] for p in peaks])) + peak_count_score = min(1.0, len(peaks) / max(1.0, min_peaks)) + peak_prom_score = avg_peak + h_lines = 0 + for (x1,y1,x2,y2) in lines: + dy = abs(y2-y1); dx = abs(x2-x1)+1e-9 + slope = dy/dx + if slope < 0.3 and abs(x2-x1) > 0.6*w: + h_lines += 1 + sep_score = min(1.0, h_lines / 2.0) + span_ok = 0 + for p in peaks: + row = sal[max(0,p-2):min(h,p+3), :] + if row.sum()<=1e-9: continue + col_sum = row.sum(axis=0) + frac_nonzero = (col_sum > (col_sum.max()*0.05)).sum() / float(w) + if frac_nonzero > 0.65: + span_ok += 1 + span_score = min(1.0, span_ok / max(1.0, len(peaks))) + score = 0.45*peak_count_score + 0.25*peak_prom_score + 0.15*sep_score + 0.15*span_score + return float(np.clip(score,0.0,1.0)), {'peaks': len(peaks), 'avg_peak': avg_peak, 'h_lines': h_lines, 'span_ok': span_ok} + +def triptych_score(sal, lines): + h,w = sal.shape + total = sal.sum() + 1e-9 + bands = [sal[:, i*w//3:(i+1)*w//3].sum() / total for i in range(3)] + mn, mx = min(bands), max(bands) + balance = 1.0 - (mx - mn) + balance = np.clip(balance, 0.0, 1.0) + v_hits = 0 + for (x1,y1,x2,y2) in lines: + dx = abs(x2-x1)+1e-9; dy = abs(y2-y1) + slope = dy / dx if dx>1 else 1e9 + if slope > 3 and dy > 0.6*h: + cx = (x1+x2)/2 + if abs(cx - w/3) < 0.08*w or abs(cx - 2*w/3) < 0.08*w: + v_hits += 1 + line_score = np.tanh(v_hits/2.0) + band_floor = min(1.0, mn / 0.15) + score = 0.55*balance + 0.25*line_score + 0.20*band_floor + return float(np.clip(score,0.0,1.0)), {'band_fracs': bands, 'v_line_hits': v_hits} + +def full_bleed_score(sal, img, edge_margin_frac=0.03, + edge_sal_thresh=0.25, + sal_bbox_thresh=0.2, + white_thresh=245, + white_frac_thresh=0.80, + edge_sal_frac_thresh=0.50, + bbox_touch_required=0.75): + """ + Binary full-bleed detector (returns 0 or 1, plus info). + - Rejects as full-bleed if a white margin/border is detected. + - Uses saliency near edges and whether saliency bbox touches edges. + """ + h, w = sal.shape + m = max(1, int(min(h, w) * edge_margin_frac)) + + + mask_edge = np.zeros_like(sal, dtype=bool) + mask_edge[:m, :] = True + mask_edge[-m:, :] = True + mask_edge[:, :m] = True + mask_edge[:, -m:] = True + edge_sal_frac = float((sal[mask_edge] > edge_sal_thresh).sum()) / (mask_edge.sum() + 1e-9) + + + mask2 = sal > sal_bbox_thresh + if mask2.sum() == 0: + bbox_touch = 0.0 + else: + ys, xs = np.where(mask2) + minx, maxx = int(xs.min()), int(xs.max()) + miny, maxy = int(ys.min()), int(ys.max()) + touches = int(minx <= 1) + int(maxx >= w - 2) + int(miny <= 1) + int(maxy >= h - 2) + bbox_touch = touches / 4.0 + + + band_top = img[:m, :, :] if m > 0 else np.zeros((0, w, 3), dtype=img.dtype) + band_bottom = img[-m:, :, :] if m > 0 else np.zeros((0, w, 3), dtype=img.dtype) + band_left = img[:, :m, :] if m > 0 else np.zeros((h, 0, 3), dtype=img.dtype) + band_right = img[:, -m:, :] if m > 0 else np.zeros((h, 0, 3), dtype=img.dtype) + + + border_pixels = np.concatenate([band_top.reshape(-1, 3), + band_bottom.reshape(-1, 3), + band_left.reshape(-1, 3), + band_right.reshape(-1, 3)], axis=0) + if border_pixels.size == 0: + white_border_frac = 0.0 + else: + white_mask = np.all(border_pixels >= white_thresh, axis=1) + white_border_frac = float(white_mask.sum()) / float(border_pixels.shape[0]) + + border_std = float(np.std(border_pixels)) if border_pixels.size > 0 else 0.0 + border_uniform = border_std < 6.0 # small std => uniform border + + cond_edge_sal = edge_sal_frac >= edge_sal_frac_thresh + cond_bbox_touch = bbox_touch >= bbox_touch_required + + is_full_bleed = (cond_edge_sal or cond_bbox_touch) and (white_border_frac < white_frac_thresh) and (not (border_uniform and white_border_frac > 0.25)) + + score = 1 if is_full_bleed else 0 + + info = { + 'edge_margin_px': m, + 'edge_sal_frac': edge_sal_frac, + 'edge_sal_thresh': edge_sal_thresh, + 'cond_edge_sal': cond_edge_sal, + 'bbox_touch': bbox_touch, + 'bbox_touch_required': bbox_touch_required, + 'white_border_frac': white_border_frac, + 'white_thresh': white_thresh, + 'white_frac_thresh': white_frac_thresh, + 'border_std': border_std, + 'border_uniform': border_uniform, + 'decision_full_bleed': bool(is_full_bleed) + } + return score, info + + +def tight_crop_score(sal, gray, sal_thresh=0.1): + h,w = sal.shape + mask = sal > sal_thresh + if mask.sum()==0: + return 0.0, {'area_frac':0.0, 'tightness':1.0} + ys,xs = np.where(mask) + minx,maxx = int(xs.min()), int(xs.max()) + miny,maxy = int(ys.min()), int(ys.max()) + bbox_area = (maxx-minx+1)*(maxy-miny+1) + area_frac = bbox_area / float(w*h) + left = minx; right = w-1-maxx; top = miny; bottom = h-1-maxy + tightness = min(left,right,top,bottom) / float(min(w,h)+1e-9) + area_score = np.clip((area_frac - 0.15) / (0.6 - 0.15), 0.0, 1.0) + tight_score = 1.0 - np.clip(tightness / 0.08, 0.0, 1.0) + score = 0.6*area_score + 0.4*tight_score + return float(np.clip(score,0.0,1.0)), {'area_frac': area_frac, 'tightness': tightness, 'margins': (left,right,top,bottom)} + +def layered_foreground_score(sal, gray): + mask = (sal > 0.25).astype(np.uint8) + kernel = np.ones((5,5), np.uint8) + mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel) + mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel) + contours, _ = cv2.findContours((mask*255).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + bboxes=[]; sharpness=[] + for c in contours: + x,y,w,h = cv2.boundingRect(c) + if w*h < 0.005*sal.size: continue + bboxes.append((x,y,x+w-1,y+h-1)) + roi = gray[y:y+h, x:x+w] + lap = cv2.Laplacian(roi, cv2.CV_64F) + sharpness.append(float(np.var(lap))) + n = len(bboxes) + if n < 2: + return 0.0, {'num_components': n} + s = np.array(sharpness) + 1e-9 + s_norm = (s - s.mean()) / (s.std()+1e-9) + sharp_std = float(np.std(s_norm)) + overlaps = 0 + for i in range(len(bboxes)): + for j in range(i+1, len(bboxes)): + a=bboxes[i]; b=bboxes[j] + ix1 = max(a[0], b[0]); iy1 = max(a[1], b[1]) + ix2 = min(a[2], b[2]); iy2 = min(a[3], b[3]) + if ix2 >= ix1 and iy2 >= iy1: + inter = (ix2-ix1+1)*(iy2-iy1+1) + area_a = (a[2]-a[0]+1)*(a[3]-a[1]+1) + if inter > 0.05*area_a: + overlaps += 1 + comp_score = np.clip((n-1)/4.0, 0, 1) + sharp_score = np.tanh(sharp_std) + overlap_score = np.tanh(overlaps/2.0) + score = 0.45*comp_score + 0.35*sharp_score + 0.20*overlap_score + return float(np.clip(score,0.0,1.0)), {'num_components': n, 'sharp_std': sharp_std, 'overlaps': overlaps} + + +def balance_score(img_bgr, sal): + img_lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB).astype(np.float32) + L = img_lab[:,:,0] / 255.0 + img_hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV).astype(np.float32) + S = img_hsv[:,:,1] / 255.0 + weight = sal * (L + 0.5 * S) + M = weight.sum() + h,w = sal.shape + if M <= 1e-6: return 0.0 + ys, xs = np.indices(sal.shape) + cx = (xs * weight).sum() / M + cy = (ys * weight).sum() / M + dx = abs(cx - w/2) / (w/2) + dy = abs(cy - h/2) / (h/2) + dist = np.sqrt(dx*dx + dy*dy) / np.sqrt(2) + return float(max(0.0, 1.0 - dist*1.4)) + +def symmetry_score(img_gray): + h,w = img_gray.shape + left = img_gray[:, :w//2] + right = img_gray[:, w - w//2:] + right_flipped = cv2.flip(right, 1) + if left.shape != right_flipped.shape: + right_flipped = cv2.resize(right_flipped, (left.shape[1], left.shape[0])) + leftf = left.astype(np.float32); rightf = right_flipped.astype(np.float32) + if leftf.std() < 1e-3 or rightf.std() < 1e-3: return 0.0 + corr = np.corrcoef(leftf.flatten(), rightf.flatten())[0,1] + return float(max(0.0, min(1.0, (corr + 1)/2))) + +def thirds_score(sal, centroid): + h,w = sal.shape + cx, cy = centroid + thirds_pts = [(w/3, h/3), (w/3, 2*h/3), (2*w/3, h/3), (2*w/3, 2*h/3)] + dists = [hypot(cx-x, cy-y) for (x,y) in thirds_pts] + maxd = hypot(w, h) + score = 1.0 - min(dists)/maxd*1.2 + vline_dist = min(abs(cx - w/3), abs(cx - 2*w/3)) / w + hline_dist = min(abs(cy - h/3), abs(cy - 2*h/3)) / h + line_score = 1.0 - min(vline_dist, hline_dist) + return float(max(0.0, min(1.0, 0.6*score + 0.4*line_score))) + +def golden_ratio_score(sal, centroid): + h,w = sal.shape + gx = [0.382, 0.618]; gy = [0.382, 0.618] + cx, cy = centroid + pts = [(g*w, g2*h) for g in gx for g2 in gy] + dists = [hypot(cx-x, cy-y) for x,y in pts] + maxd = hypot(w, h) + return float(max(0.0, min(1.0, 1.0 - min(dists)/ (maxd*0.9)))) + +def negative_space_score(sal): + low = (sal < 0.15).sum() / sal.size + return float(min(1.0, low)) + +def center_score(sal, centroid): + h,w = sal.shape + cx, cy = centroid + dx = abs(cx - w/2) / (w/2) + dy = abs(cy - h/2) / (h/2) + dist = np.sqrt(dx*dx + dy*dy) / np.sqrt(2) + return float(max(0.0, 1.0 - dist*1.2)) + + + +def analyze_single_image(img_path): + if not os.path.exists(img_path): + print(f"Error: File not found at {img_path}") + return + + img_full = cv2.imread(img_path) + if img_full is None: + print(f"Error: Could not load image. Check format.") + return + + print(f"Analyzing: {os.path.basename(img_path)}...") + + img, scale = resize_for_fast_processing(img_full, max_side=640) + h, w = img.shape[:2] + gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) + + sal = compute_saliency_gray(img) + if sal.shape != gray.shape: + sal = cv2.resize(sal, (w, h), interpolation=cv2.INTER_LINEAR) + centroid = saliency_centroid(sal) + + edges, lines = detect_lines(gray) + _, thr = cv2.threshold((gray).astype(np.uint8), 0, 255, cv2.THRESH_OTSU + cv2.THRESH_BINARY_INV) + contours, _ = cv2.findContours(thr, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + composition_scores = { + 'Rule of Thirds': thirds_score(sal, centroid), + 'Golden Ratio': golden_ratio_score(sal, centroid), + 'Symmetry': symmetry_score(gray), + 'Negative Space': negative_space_score(sal), + 'Center Composition': center_score(sal, centroid) + } + + v_score, v_info = vertical_stack_score(sal, lines) + t_score, t_info = triptych_score(sal, lines) + fb_score, fb_info = full_bleed_score(sal, img) + tc_score, tc_info = tight_crop_score(sal, gray) + lf_score, lf_info = layered_foreground_score(sal, gray) + + composition_scores['Vertical Stack'] = v_score + composition_scores['Triptych'] = t_score + composition_scores['Full Bleed'] = fb_score + composition_scores['Tight Crop'] = tc_score + composition_scores['Layered Foreground'] = lf_score + + print("-" * 40) + print(f"Top Composition Scores:") + sorted_comp = sorted(composition_scores.items(), key=lambda x: x[1], reverse=True) + for k, v in sorted_comp[:10]: + print(f" - {k}: {v:.2f}") + + print("\n All Composition Scores:") + for k, v in sorted_comp: + print(f" {k:<25} : {v:.3f}") + print("-" * 40) + + diagnostics = { + 'Vertical Stack_info': v_info, + 'Triptych_info': t_info, + 'Full Bleed_info': fb_info, + 'Tight Crop_info': tc_info, + 'Layered Foreground_info': lf_info + } + + +# MAIN EXECUTION + +if __name__ == '__main__': + target_image_path = "/content/Screenshot 2025-11-24 200949.png" + analyze_single_image(target_image_path) +\ No newline at end of file diff --git a/models/emotions.py b/models/emotions.py @@ -0,0 +1,164 @@ +# -*- coding: utf-8 -*- +"""Untitled4.ipynb + +Automatically generated by Colab. + +Original file is located at + https://colab.research.google.com/drive/1SgZgWv7uK3kibT3KW1zR5jp5pH46AQE0 +""" + +!pip install git+https://github.com/openai/CLIP.git --q +!pip install torch torchvision pillow --q + +!pip install ftfy regex tqdm + +import torch +import clip +from PIL import Image +import numpy as np + +device = "cuda" if torch.cuda.is_available() else "cpu" +model, preprocess = clip.load("ViT-B/32", device=device) + +print("Device:", device) + +EMOTIONAL_CLASS = [ + "sensual", + "playful", + "edgy", + "bold", + "luxurious", + "cheeky", + "ironic", + "mysterious" +] + +EMOTIONAL_PHRASES = { + "sensual": "a sensual emotional tone, conveying intimacy and softness", + "playful": "a playful and fun emotional mood with cheerful energy", + "edgy": "an edgy emotional feel with intensity and attitude", + "bold": "a bold, fearless emotional expression with strong presence", + "luxurious": "a luxurious emotional mood with elegance and refinement", + "cheeky": "a cheeky, mischievous emotional attitude with confidence", + "ironic": "an ironic, detached emotional vibe with witty contrast", + "mysterious": "a mysterious emotional tone filled with intrigue and ambiguity", +} + + +EMOTION_PROMPT_TEMPLATES = [ + "a photo that conveys {}", + "a portrait expressing {}", + "a cinematic shot capturing {}", + "an artistic image filled with {}", + "a dramatic photograph embodying {}", + "a close-up shot that reveals {}", + "a high-fashion editorial with {}", + "a lifestyle photo characterized by {}", + "a conceptual art image exploring {}", + "a professional photoshoot showcasing {}", +] + +def build_text_bank(output_path="lighting_text_bank"): + all_prompt_texts = [] + all_prompt_labels = [] + all_prompt_features = [] + + class_centroids = [] + + model.eval() + with torch.no_grad(): + for lighting in EMOTIONAL_CLASS: + phrase = EMOTIONAL_PHRASES[lighting] + + prompts = [tpl.format(phrase) for tpl in EMOTION_PROMPT_TEMPLATES] + tokens = clip.tokenize(prompts).to(device) + + text_features = model.encode_text(tokens) + text_features = text_features / text_features.norm(dim=-1, keepdim=True) + + all_prompt_texts.extend(prompts) + all_prompt_labels.extend([lighting] * len(prompts)) + all_prompt_features.append(text_features.cpu().numpy()) + + centroid = text_features.mean(dim=0) + centroid = centroid / centroid.norm() + class_centroids.append(centroid.cpu().numpy()) + + all_prompt_features = np.concatenate(all_prompt_features, axis=0) + class_centroids = np.stack(class_centroids, axis=0) + + np.savez( + output_path, + prompt_texts=np.array(all_prompt_texts), + prompt_labels=np.array(all_prompt_labels), + prompt_features=all_prompt_features, + class_centroids=class_centroids, + class_names=np.array(EMOTIONAL_CLASS), + ) + print(f"Saved text bank → {output_path}") + + +build_text_bank() + +def get_image_embedding(image_path: str) -> torch.Tensor: + image = Image.open(image_path).convert("RGB") + img_tensor = preprocess(image).unsqueeze(0).to(device) + + model.eval() + with torch.no_grad(): + img_feat = model.encode_image(img_tensor) + img_feat = img_feat / img_feat.norm(dim=-1, keepdim=True) + + return img_feat.squeeze(0) + +DEFAULT_THRESHOLD = 0.22 + + +def load_text_bank(path="/content/lighting_text_bank.npz"): + bank = np.load(path, allow_pickle=True) + centroids = torch.tensor(bank["class_centroids"]).to(device) + class_names = bank["class_names"] + + centroids = centroids / centroids.norm(dim=-1, keepdim=True) + return centroids, class_names + + +def classify_lighting(image_path, threshold=DEFAULT_THRESHOLD): + img_feat = get_image_embedding(image_path).to(device) + + centroids, class_names = load_text_bank() + + img_feat = img_feat.to(centroids.dtype) + + img_feat = img_feat / img_feat.norm() + + sims = centroids @ img_feat + best_score, best_idx = torch.max(sims, dim=0) + + best_score = float(best_score.detach().cpu()) + sims_list = sims.detach().cpu().tolist() + + sims_dict = { + str(cls): float(s) + for cls, s in zip(class_names, sims_list) + } + + if best_score < threshold: + predicted = "no_specific_emotion" + else: + predicted = str(class_names[best_idx]) + + return predicted, best_score, sims_dict + +from google.colab import files +uploaded = files.upload() + +image_path = list(uploaded.keys())[0] + +pred, score, sims = classify_lighting(image_path) + +print("Prediction:", pred) +print("Best score:", score) +print("\nAll similarities:") +for k, v in sims.items(): + print(f"{k:20s}: {v:.3f}") +\ No newline at end of file diff --git a/models/era_cult_ref.py b/models/era_cult_ref.py @@ -0,0 +1,139 @@ +# -*- coding: utf-8 -*- +"""Era/Cult_Ref.ipynb + +Automatically generated by Colab. + +Original file is located at + https://colab.research.google.com/drive/15k15A0X38sOgwqGR5j7QhDBS8_61R_ca +""" + +!pip install git+https://github.com/openai/CLIP.git --q +!pip install torch torchvision ftfy regex tqdm --q + +import torch +import clip +from PIL import Image +import numpy as np +from google.colab import files +import io + +device = "cuda" if torch.cuda.is_available() else "cpu" +print(f"Using device: {device}") + +model, preprocess = clip.load("ViT-B/32", device=device) + +ERA_CLASSES = [ + "vintage-60s", + "modern", + "90s-magazine", + "art-deco", + "punk", + "avant-garde" +] + +ERA_DESCRIPTIONS = { + "vintage-60s": "1960s aesthetic, psychedelic art, mod fashion, kodachrome film look, flower power era", + "modern": "modern contemporary aesthetic, clean design, digital photography, high resolution, 21st century style", + "90s-magazine": "1990s fashion magazine style, grunge aesthetic, harsh flash photography, teen vogue 90s, glossy editorial", + "art-deco": "art deco style, 1920s geometric patterns, gold and black luxury, great gatsby aesthetic, streamline moderne", + "punk": "punk rock aesthetic, diy zine style, xeroxed textures, safety pins, rebellious and chaotic visual style", + "avant-garde": "avant-garde fashion, experimental art, high concept, unconventional silhouette, futuristic and edgy" +} + +PROMPT_TEMPLATES = [ + "a photo representing the {} aesthetic", + "artistic rendering in {} style", + "an image from the {} era", + "visuals depicting {}", + "a high quality example of {} culture", + "a magazine scan showing {}", + "fashion photography in the style of {}", + "a poster with {} design elements", +] + + +class EraClassifier: + def __init__(self): + self.centroids = None + self.class_names = ERA_CLASSES + self.build_text_features() + + def build_text_features(self): + """ + Pre-computes the embedding centroids for all eras. + """ + print("Building text embeddings for eras...") + class_centroids = [] + + with torch.no_grad(): + for label in self.class_names: + desc = ERA_DESCRIPTIONS[label] + + prompts = [tpl.format(desc) for tpl in PROMPT_TEMPLATES] + + tokens = clip.tokenize(prompts).to(device) + text_features = model.encode_text(tokens) + + text_features = text_features / text_features.norm(dim=-1, keepdim=True) + + centroid = text_features.mean(dim=0) + centroid = centroid / centroid.norm() + + class_centroids.append(centroid) + + self.centroids = torch.stack(class_centroids).to(device) + print("Text embeddings built successfully.") + + def predict(self, image_path, threshold=0.20): + image = Image.open(image_path).convert("RGB") + img_tensor = preprocess(image).unsqueeze(0).to(device) + + with torch.no_grad(): + img_feat = model.encode_image(img_tensor) + img_feat = img_feat / img_feat.norm(dim=-1, keepdim=True) + + similarity = (100.0 * img_feat @ self.centroids.T) + + probs = similarity.softmax(dim=-1).cpu().numpy()[0] + raw_scores = similarity.cpu().numpy()[0] / 100.0 + + results = [] + for i, label in enumerate(self.class_names): + results.append({ + "label": label, + "score": float(raw_scores[i]), + "confidence_pct": float(probs[i] * 100) + }) + + results.sort(key=lambda x: x["score"], reverse=True) + top_result = results[0] + + if top_result["score"] < threshold: + prediction = "undefined/mixed" + else: + prediction = top_result["label"] + + return prediction, results + +# RUNNER + +classifier = EraClassifier() + +print("\nPlease upload an image to analyze...") +uploaded = files.upload() + +if len(uploaded) > 0: + img_path = list(uploaded.keys())[0] + + prediction, details = classifier.predict(img_path) + + print("\n" + "="*30) + print(f"PREDICTED ERA: {prediction.upper()}") + print("="*30) + + print(f"\nDetailed Breakdown:") + print(f"{'Era/Style':<15} | {'Cosine Sim':<12} | {'Rel Confidence':<15}") + print("-" * 45) + + for item in details: + print(f"{item['label']:<15} | {item['score']:.4f} | {item['confidence_pct']:.1f}%") +\ No newline at end of file diff --git a/models/inpainting_api&local.py b/models/inpainting_api&local.py @@ -0,0 +1,275 @@ +# -*- coding: utf-8 -*- +"""inpainting (1) (1).ipynb + +Automatically generated by Colab. + +Original file is located at + https://colab.research.google.com/drive/15HZcoDGdzTaDO9d1uIUMI8LQToCP4ryC +""" + +import numpy as np +import pandas as pd + +!pip install fal-client + +import fal_client +import requests +import os +import numpy as np +import scipy.ndimage +from io import BytesIO +from PIL import Image, ImageFilter +from IPython.display import display + +os.environ["FAL_KEY"] = "YOUR_API_KEY" + +CLEAN_IMAGE_PATH = "clean_image_path" +DRAWN_IMAGE_PATH = "drawn_image_path" +MASK_OUTPUT_PATH = "solid_mask.png" + +def generate_and_save_mask(clean_path, drawn_path, save_path): + print("🛠️ Generating mask with robust filtering...") + + if not os.path.exists(clean_path) or not os.path.exists(drawn_path): + print(" Error: Input images not found.") + return False + + img_clean = Image.open(clean_path).convert("RGB").resize((1024, 1024)) + img_drawn = Image.open(drawn_path).convert("RGB").resize((1024, 1024)) + + clean_blur = np.array(img_clean.filter(ImageFilter.GaussianBlur(2)), dtype=np.int16) + drawn_blur = np.array(img_drawn.filter(ImageFilter.GaussianBlur(2)), dtype=np.int16) + + diff_arr = np.abs(drawn_blur - clean_blur) + mask_arr = np.max(diff_arr, axis=2) + + mask_binary = mask_arr > 30 + + mask_filled = scipy.ndimage.binary_fill_holes(mask_binary) + mask = Image.fromarray((mask_filled * 255).astype(np.uint8)) + mask = mask.filter(ImageFilter.MaxFilter(9)) + + mask.save(save_path) + img_clean.save("resized_clean.png") + + print(" The AI will ONLY edit the WHITE area below:") + display(mask.resize((256, 256))) + + return True + + +def run_fal_inpainting(image_path, mask_path): + print(" Uploading images to Fal.ai...") + + image_url = fal_client.upload_file(image_path) + mask_url = fal_client.upload_file(mask_path) + + print(" Running Flux Dev Fill...") + + try: + handler = fal_client.submit( + "fal-ai/flux-lora-fill", + arguments={ + "prompt": "The image shows a river running through a lush green valley surrounded by trees, plants, grass, and poles. In the background, the sky is filled with clouds, creating a peaceful atmosphere.", + "image_url": image_url, + "mask_url": mask_url, + "guidance_scale": 30, + "num_inference_steps": 28, + "enable_safety_checker": False + } + ) + + result = handler.get() + + if 'images' in result: + output_url = result['images'][0]['url'] + print(f"✨ Success! Image generated: {output_url}") + + response = requests.get(output_url) + img = Image.open(BytesIO(response.content)) + img.save("final_output.png") + display(img) + else: + print(" API returned no images.") + print(result) + + except Exception as e: + print(f"Error during API call: {e}") + + +# MAIN EXECUTION + +if __name__ == "__main__": + success = generate_and_save_mask(CLEAN_IMAGE_PATH, DRAWN_IMAGE_PATH, MASK_OUTPUT_PATH) + if success: + run_fal_inpainting("resized_clean.png", MASK_OUTPUT_PATH) + +!pip install torch diffusers transformers accelerate gradio + +import torch +import numpy as np +import matplotlib.pyplot as plt +from PIL import Image, ImageFilter +from diffusers import AutoPipelineForInpainting +import os +import scipy.ndimage + +def load_model(): + model_id = "runwayml/stable-diffusion-inpainting" + device = "cuda" if torch.cuda.is_available() else "cpu" + + print(f"⏳ Loading model to {device} (this may take a minute)...") + try: + pipe = AutoPipelineForInpainting.from_pretrained( + model_id, + torch_dtype=torch.float16, + variant="fp16" + ).to(device) + pipe.enable_attention_slicing() + return pipe + except Exception as e: + print(f"Error loading model: {e}") + return None + +def run_sketch_to_image(pipe, clean_path, drawn_path, prompt, negative_prompt="", strength=0.85, seed=42): + if not os.path.exists(clean_path) or not os.path.exists(drawn_path): + print(" Error: Images not found.") + return + + img_clean = Image.open(clean_path).convert("RGB").resize((512, 512)) + img_drawn = Image.open(drawn_path).convert("RGB").resize((512, 512)) + + + clean_blur = np.array(img_clean.filter(ImageFilter.GaussianBlur(radius=2)), dtype=np.int16) + drawn_blur = np.array(img_drawn.filter(ImageFilter.GaussianBlur(radius=2)), dtype=np.int16) + + diff_arr = np.abs(drawn_blur - clean_blur) + + mask_arr = np.max(diff_arr, axis=2) + + mask_binary = mask_arr > 30 + + mask_filled = scipy.ndimage.binary_fill_holes(mask_binary) + mask = Image.fromarray((mask_filled * 255).astype(np.uint8)) + + mask = mask.filter(ImageFilter.MaxFilter(9)) + + generator = torch.Generator(device="cuda").manual_seed(seed) + + print(f" Generating with Strength {strength}...") + with torch.inference_mode(): + output = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + image=img_drawn, + mask_image=mask, + strength=strength, + guidance_scale=8.0, + num_inference_steps=50, + generator=generator + ).images[0] + + fig, axs = plt.subplots(1, 4, figsize=(20, 6)) + axs[0].imshow(img_clean); axs[0].set_title("Initial") + axs[1].imshow(img_drawn); axs[1].set_title("Sketch") + axs[2].imshow(mask, cmap="gray"); axs[2].set_title("Calculated Mask (Fixed)") + axs[3].imshow(output); axs[3].set_title("Result") + plt.tight_layout() + plt.show() + return output + +if __name__ == "__main__": + if 'pipeline' not in globals(): + pipeline = load_model() + + if pipeline: + run_sketch_to_image( + pipe=pipeline, + clean_path="/kaggle/input/sample/scenery.jpg.png", + drawn_path="/kaggle/input/sample/scenery_sketch.jpg", + prompt="The image shows a river running through a lush green valley surrounded by trees, plants, grass, and poles. In the background, the sky is filled with clouds, creating a peaceful atmosphere.", + negative_prompt="", + strength=0.85, + seed=100 + ) + + + +!pip install -q einops timm + +import torch +from transformers import AutoProcessor, AutoModelForCausalLM +from PIL import Image + +device = "cuda" if torch.cuda.is_available() else "cpu" +model_id = "microsoft/Florence-2-base" + +print(f" Loading Model: {model_id}...") + +model = AutoModelForCausalLM.from_pretrained( + model_id, + trust_remote_code=True, + torch_dtype=torch.float16 +).to(device) + +processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) + +print(" Model Loaded & Moved to GPU. Ready for Inference.") + +import time +import torch +from PIL import Image + +def generate_concise_caption(image_path): + device = "cuda" if torch.cuda.is_available() else "cpu" + dtype = torch.float16 + + torch.cuda.reset_peak_memory_stats() + torch.cuda.synchronize() + start_time = time.time() + + image = Image.open(image_path).convert("RGB") + + + prompt_task = "<DETAILED_CAPTION>" + + inputs = processor(text=prompt_task, images=image, return_tensors="pt").to(device, dtype) + + generated_ids = model.generate( + input_ids=inputs["input_ids"], + pixel_values=inputs["pixel_values"], + max_new_tokens=64, + num_beams=1, + do_sample=False + ) + + generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0] + caption = processor.post_process_generation( + generated_text, + task=prompt_task, + image_size=(image.width, image.height) + )[prompt_task] + + caption = caption.strip() + if not caption.endswith("."): + last_period_index = caption.rfind(".") + if last_period_index != -1: + caption = caption[:last_period_index + 1] + + torch.cuda.synchronize() + end_time = time.time() + latency = end_time - start_time + mem_gb = torch.cuda.max_memory_allocated() / (1024**3) + + print("-" * 30) + print(f"📝 Result: {caption}") + print(f"📏 Tokens: ~{len(caption.split())} words") # Approx count + print("-" * 30) + print(f"⚡ Latency: {latency:.4f}s") + print(f"💾 VRAM: {mem_gb:.2f} GB") + print("-" * 30) + + return caption + +text_output = generate_concise_caption("/kaggle/input/test1/tree_sktech.jpg") + diff --git a/models/lighting.py b/models/lighting.py @@ -0,0 +1,319 @@ +# -*- coding: utf-8 -*- +"""lighting (1).ipynb + +Automatically generated by Colab. + +Original file is located at + https://colab.research.google.com/drive/1edtbuTtWprRBw_lOvGJuEEDu5HGc-KA3 +""" + +!pip install git+https://github.com/openai/CLIP.git --q +!pip install torch torchvision pillow --q + +import torch +import clip +from PIL import Image +import numpy as np + +LIGHTING_CLASSES = [ + "soft-light", + "specular-highlights", + "backlit", + "studio-lighting", + "flat-lighting", + "dramatic-contrast", + "diffused", +] + +LIGHTING_PHRASES = { + "soft-light": "soft light", + "specular-highlights": "strong specular highlights", + "backlit": "backlit lighting from behind the subject", + "studio-lighting": "professional studio lighting", + "flat-lighting": "flat, low-contrast lighting", + "dramatic-contrast": "dramatic high-contrast lighting", + "diffused": "soft diffused lighting with low shadows", +} + +PROMPT_TEMPLATES = [ + "a photo with {}", + "a portrait lit with {}", + "a high quality photograph with {}", + "a cinematic shot using {}", + "an image captured in {}", + "a studio photo with {}", + "a product photo taken with {}", + "a landscape scene under {}", + "a close-up shot with {}", + "a professional photography setup using {}", +] + +import torch +import clip +import numpy as np + +device = "cuda" if torch.cuda.is_available() else "cpu" +model, preprocess = clip.load("ViT-B/32", device=device) + +def build_text_bank(output_path="lighting_text_bank.npz"): + all_prompt_texts = [] + all_prompt_labels = [] + all_prompt_features = [] + + class_centroids = [] + + with torch.no_grad(): + for lighting in LIGHTING_CLASSES: + phrase = LIGHTING_PHRASES[lighting] + + prompts = [tpl.format(phrase) for tpl in PROMPT_TEMPLATES] + + tokens = clip.tokenize(prompts).to(device) + text_features = model.encode_text(tokens) + text_features = text_features / text_features.norm(dim=-1, keepdim=True) + + all_prompt_texts.extend(prompts) + all_prompt_labels.extend([lighting] * len(prompts)) + all_prompt_features.append(text_features.cpu().numpy()) + + centroid = text_features.mean(dim=0) + centroid = centroid / centroid.norm() + class_centroids.append(centroid.cpu().numpy()) + + all_prompt_features = np.concatenate(all_prompt_features, axis=0); + class_centroids = np.stack(class_centroids, axis=0); + + np.savez( + output_path, + prompt_texts=np.array(all_prompt_texts), + prompt_labels=np.array(all_prompt_labels), + prompt_features=all_prompt_features, + class_centroids=class_centroids, + class_names=np.array(LIGHTING_CLASSES), + ) + print(f"Saved text bank to {output_path}") + + +if __name__ == "__main__": + build_text_bank() + +import torch +import clip +from PIL import Image + +device = "cuda" if torch.cuda.is_available() else "cpu" +model, preprocess = clip.load("ViT-B/32", device=device) + +def get_image_embedding(image_path: str) -> torch.Tensor: + image = Image.open(image_path).convert("RGB") + img_tensor = preprocess(image).unsqueeze(0).to(device) + + with torch.no_grad(): + img_feat = model.encode_image(img_tensor) + img_feat = img_feat / img_feat.norm(dim=-1, keepdim=True) + + return img_feat.squeeze(0) + +import numpy as np +import torch + +import clip +from PIL import Image + +def get_image_embedding(image_path: str) -> torch.Tensor: + image = Image.open(image_path).convert("RGB") + img_tensor = preprocess(image).unsqueeze(0).to(device) + + with torch.no_grad(): + img_feat = model.encode_image(img_tensor) + img_feat = img_feat / img_feat.norm(dim=-1, keepdim=True) + + return img_feat.squeeze(0) + +def load_text_bank(path="lighting_text_bank.npz"): + data = np.load(path, allow_pickle=True) + centroids = data["class_centroids"] + class_names = list(data["class_names"]) + centroids = centroids / np.linalg.norm(centroids, axis=1, keepdims=True) + return centroids, class_names + +def classify_lighting(image_path: str, + text_bank_path="lighting_text_bank.npz", + threshold: float = 0.25): + centroids, class_names = load_text_bank(text_bank_path) + + img_feat = get_image_embedding(image_path) + img_np = img_feat.cpu().numpy()[None, :] + + sims = img_np @ centroids.T + sims = sims.squeeze(0) + + sorted_indices = np.argsort(sims)[::-1] + + top_n = 3 + results = [] + for i in range(min(top_n, len(class_names))): + idx = sorted_indices[i] + lighting_style = class_names[idx] + similarity = float(sims[idx]) + results.append({"lighting_style": lighting_style, "similarity": similarity}) + + return results + +if __name__ == "__main__": + img_path = "/content/Screenshot 2025-11-24 194825.png" + result = classify_lighting(img_path) + print(result) + +!pip install ftfy regex tqdm + +import torch +import clip +from PIL import Image +import numpy as np + +device = "cuda" if torch.cuda.is_available() else "cpu" +model, preprocess = clip.load("ViT-B/32", device=device) + +print("Device:", device) + +LIGHTING_CLASSES = [ + "soft-light", + "specular-highlights", + "backlit", + "studio-lighting", + "flat-lighting", + "dramatic-contrast", + "diffused", +] + + +LIGHTING_PHRASES = { + "soft-light": "soft light", + "specular-highlights": "strong specular highlights", + "backlit": "backlit lighting from behind the subject", + "studio-lighting": "professional studio lighting", + "flat-lighting": "flat, low-contrast lighting", + "dramatic-contrast": "dramatic high-contrast lighting", + "diffused": "soft diffused lighting with low shadows", +} + + +PROMPT_TEMPLATES = [ + "a photo with {}", + "a portrait lit with {}", + "a high quality photograph with {}", + "a cinematic shot using {}", + "an image captured in {}", + "a studio photo with {}", + "a product photo taken with {}", + "a landscape scene under {}", + "a close-up shot with {}", + "a professional photography setup using {}", +] + +def build_text_bank(output_path="lighting_text_bank.npz"): + all_prompt_texts = [] + all_prompt_labels = [] + all_prompt_features = [] + + class_centroids = [] + + model.eval() + with torch.no_grad(): + for lighting in LIGHTING_CLASSES: + phrase = LIGHTING_PHRASES[lighting] + + prompts = [tpl.format(phrase) for tpl in PROMPT_TEMPLATES] + tokens = clip.tokenize(prompts).to(device) + + text_features = model.encode_text(tokens) + text_features = text_features / text_features.norm(dim=-1, keepdim=True) + + all_prompt_texts.extend(prompts) + all_prompt_labels.extend([lighting] * len(prompts)) + all_prompt_features.append(text_features.cpu().numpy()) + + centroid = text_features.mean(dim=0) + centroid = centroid / centroid.norm() + class_centroids.append(centroid.cpu().numpy()) + + all_prompt_features = np.concatenate(all_prompt_features, axis=0) + class_centroids = np.stack(class_centroids, axis=0) + + np.savez( + output_path, + prompt_texts=np.array(all_prompt_texts), + prompt_labels=np.array(all_prompt_labels), + prompt_features=all_prompt_features, + class_centroids=class_centroids, + class_names=np.array(LIGHTING_CLASSES), + ) + print(f"Saved text bank → {output_path}") + + +build_text_bank() + +def get_image_embedding(image_path: str) -> torch.Tensor: + image = Image.open(image_path).convert("RGB") + img_tensor = preprocess(image).unsqueeze(0).to(device) + + model.eval() + with torch.no_grad(): + img_feat = model.encode_image(img_tensor) + img_feat = img_feat / img_feat.norm(dim=-1, keepdim=True) + + return img_feat.squeeze(0) + +DEFAULT_THRESHOLD = 0.28 + + +def load_text_bank(path="lighting_text_bank.npz"): + bank = np.load(path, allow_pickle=True) + centroids = torch.tensor(bank["class_centroids"]).to(device) + class_names = bank["class_names"] + + centroids = centroids / centroids.norm(dim=-1, keepdim=True) + return centroids, class_names + + +def classify_lighting(image_path, threshold=DEFAULT_THRESHOLD): + img_feat = get_image_embedding(image_path).to(device) + + centroids, class_names = load_text_bank() + + img_feat = img_feat.to(centroids.dtype) + + img_feat = img_feat / img_feat.norm() + + sims = centroids @ img_feat + best_score, best_idx = torch.max(sims, dim=0) + + best_score = float(best_score.detach().cpu()) + sims_list = sims.detach().cpu().tolist() + + sims_dict = { + str(cls): float(s) + for cls, s in zip(class_names, sims_list) + } + + if best_score < threshold: + predicted = "no_specific_lighting" + else: + predicted = str(class_names[best_idx]) + + return predicted, best_score, sims_dict + +from google.colab import files +uploaded = files.upload() + +image_path = list(uploaded.keys())[0] + +pred, score, sims = classify_lighting(image_path) + +print("Prediction:", pred) +print("Best score:", score) +print("\nAll similarities:") +for k, v in sims.items(): + print(f"{k:20s}: {v:.3f}") + diff --git a/models/sketch2img_local_server.py b/models/sketch2img_local_server.py @@ -0,0 +1,61 @@ +# -*- coding: utf-8 -*- +"""sketch2img_local_server (1).ipynb + +Automatically generated by Colab. + +Original file is located at + https://colab.research.google.com/drive/1rE5rC71QRvdyxlhnjZUZ320xOV0xrdgx +""" + +!pip install diffusers==0.30.0 transformers accelerate safetensors xformers torch --upgrade + +from diffusers import StableDiffusionPipeline +import torch +from PIL import Image + +model_id = "runwayml/stable-diffusion-v1-5" + +pipe = StableDiffusionPipeline.from_pretrained( + model_id, + torch_dtype=torch.float16, + use_safetensors=True +).to("cuda") + +pipe.enable_xformers_memory_efficient_attention() +pipe.enable_attention_slicing() + +import time +from PIL import Image +import subprocess + +prompt = "A luxury wristwatch placed on a black velvet pedestal, enclosed within a subtle metallic frame. Low-key dramatic lighting with controlled highlights. High-end jewelry product atmosphere. Deep neutral palette with soft contrast transitions." +negative_prompt = "" + +height, width = 512, 512 +steps = 30 +cfg = 7 + +torch.cuda.empty_cache() +start = time.time() + +image = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + guidance_scale=cfg, + num_inference_steps=steps +).images[0] + +end = time.time() + +image.save("output.png") + +print(f"⏱️ Generation Time: {end - start:.2f} seconds") + +image + +!nvidia-smi + +gpu_usage = subprocess.getoutput('nvidia-smi').split('\n') +print("\n".join(gpu_usage[:10])) +\ No newline at end of file diff --git a/models/style.py b/models/style.py @@ -0,0 +1,382 @@ +# -*- coding: utf-8 -*- +"""style.ipynb + +Automatically generated by Colab. + +Original file is located at + https://colab.research.google.com/drive/1YFqRnpIrbZiN2fwOVcr3FUMdHS3WZdwf +""" + +!pip install git+https://github.com/openai/CLIP.git --q +!pip install torch torchvision ftfy regex tqdm --q + +import torch +import clip +from PIL import Image +import numpy as np +from google.colab import files +import io + +device = "cuda" if torch.cuda.is_available() else "cpu" +print(f"Using device: {device}") + +model, preprocess = clip.load("ViT-B/32", device=device) + +STYLE_CLASSES = [ + # --- Pop / Comic / Print --- + "pop-art", + "halftone-print", + "comic-book", + "graphic-novel", + "newspaper-comic-strip", + + # --- Cartoon / Animation --- + "cartoon-style", + "minimal-flat-cartoon", + "rubber-hose", + "retro-cartoon", + "anime", + "chibi", + "manga-inking", + "ghibli-style", + "superhero-comic", + "disney-style", + + # --- Line-Based --- + "line-art", + "continuous-line", + "blind-contour", + "cross-hatching", + "stippling", + "ink-drawing", + "technical-blueprint", + "engraving", + "monoline-icons", + + # --- Traditional Media --- + "watercolor", + "gouache", + "oil-painting", + "acrylic-paint", + "pastel-chalk", + "oil-pastel", + "marker-render", + "crayon", + "ink-wash", + "colored-pencil", + "mixed-media", + + # --- Realism / Naturalism --- + "realism", + "naturalism", + "photorealism", + "hyperrealism", + "portrait-realism", + "botanical-illustration", + "wildlife-realism", + "scientific-illustration", + "documentary-illustration", + + # --- Abstract / Modernist --- + "abstract-expressionism", + "minimal-abstract", + "geometric-abstraction", + "color-field", + "cubism", + "futurism", + "bauhaus", + "constructivism", + "surreal", + "magical-realism", + "op-art", + "escher-impossible", + + # --- Collage / Cut --- + "paper-cutout", + "analog-collage", + "digital-collage", + "photomontage", + + # --- Historical / Cultural --- + "ukiyo-e", + "indian-miniature", + "persian-miniature", + "art-nouveau", + "art-deco", + "medieval-illustration", + "byzantine-icon", + + # --- Retro / Vintage --- + "vintage", + "mid-century-modern", + "psychedelic-60s", + "groovy-70s", + "vaporwave", + "synthwave", + "retro-futurism", + "risograph", + "screenprint", + + # --- Street / Subculture --- + "graffiti", + "street-art", + "lowbrow-pop", + "sticker-bomb", + "zine-grunge", + + # --- Digital / 3D / Tech --- + "pixel-art", + "isometric-pixel", + "voxel-art", + "low-poly-3d", + "hyperreal-3d", + "clay-render", + "cel-shaded-3d", + "holographic-gradient", + "glitch-art", + "chromatic-aberration", + "datamosh", + "vhs-analog" +] + + +STYLE_DESCRIPTIONS = { + + # --- Pop / Comic / Print --- + "pop-art": "pop art, comic book style, bold colors, halftone dots, thick outlines", + "halftone-print": "halftone dot shading, CMYK print texture, vintage comic printing", + "comic-book": "dynamic panels, bold inking, action lines, superhero comic style", + "graphic-novel": "moody shadows, dramatic black ink, narrative panels, serious tone", + "newspaper-comic-strip": "simple line work, flat shading, classic newspaper strip style", + + # --- Cartoon / Animation --- + "cartoon-style": "bold outlines, flat colors, exaggerated characters and expressions", + "minimal-flat-cartoon": "simple vector shapes, flat color fills, minimal details", + "rubber-hose": "1930s animation, noodle limbs, vintage rubber hose characters", + "retro-cartoon": "old-school grain, hand-drawn frames, muted vintage palette", + "anime": "clean linework, cel-shading, large expressive eyes, Japanese animation", + "chibi": "super-deformed characters, big heads, tiny bodies, cute aesthetic", + "manga-inking": "black and white screentones, clean ink lines, manga style", + "ghibli-style": "soft painterly shading, warm palettes, whimsical nature scenes", + "superhero-comic": "muscular characters, dramatic shading, intense action poses", + "disney-style": "smooth polished shading, expressive eyes, classic animation aesthetic", + + # --- Line-Based --- + "line-art": "clean linear outlines, no shading, monochrome strokes", + "continuous-line": "one unbroken flowing line forming the subject", + "blind-contour": "distorted outlines, drawn without looking, loose sketch style", + "cross-hatching": "intersecting lines for shading, dense ink texture", + "stippling": "dot-based shading, gradual tonal gradients from dots", + "ink-drawing": "strong black ink strokes, high contrast, brush or pen feel", + "technical-blueprint": "precise technical diagrams, blue background, white lines", + "engraving": "fine parallel carved lines, antique print look", + "monoline-icons": "uniform line weight, minimal shapes, iconographic design", + + # --- Traditional Media --- + "watercolor": "transparent washes, soft gradients, bleeding pigment, paper texture", + "gouache": "opaque matte colors, flat coverage, velvety paint texture", + "oil-painting": "rich color blends, heavy brushstrokes, textured canvas feel", + "acrylic-paint": "opaque fast-drying strokes, strong colors, painterly texture", + "pastel-chalk": "soft dusty texture, delicate blending, chalky edges", + "oil-pastel": "waxy thick strokes, vibrant smears, layered pigments", + "marker-render": "streaky marker fills, smooth gradients, alcohol ink texture", + "crayon": "childlike wax texture, rough strokes, uneven coloring", + "ink-wash": "monochrome fluid washes, sumi-e brush style, soft gradients", + "colored-pencil": "fine pencil texture, layered shading, visible grain", + "mixed-media": "combined materials, layered textures, hybrid collage effects", + + # --- Realism / Naturalism --- + "realism": "accurate representation, natural lighting, lifelike detail", + "naturalism": "organic colors, soft shading, true-to-life textures", + "photorealism": "photo-like accuracy, crisp details, polished lighting", + "hyperrealism": "ultra-detailed textures, exaggerated clarity, surreal realism", + "portrait-realism": "true facial anatomy, realistic skin tones, fine shading", + "botanical-illustration": "high-detail plants, scientific precision, fine lines", + "wildlife-realism": "accurate animal anatomy, natural fur/feather rendering", + "scientific-illustration": "precise labeled diagrams, technical accuracy", + "documentary-illustration": "real-life scenes, muted tones, reportage style", + + # --- Abstract / Modernist --- + "abstract-expressionism": "gestural strokes, energetic splashes, emotional color", + "minimal-abstract": "simple geometric forms, large negative space, clarity", + "geometric-abstraction": "strict shapes, patterns, bold geometry", + "color-field": "large flat color areas, subtle tonal transitions", + "cubism": "fragmented shapes, multiple viewpoints, angular abstraction", + "futurism": "dynamic angles, motion streaks, mechanical aesthetic", + "bauhaus": "primary colors, clean geometry, functional composition", + "constructivism": "bold diagonals, industrial forms, propaganda flavor", + "surreal": "dreamlike imagery, strange juxtapositions, impossible forms", + "magical-realism": "ordinary scenes with subtle fantastical elements", + "op-art": "optical illusions, vibrating lines, high-contrast patterns", + "escher-impossible": "paradoxical architecture, infinite loops, tessellations", + + # --- Collage / Cut --- + "paper-cutout": "flat colored shapes, crisp edges, layered paper look", + "analog-collage": "torn paper edges, magazine scraps, glue texture", + "digital-collage": "layered photography fragments, digital assembly", + "photomontage": "multiple photos combined into surreal compositions", + + # --- Historical / Cultural --- + "ukiyo-e": "japanese woodblock style, flat colors, bold outlines, wave patterns", + "indian-miniature": "fine detailing, ornate patterns, traditional storytelling", + "persian-miniature": "vibrant colors, intricate decoration, ornate scenes", + "art-nouveau": "flowing curves, floral motifs, elegant decorative lines", + "art-deco": "geometric symmetry, metallic accents, luxury aesthetic", + "medieval-illustration": "flat perspective, gold leaf accents, illuminated style", + "byzantine-icon": "gold backgrounds, frontal figures, religious icon style", + + # --- Retro / Vintage --- + "vintage": "aged tones, film grain, retro print texture, nostalgic aesthetic", + "mid-century-modern": "retro shapes, muted tones, playful geometric forms", + "psychedelic-60s": "vivid gradients, swirling patterns, trippy colors", + "groovy-70s": "warm earthy colors, wavy lines, bohemian retro vibe", + "vaporwave": "neon pinks and blues, retro grids, greek busts, glitchy nostalgia", + "synthwave": "neon purples, chrome shine, retrofuturistic sunset grids", + "retro-futurism": "futuristic style imagined in the past, old sci-fi aesthetic", + "risograph": "duotone neon inks, grainy overlays, misaligned color layers", + "screenprint": "flat inks, bold shapes, layered silkscreen texture", + + # --- Street / Subculture --- + "graffiti": "spray paint texture, drips, tags, bold lettering", + "street-art": "urban wall murals, stencil work, mixed street media", + "lowbrow-pop": "quirky surrealism, outsider cartoon style, punk influence", + "sticker-bomb": "dense overlapping stickers, colorful chaotic layers", + "zine-grunge": "photocopy texture, rough cut edges, DIY punk style", + + # --- Digital / 3D / Tech --- + "pixel-art": "pixelated blocks, low resolution, retro game style", + "isometric-pixel": "45-degree pixel perspective, 3D pixel illusion", + "voxel-art": "3D cubes, blocky shapes, volumetric pixel aesthetic", + "low-poly-3d": "simple polygon shapes, faceted surfaces, flat shading", + "hyperreal-3d": "lifelike 3D rendering, photoreal materials and lighting", + "clay-render": "soft clay surface, sculpted feel, neutral lighting", + "cel-shaded-3d": "toon outlines, flat shading, stylized animation look", + "holographic-gradient": "iridescent metallic rainbow gradients, glossy shine", + "glitch-art": "broken pixels, digital errors, RGB misalignment", + "chromatic-aberration": "red-blue fringing, lens distortion edges", + "datamosh": "compression smears, frame artifacts, melting visuals", + "vhs-analog": "scanlines, tape noise, retro analog glitch" +} + + +PROMPT_TEMPLATES = [ + "a photo in the style of {}", + "artistic rendering of {}", + "a {} style image", + "a painting in the style of {}", + "a high quality example of {} art", + "visuals depicting {}", + "artwork created using {}", + "a poster with {}", + "a photograph looking like {}", + "an illustration inspired by {}", + "a digital artwork following the rules of {}", + "a creative piece showcasing {}", + "a visual composition designed in {} style", + "an artistic composition influenced by {}", + "a scene presented with {} aesthetics", + "an image heavily influenced by {} techniques", + "a graphical representation in {} aesthetics", + "a stylized depiction matching {}", + "the visual language of {} applied to an artwork", + "a concept image using {} characteristics", + "a stylized portrait in the manner of {}", + "a landscape interpreted through {} style", + "a figure drawing executed in {} technique", + "a creative poster influenced by {}", + "a concept sketch using {} visual rules", + "an editorial illustration drawn in {} style", + "a rendered scene with {} visual identity", + "an artistic study done with {} aesthetics", + "an experiment in {} visual style", + "a fully rendered image designed around {} style" +] + + +class StyleClassifier: + def __init__(self): + self.centroids = None + self.class_names = STYLE_CLASSES + self.build_text_features() + + def build_text_features(self): + """ + Pre-computes the embedding centroids for all styles. + """ + print("Building text embeddings for styles...") + class_centroids = [] + + with torch.no_grad(): + for style in self.class_names: + desc = STYLE_DESCRIPTIONS[style] + + prompts = [tpl.format(desc) for tpl in PROMPT_TEMPLATES] + + tokens = clip.tokenize(prompts).to(device) + + text_features = model.encode_text(tokens) + + text_features = text_features / text_features.norm(dim=-1, keepdim=True) + + centroid = text_features.mean(dim=0) + + centroid = centroid / centroid.norm() + + class_centroids.append(centroid) + + self.centroids = torch.stack(class_centroids).to(device) + print("Text embeddings built successfully.") + + def predict(self, image_path, threshold=0.22): + try: + image = Image.open(image_path).convert("RGB") + except Exception as e: + print(f"Error loading image: {e}") + return "error", [] + + img_tensor = preprocess(image).unsqueeze(0).to(device) + + with torch.no_grad(): + img_feat = model.encode_image(img_tensor) + img_feat = img_feat / img_feat.norm(dim=-1, keepdim=True) + + similarity = (100.0 * img_feat @ self.centroids.T) + + probs = similarity.softmax(dim=-1).cpu().numpy()[0] + raw_scores = similarity.cpu().numpy()[0] / 100.0 + + results = [] + for i, style in enumerate(self.class_names): + results.append({ + "style": style, + "score": float(raw_scores[i]), + "confidence_pct": float(probs[i] * 100) + }) + + results.sort(key=lambda x: x["score"], reverse=True) + + top_result = results[0] + + if top_result["score"] < threshold: + prediction = "undefined/mixed" + else: + prediction = top_result["style"] + + return prediction, results + +# EXECUTION + +classifier = StyleClassifier() + +print("\nPlease upload an image to analyze:") +uploaded = files.upload() + +for filename in uploaded.keys(): + print(f"\nAnalyzing: {filename}...") + prediction, details = classifier.predict(filename) + + print(f"TOP PREDICTION: {prediction.upper()}") + print("-" * 30) + print("Full Breakdown:") + for res in details[:5]: + print(f"{res['style']:<20} | Score: {res['score']:.4f} | Conf: {res['confidence_pct']:.2f}%") +\ No newline at end of file diff --git a/models/stylesheet_gen&prompt_gen.py b/models/stylesheet_gen&prompt_gen.py @@ -0,0 +1,671 @@ +# -*- coding: utf-8 -*- +"""Untitled6.ipynb + +Automatically generated by Colab. + +Original file is located at + https://colab.research.google.com/drive/1eZLWqhB0KKFLhCA0E3qDHXpdvEIH44yE +""" + +import json +import math +import re +import numpy as np +from collections import defaultdict +from typing import List, Dict, Tuple, Any, Optional +from sklearn.cluster import KMeans + +RANK_DECAY_FACTOR = 0.85 +MIN_CONSENSUS_THRESHOLD = 0.05 +PALETTE_SIZE = 5 + + +def hex_to_rgb(hex_str: str) -> List[int]: + """Converts '#FF5733' to [255, 87, 51] for math operations.""" + try: + hex_str = hex_str.strip().lstrip('#') + if len(hex_str) != 6: return [0, 0, 0] + return [int(hex_str[i:i+2], 16) for i in (0, 2, 4)] + except: + return [0, 0, 0] + +def rgb_to_hex(rgb: List[int]) -> str: + """Converts [255, 87, 51] back to '#FF5733'.""" + return '#{:02x}{:02x}{:02x}'.format(int(rgb[0]), int(rgb[1]), int(rgb[2])) + + +class FontClusterEngine: + def __init__(self): + self.full_font_map = { + "abeezee": "Sans Serif (Generic)", "abel": "Condensed Sans", "abhayalibre": "Serif (Generic)", + "abrilfatface": "Modern Serif", "abyssinicasil": "Serif (Generic)", "aclonica": "Display/Other", + "acme": "Display/Other", "actor": "Sans Serif (Generic)", "adamina": "Serif (Generic)", + "adventpro": "Display/Other", "aguafinascript": "Script/Handwritten", "akronim": "Display/Other", + "aladin": "Display/Other", "aldrich": "Display/Other", "alef": "Sans Serif (Generic)", + "alegreya": "Old Style Serif", "alegreyasans": "Humanist Sans", "alegreyasanssc": "Sans Serif (Generic)", + "alegreyasc": "Old Style Serif", "alexbrush": "Script/Handwritten", "alfaslabone": "Slab Serif", + "alice": "Serif (Generic)", "alike": "Serif (Generic)", "alikeangular": "Serif (Generic)", + "allan": "Display/Other", "allerta": "Sans Serif (Generic)", "allertastencil": "Display/Other", + "allura": "Script/Handwritten", "almendra": "Display/Other", "almendradisplay": "Display/Other", + "almendrasc": "Display/Other", "amarante": "Display/Other", "amaranth": "Display/Other", + "amaticasc": "Script/Handwritten", "amaticsc": "Script/Handwritten", "amethysta": "Serif (Generic)", + "amiko": "Sans Serif (Generic)", "amiri": "Serif (Generic)", "amita": "Script/Handwritten", + "anaheim": "Sans Serif (Generic)", "andada": "Serif (Generic)", "andadasc": "Serif (Generic)", + "andika": "Sans Serif (Generic)", "annieuseyourtelescope": "Script/Handwritten", "anonymouspro": "Monospace", + "antic": "Sans Serif (Generic)", "anticdidone": "Serif (Generic)", "anticslab": "Slab Serif", + "anton": "Grotesque Sans", "antonio": "Sans Serif (Generic)", "arapey": "Serif (Generic)", + "arbutus": "Display/Other", "arbutusslab": "Slab Serif", "architectsdaughter": "Script/Handwritten", + "archivo": "Grotesque Sans", "archivoblack": "Grotesque Sans", "archivonarrow": "Condensed Sans", + "archivovfbeta": "Grotesque Sans", "arefruqaa": "Serif (Generic)", "arimamadurai": "Display/Other", + "arimo": "Neo-Grotesque Sans", "arizonia": "Script/Handwritten", "armata": "Sans Serif (Generic)", + "arsenal": "Sans Serif (Generic)", "artifika": "Serif (Generic)", "arvo": "Slab Serif", + "arya": "Sans Serif (Generic)", "asap": "Rounded Sans", "asapcondensed": "Condensed Sans", + "asapvfbeta": "Rounded Sans", "asar": "Serif (Generic)", "asset": "Display/Other", + "assistant": "Sans Serif (Generic)", "astloch": "Display/Other", "asul": "Sans Serif (Generic)", + "athiti": "Sans Serif (Generic)", "atma": "Display/Other", "atomicage": "Display/Other", + "aubrey": "Display/Other", "audiowide": "Display/Other", "autourone": "Display/Other", + "average": "Serif (Generic)", "averagesans": "Sans Serif (Generic)", "averiagruesalibre": "Display/Other", + "averialibre": "Display/Other", "averiasanslibre": "Sans Serif (Generic)", "averiaseriflibre": "Serif (Generic)", + "badscript": "Script/Handwritten", "bahiana": "Display/Other", "baloo": "Display/Other", + "baloobhai": "Display/Other", "baloobhaijaan": "Display/Other", "baloobhaina": "Display/Other", + "baloochettan": "Display/Other", "balooda": "Display/Other", "baloopaaji": "Display/Other", + "balootamma": "Display/Other", "balootammudu": "Display/Other", "baloothambi": "Display/Other", + "balthazar": "Serif (Generic)", "bangers": "Display", "barlow": "Grotesque Sans", + "barlowcondensed": "Condensed Sans", "barlowsemicondensed": "Condensed Sans", "barrio": "Display/Other", + "basic": "Sans Serif (Generic)", "baumans": "Display/Other", "belgrano": "Serif (Generic)", + "bellefair": "Serif (Generic)", "belleza": "Sans Serif (Generic)", "benchnine": "Sans Serif (Generic)", + "bentham": "Serif (Generic)", "berkshireswash": "Script/Handwritten", "bevan": "Slab Serif", + "bhavuka": "Display/Other", "bigelowrules": "Display/Other", "bigshotone": "Display/Other", + "bilbo": "Script/Handwritten", "bilboswashcaps": "Script/Handwritten", "biorhyme": "Slab Serif", + "biorhymeexpanded": "Slab Serif", "biryani": "Sans Serif (Generic)", "bitter": "Slab Serif", + "blackandwhitepicture": "Display/Other", "blackhansans": "Sans Serif (Generic)", "blackopsone": "Display/Other", + "bonbon": "Display/Other", "boogaloo": "Display/Other", "bowlbyone": "Display/Other", + "bowlbyonesc": "Display/Other", "brawler": "Serif (Generic)", "breeserif": "Slab Serif", + "brunoace": "Display/Other", "brunoacesc": "Display/Other", "bubblegumsans": "Display/Other", + "bubblerone": "Display/Other", "buenard": "Serif (Generic)", "bungee": "Display/Other", + "bungeehairline": "Display/Other", "bungeeinline": "Display/Other", "bungeeoutline": "Display/Other", + "bungeeshade": "Display/Other", "butcherman": "Display/Other", "butchermancaps": "Display/Other", + "butterflykids": "Script/Handwritten", "cabin": "Humanist Sans", "cabincondensed": "Condensed Sans", + "cabinsketch": "Display/Other", "cabinvfbeta": "Humanist Sans", "caesardressing": "Display/Other", + "cagliostro": "Sans Serif (Generic)", "cairo": "Grotesque Sans", "calligraffitti": "Script/Handwritten", + "cambay": "Sans Serif (Generic)", "cambo": "Serif (Generic)", "cantarell": "Sans Serif (Generic)", + "cantataone": "Serif (Generic)", "cantoraone": "Sans Serif (Generic)", "capriola": "Sans Serif (Generic)", + "cardo": "Serif (Generic)", "carme": "Sans Serif (Generic)", "carroisgothic": "Sans Serif (Generic)", + "carroisgothicsc": "Sans Serif (Generic)", "catamaran": "Sans Serif", "caudex": "Serif (Generic)", + "caveat": "Script/Handwritten", "caveatbrush": "Script/Handwritten", "cevicheone": "Display/Other", + "changa": "Sans Serif (Generic)", "changaone": "Display/Other", "chango": "Display/Other", + "chathura": "Sans Serif (Generic)", "chauphilomeneone": "Display/Other", "chelaone": "Display/Other", + "chelseamarket": "Display/Other", "cherrycreamsoda": "Display/Other", "cherryswash": "Display/Other", + "chewy": "Display", "chicle": "Display/Other", "chivo": "Grotesque Sans", + "chonburi": "Display/Other", "cinzel": "Serif", "cinzeldecorative": "Display/Other", + "clickerscript": "Script/Handwritten", "coda": "Display/Other", "codystar": "Display/Other", + "coiny": "Display/Other", "combo": "Display/Other", "comfortaa": "Rounded Sans", + "comingsoon": "Script/Handwritten", "concertone": "Display/Other", "condiment": "Script/Handwritten", + "contrailone": "Display/Other", "convergence": "Sans Serif (Generic)", "cookie": "Script", + "copse": "Slab Serif", "cormorant": "Old Style Serif", "cormorantgaramond": "Old Style Serif", + "cormorantinfant": "Old Style Serif", "cormorantsc": "Old Style Serif", "cormorantunicase": "Old Style Serif", + "cormorantupright": "Old Style Serif", "courgette": "Script", "cousine": "Monospace", + "coustard": "Slab Serif", "craftygirls": "Script/Handwritten", "creepster": "Display/Other", + "creepstercaps": "Display/Other", "creteround": "Slab Serif", "crimsontext": "Old Style Serif", + "croissantone": "Display/Other", "crushed": "Display/Other", "cuprum": "Sans Serif", + "cutefont": "Display/Other", "cutive": "Serif (Generic)", "cutivemono": "Monospace", + "damion": "Script/Handwritten", "dancingscript": "Script/Handwritten", "daysone": "Display/Other", + "decovaralpha": "Display/Other", "dekko": "Script/Handwritten", "delius": "Script/Handwritten", + "deliusswashcaps": "Script/Handwritten", "deliusunicase": "Script/Handwritten", "dellarespira": "Serif (Generic)", + "denkone": "Display/Other", "devonshire": "Script/Handwritten", "dhurjati": "Display/Other", + "dhyana": "Display/Other", "didactgothic": "Sans Serif (Generic)", "digitalnumbers": "Display/Other", + "diplomata": "Display/Other", "diplomatasc": "Display/Other", "dohyeon": "Sans Serif (Generic)", + "dokdo": "Script/Handwritten", "domine": "Serif", "donegalone": "Display/Other", + "doppioone": "Sans Serif (Generic)", "dorsa": "Display/Other", "dosis": "Rounded Sans", + "drsugiyama": "Script/Handwritten", "durusans": "Sans Serif (Generic)", "dynalight": "Script/Handwritten", + "eaglelake": "Script/Handwritten", "eastseadokdo": "Script/Handwritten", "eater": "Display/Other", + "eatercaps": "Display/Other", "ebgaramond": "Old Style Serif", "economica": "Sans Serif (Generic)", + "eczar": "Display/Other", "ekmukta": "Sans Serif (Generic)", "electrolize": "Display/Other", + "elmessiri": "Sans Serif (Generic)", "elsie": "Display/Other", "elsieswashcaps": "Display/Other", + "emblemaone": "Display/Other", "emilyscandy": "Display/Other", "encodesans": "Sans Serif", + "encodesanscondensed": "Condensed Sans", "encodesansexpanded": "Sans Serif (Generic)", "encodesanssemicondensed": "Condensed Sans", + "encodesanssemiexpanded": "Sans Serif (Generic)", "engagement": "Script/Handwritten", "englebert": "Display/Other", + "enriqueta": "Slab Serif", "ericaone": "Display/Other", "esteban": "Serif (Generic)", + "euphoriascript": "Script/Handwritten", "ewert": "Display/Other", "exo": "Geometric Sans", + "exo2": "Geometric Sans", "expletussans": "Display/Other", "fanwoodtext": "Serif (Generic)", + "farsan": "Display/Other", "fascinate": "Display/Other", "fascinateinline": "Display/Other", + "fasterone": "Display/Other", "fasthand": "Script/Handwritten", "faunaone": "Serif (Generic)", + "faustina": "Serif (Generic)", "faustinavfbeta": "Serif (Generic)", "federant": "Display/Other", + "federo": "Display/Other", "felipa": "Script/Handwritten", "fenix": "Serif (Generic)", + "fingerpaint": "Display/Other", "firasans": "Humanist Sans", "firasanscondensed": "Condensed Sans", + "firasansextracondensed": "Condensed Sans", "fjallaone": "Grotesque Sans", "fjordone": "Serif (Generic)", + "flamenco": "Display/Other", "flavors": "Display/Other", "fondamento": "Script/Handwritten", + "fontdinerswanky": "Display/Other", "forum": "Display/Other", "francoisone": "Sans Serif", + "frankruhllibre": "Serif (Generic)", "freckleface": "Display/Other", "frederickathegreat": "Display/Other", + "fredokaone": "Rounded Display", "fresca": "Sans Serif (Generic)", "frijole": "Display/Other", + "fruktur": "Display/Other", "fugazone": "Display/Other", "gabriela": "Serif (Generic)", + "gaegu": "Script/Handwritten", "gafata": "Sans Serif (Generic)", "galada": "Display/Other", + "galdeano": "Sans Serif (Generic)", "galindo": "Display/Other", "gamjaflower": "Script/Handwritten", + "gemunulibre": "Sans Serif (Generic)", "genbasb": "Serif (Generic)", "genbasbi": "Serif (Generic)", + "genbkbasb": "Serif (Generic)", "genbkbasbi": "Serif (Generic)", "geo": "Display/Other", + "geostar": "Display/Other", "geostarfill": "Display/Other", "germaniaone": "Display/Other", + "gfsdidot": "Serif (Generic)", "gfsneohellenic": "Sans Serif (Generic)", "gfsneohellenicbold": "Sans Serif (Generic)", + "gfsneohellenicbolditalic": "Sans Serif (Generic)", "gidugu": "Display/Other", "gildadisplay": "Serif (Generic)", + "glassantiqua": "Display/Other", "glegoo": "Slab Serif", "gochihand": "Script/Handwritten", + "gorditas": "Display/Other", "gothica1": "Sans Serif (Generic)", "graduate": "Display/Other", + "grandhotel": "Script/Handwritten", "greatvibes": "Script", "griffy": "Display/Other", + "gruppo": "Display/Other", "gudea": "Sans Serif (Generic)", "gugi": "Display/Other", + "gurajada": "Display/Other", "habibi": "Serif (Generic)", "halant": "Serif (Generic)", + "hammersmithone": "Sans Serif (Generic)", "hanalei": "Display/Other", "hanaleifill": "Display/Other", + "handlee": "Script/Handwritten", "hannari": "Display/Other", "hanuman": "Serif (Generic)", + "happymonkey": "Display/Other", "harmattan": "Sans Serif (Generic)", "headlandone": "Serif (Generic)", + "heebo": "Sans Serif", "hennypenny": "Display/Other", "hermeneusone": "Sans Serif (Generic)", + "herrvonmuellerhoff": "Script/Handwritten", "himelody": "Script/Handwritten", "hind": "Humanist Sans", + "hindcolombo": "Sans Serif (Generic)", "hindguntur": "Sans Serif (Generic)", "hindjalandhar": "Sans Serif (Generic)", + "hindkochi": "Sans Serif (Generic)", "hindmadurai": "Sans Serif (Generic)", "hindmysuru": "Sans Serif (Generic)", + "hindsiliguri": "Sans Serif (Generic)", "hindvadodara": "Sans Serif (Generic)", "homemadeapple": "Script/Handwritten", + "homenaje": "Sans Serif (Generic)", "ibmplexmono": "Monospace", "ibmplexsans": "Neo-Grotesque Sans", + "ibmplexsanscondensed": "Condensed Sans", "ibmplexserif": "Transitional Serif", "iceberg": "Display/Other", + "iceland": "Display/Other", "imprima": "Sans Serif (Generic)", "inconsolata": "Monospace", + "inder": "Sans Serif (Generic)", "indieflower": "Handwritten", "inika": "Serif (Generic)", + "inknutantiqua": "Serif (Generic)", "irishgrover": "Display/Other", "istokweb": "Sans Serif (Generic)", + "italiana": "Serif (Generic)", "italianno": "Script/Handwritten", "itim": "Script/Handwritten", + "jacquesfrancois": "Display/Other", "jacquesfrancoisshadow": "Display/Other", "jaldi": "Sans Serif (Generic)", + "jejugothic": "Sans Serif (Generic)", "jejuhallasan": "Display/Other", "jejumyeongjo": "Serif (Generic)", + "jimnightshade": "Display/Other", "jockeyone": "Display/Other", "jollylodger": "Display/Other", + "jomhuria": "Display/Other", "josefinsans": "Geometric Sans", "josefinslab": "Slab Serif", + "jotione": "Display/Other", "jua": "Sans Serif (Generic)", "judson": "Serif (Generic)", + "julee": "Script/Handwritten", "juliussansone": "Sans Serif (Generic)", "junge": "Serif (Generic)", + "jura": "Sans Serif (Generic)", "justanotherhand": "Script/Handwritten", "kadwa": "Serif (Generic)", + "kalam": "Handwritten", "kameron": "Serif (Generic)", "kanit": "Sans Serif", + "kantumruy": "Sans Serif (Generic)", "karla": "Sans Serif (Generic)", "karlatamilinclined": "Sans Serif (Generic)", + "karlatamilupright": "Sans Serif (Generic)", "karma": "Serif (Generic)", "katibeh": "Display/Other", + "kaushanscript": "Script", "kavivanar": "Script/Handwritten", "kavoon": "Display/Other", + "kdamthmor": "Display/Other", "keaniaone": "Display/Other", "kellyslab": "Display/Other", + "kenia": "Display/Other", "khand": "Sans Serif", "khula": "Sans Serif (Generic)", + "khyay": "Display/Other", "kiranghaerang": "Display/Other", "kiteone": "Sans Serif (Generic)", + "knewave": "Display/Other", "kokoro": "Display/Other", "kopubbatang": "Serif (Generic)", + "kottaone": "Serif (Generic)", "kranky": "Display/Other", "kreon": "Slab Serif", + "kristi": "Script/Handwritten", "kronaone": "Sans Serif (Generic)", "kurale": "Serif (Generic)", + "laila": "Serif (Generic)", "lakkireddy": "Display/Other", "lalezar": "Display/Other", + "lancelot": "Display/Other", "laomuangdon": "Display/Other", "laomuangkhong": "Display/Other", + "laosanspro": "Sans Serif (Generic)", "lato": "Humanist Sans", "leckerlione": "Display/Other", + "ledger": "Serif (Generic)", "lekton": "Monospace", "lemon": "Display/Other", + "lemonada": "Display/Other", "lemonadavfbeta": "Display/Other", "librebaskerville": "Serif (Generic)", + "librecaslontext": "Serif (Generic)", "librefranklin": "Neo-Grotesque Sans", "lifesavers": "Display/Other", + "lilitaone": "Display/Other", "lilyscriptone": "Display/Other", "limelight": "Display/Other", + "lindenhill": "Serif (Generic)", "lobster": "Script", "lobstertwo": "Script/Handwritten", + "londrinaoutline": "Display/Other", "londrinashadow": "Display/Other", "londrinasketch": "Display/Other", + "londrinasolid": "Display/Other", "lora": "Old Style Serif", "loversquarrel": "Script/Handwritten", + "loveyalikeasister": "Script/Handwritten", "luckiestguy": "Display/Other", "lusitana": "Serif (Generic)", + "lustria": "Serif (Generic)", "macondo": "Display/Other", "macondoswashcaps": "Display/Other", + "mada": "Sans Serif (Generic)", "magra": "Sans Serif (Generic)", "maidenorange": "Display/Other", + "maitree": "Serif (Generic)", "mako": "Sans Serif (Generic)", "mallanna": "Display/Other", + "mandali": "Sans Serif (Generic)", "manuale": "Serif (Generic)", "marcellus": "Serif (Generic)", + "marcellussc": "Serif (Generic)", "marckscript": "Script/Handwritten", "margarine": "Display/Other", + "markazitext": "Serif (Generic)", "markoone": "Serif (Generic)", "marmelad": "Sans Serif (Generic)", + "martel": "Serif (Generic)", "martelsans": "Sans Serif (Generic)", "marvel": "Sans Serif (Generic)", + "mate": "Serif (Generic)", "matesc": "Serif (Generic)", "mavenpro": "Geometric Sans", + "mavenprovfbeta": "Geometric Sans", "mclaren": "Display/Other", "medulaone": "Display/Other", + "meerainimai": "Display/Other", "meiescript": "Script/Handwritten", "mergeone": "Sans Serif (Generic)", + "merienda": "Display/Other", "meriendaone": "Display/Other", "merriweather": "Transitional Serif", + "merriweathersans": "Humanist Sans", "mervalescript": "Script/Handwritten", "metalmania": "Display/Other", + "metamorphous": "Display/Other", "metrophobic": "Sans Serif (Generic)", "miama": "Script/Handwritten", + "milonga": "Display/Other", "miltonian": "Display/Other", "miltoniantattoo": "Display/Other", + "mina": "Sans Serif (Generic)", "miniver": "Display/Other", "miriamlibre": "Sans Serif (Generic)", + "mirza": "Display/Other", "missfajardose": "Script/Handwritten", "mitr": "Sans Serif (Generic)", + "modak": "Display/Other", "modernantiqua": "Serif (Generic)", "mogra": "Display/Other", + "molengo": "Sans Serif (Generic)", "molle": "Display/Other", "monda": "Sans Serif (Generic)", + "monoton": "Display/Other", "monsieurladoulaise": "Script/Handwritten", "montaga": "Serif (Generic)", + "montez": "Script/Handwritten", "montserrat": "Geometric Sans", "montserratalternates": "Geometric Sans", + "montserratsubrayada": "Geometric Sans", "mountainsofchristmas": "Display/Other", "mousememoirs": "Display/Other", + "mplus1p": "Rounded Sans", "mrbedfort": "Script/Handwritten", "mrdafoe": "Script/Handwritten", + "mrdehaviland": "Script/Handwritten", "mrssaintdelafield": "Script/Handwritten", "mrssheppards": "Script/Handwritten", + "mukta": "Sans Serif", "muktamahee": "Sans Serif (Generic)", "muktamalar": "Sans Serif (Generic)", + "muktavaani": "Sans Serif (Generic)", "muli": "Humanist Sans", "myanmarsanspro": "Sans Serif (Generic)", + "mysteryquest": "Display/Other", "nanumbrushscript": "Script/Handwritten", "nanumgothic": "Sans Serif (Generic)", + "nanumgothiccoding": "Monospace", "nanummyeongjo": "Serif (Generic)", "nanumpenscript": "Script/Handwritten", + "nats": "Display/Other", "neuton": "Serif (Generic)", "newrocker": "Display/Other", + "newscycle": "Sans Serif (Generic)", "nicomoji": "Display/Other", "niconne": "Script/Handwritten", + "nikukyu": "Display/Other", "nixieone": "Display/Other", "nobile": "Sans Serif (Generic)", + "nokora": "Sans Serif (Generic)", "norican": "Script/Handwritten", "nosifer": "Display/Other", + "nosifercaps": "Display/Other", "notable": "Display/Other", "noticiatext": "Serif (Generic)", + "notosans": "Humanist Sans", "notosanstamil": "Humanist Sans", "notoserif": "Transitional Serif", + "novascript": "Display/Other", "ntr": "Sans Serif (Generic)", "numans": "Sans Serif (Generic)", + "nunito": "Rounded Sans", "nunitosans": "Rounded Sans", "offside": "Display/Other", + "oldenburg": "Display/Other", "oldstandard": "Serif (Generic)", "oleoscript": "Script/Handwritten", + "oleoscriptswashcaps": "Script/Handwritten", "opensans": "Humanist Sans", "opensanscondensed": "Condensed Sans", + "opensanshebrew": "Humanist Sans", "opensanshebrewcondensed": "Condensed Sans", "oranienbaum": "Serif (Generic)", + "orbitron": "Display/Other", "oregano": "Script/Handwritten", "orienta": "Sans Serif (Generic)", + "originalsurfer": "Display/Other", "oswald": "Grotesque Sans", "overlock": "Display/Other", + "overlocksc": "Display/Other", "overpass": "Sans Serif (Generic)", "overpassmono": "Monospace", + "ovo": "Serif (Generic)", "oxygen": "Geometric Sans", "oxygenmono": "Monospace", + "pacifico": "Script", "palanquin": "Sans Serif (Generic)", "palanquindark": "Sans Serif (Generic)", + "pangolin": "Script/Handwritten", "paprika": "Display/Other", "parisienne": "Script/Handwritten", + "passeroone": "Display/Other", "passionone": "Display", "pathwaygothicone": "Condensed Sans", + "patrickhand": "Script/Handwritten", "patrickhandsc": "Script/Handwritten", "pattaya": "Display/Other", + "patuaone": "Slab Serif", "pavanam": "Sans Serif (Generic)", "paytoneone": "Sans Serif (Generic)", + "peddana": "Serif (Generic)", "peralta": "Display/Other", "permanentmarker": "Handwritten", + "petitformalscript": "Script/Handwritten", "petrona": "Serif (Generic)", "philosopher": "Sans Serif (Generic)", + "piedra": "Display/Other", "pinyonscript": "Script/Handwritten", "pirataone": "Display/Other", + "plaster": "Display/Other", "play": "Sans Serif (Generic)", "playball": "Script/Handwritten", + "playfairdisplay": "Modern Serif", "playfairdisplaysc": "Modern Serif", "podkova": "Slab Serif", + "podkovavfbeta": "Slab Serif", "poetsenone": "Display/Other", "poiretone": "Display/Other", + "poly": "Serif (Generic)", "pompiere": "Display/Other", "ponnala": "Display/Other", + "pontanosans": "Sans Serif (Generic)", "poorstory": "Display/Other", "poppins": "Geometric Sans", + "portersansblock": "Display/Other", "portlligatsans": "Display/Other", "portlligatslab": "Display/Other", + "postnobillscolombo": "Display/Other", "postnobillsjaffna": "Display/Other", "pragatinarrow": "Condensed Sans", + "prata": "Serif (Generic)", "pressstart2p": "Display/Other", "pridi": "Serif (Generic)", + "princesssofia": "Script/Handwritten", "prociono": "Serif (Generic)", "prompt": "Sans Serif", + "prostoone": "Display/Other", "prozalibre": "Sans Serif (Generic)", "pt_sans": "Humanist Sans", + "pt_serif": "Transitional Serif", "puritan": "Sans Serif (Generic)", "purplepurse": "Display/Other", + "questrial": "Geometric Sans", "qwigley": "Script/Handwritten" + } + + def get_cluster(self, raw_name: str) -> str: + clean = raw_name.lower().split('-')[0].strip() + if clean in self.full_font_map: + return self.full_font_map[clean] + return "Display/Other" + + +class UnifiedStyleEngine: + def __init__(self): + self.feature_registry = defaultdict(lambda: defaultdict(list)) + self.font_engine = FontClusterEngine() + self.font_family_scores = defaultdict(float) + self.font_individual_scores = defaultdict(float) + self.font_members = defaultdict(list) + self.color_pool = [] + self.doc_counter = 0 + + self.material_filter = { + "glossy", "matte", "metallic", "organic", + "wood", "laminated", "porcelain", "wet-look", + "denim", "fabric", "leather" + } + + self.aliases = { + "Colour Palette": "Color Palette", + "Texture": "Background/Texture", + "Era": "Era/Cultural Reference", + "Font": "Typography", + "Fonts": "Typography", + "Text": "Typography", + "Layout": "Composition" + } + + def ingest_data(self, data: Dict): + """ + Parses the JSON structure from the app + Handles: { "success": true, "data": { "results": { ... } } } + """ + try: + self.doc_counter += 1 + + source_data = {} + if "data" in data and isinstance(data["data"], dict): + if "results" in data["data"]: + source_data = data["data"]["results"] + else: + source_data = data["data"] + elif "results" in data: + source_data = data["results"] + else: + source_data = data + + for category, payload in source_data.items(): + if category in ["filename", "meta"]: continue + + std_category = self.aliases.get(category, category) + + if std_category == "Color Palette": + if isinstance(payload, dict) and "palette" in payload: + for hex_code in payload["palette"]: + if isinstance(hex_code, str) and hex_code.startswith('#'): + self.color_pool.append(hex_to_rgb(hex_code)) + elif isinstance(payload, list): + for item in payload: + if isinstance(item, str) and item.startswith('#'): + self.color_pool.append(hex_to_rgb(item)) + + elif std_category == "Typography": + if isinstance(payload, dict) and "scores" in payload: + payload = payload["scores"] + vectors = self._normalize(payload) + + for rank, (font_name, score) in enumerate(vectors): + if "no text" in font_name.lower(): continue + + weight = math.pow(RANK_DECAY_FACTOR, rank) * score + cluster = self.font_engine.get_cluster(font_name) + + self.font_family_scores[cluster] += weight + self.font_individual_scores[font_name] += weight + if font_name not in self.font_members[cluster]: + self.font_members[cluster].append(font_name) + + else: + if isinstance(payload, dict) and "scores" in payload: + payload = payload["scores"] + vectors = self._normalize(payload) + + for rank, (label, score) in enumerate(vectors): + self.feature_registry[std_category][label].append({ + "raw_score": score, + "rank": rank + }) + + except Exception as e: + print(f"Warning during ingestion: {e}") + + def _normalize(self, payload) -> List[Tuple[str, float]]: + """ + Robustly converts inputs to [(label, score), ...] + Handles mixed types and weird nesting to prevent crashes. + """ + if isinstance(payload, dict): + clean_items = [] + for k, v in payload.items(): + if isinstance(v, (int, float)): + clean_items.append((k, float(v))) + return sorted(clean_items, key=lambda x: x[1], reverse=True) + + if isinstance(payload, list): + res = [] + for item in payload: + if isinstance(item, dict): + try: + score = float(item.get('score', 0.9)) + except (ValueError, TypeError): + score = 0.0 + + label = item.get('label', item.get('tag', 'unknown')) + if not isinstance(label, str): label = str(label) + + res.append((label, score)) + elif isinstance(item, str): + res.append((item, 1.0)) + return res + return [] + + def compute_final_stylesheet(self) -> Dict: + final_json = {"results": {}} + + for category, tags_map in self.feature_registry.items(): + candidates = [] + for tag, occurrences in tags_map.items(): + w_sum = sum(obs['raw_score'] * math.pow(RANK_DECAY_FACTOR, obs['rank']) for obs in occurrences) + score = w_sum / max(1, self.doc_counter) + + if score > MIN_CONSENSUS_THRESHOLD: + candidates.append({"label": tag, "score": round(score, 2)}) + + candidates.sort(key=lambda x: x['score'], reverse=True) + if candidates: + final_json["results"][category] = candidates[:5] + + if "Background/Texture" in final_json["results"]: + textures = [] + materials = [] + + for item in final_json["results"]["Background/Texture"]: + label = item["label"].lower().replace("_", " ").replace("-", " ") + is_material = False + if item["label"] in self.material_filter: + is_material = True + else: + for mat in self.material_filter: + if mat in label: + is_material = True + break + + if is_material: + materials.append(item) + else: + textures.append(item) + + if textures: + final_json["results"]["Background/Texture"] = textures + else: + del final_json["results"]["Background/Texture"] + + if materials: + if "Material Look" in final_json["results"]: + final_json["results"]["Material Look"].extend(materials) + final_json["results"]["Material Look"].sort(key=lambda x: x['score'], reverse=True) + else: + final_json["results"]["Material Look"] = materials + + if self.font_family_scores: + best_fam, _ = sorted(self.font_family_scores.items(), key=lambda x: x[1], reverse=True)[0] + members = self.font_members[best_fam] + members.sort(key=lambda x: self.font_individual_scores[x], reverse=True) + + typo_output = [] + for font in members[:5]: + score = self.font_individual_scores[font] / max(1, self.doc_counter) + typo_output.append({"label": font, "score": round(score, 2)}) + + final_json["results"]["Typography"] = typo_output + final_json["results"]["Typography_Family"] = best_fam + + if len(self.color_pool) > 0: + try: + k_clusters = min(len(self.color_pool), PALETTE_SIZE) + + kmeans = KMeans(n_clusters=k_clusters, n_init=1, random_state=42) + kmeans.fit(self.color_pool) + centers = sorted(kmeans.cluster_centers_.astype(int).tolist(), key=sum) + + palette_output = [] + for i, rgb in enumerate(centers): + palette_output.append({ + "label": rgb_to_hex(rgb), + "score": round(1.0 - (i * 0.1), 2) + }) + final_json["results"]["Color Palette"] = palette_output + except Exception as e: + print(f"Color clustering failed: {e}") + + final_json["results"]["meta"] = {"source_files": self.doc_counter} + return final_json + + +def generate_stylesheet(json_strings_list: Any) -> str: + """ + Entry point + Accepts a Java ArrayList of JSON strings. + """ + engine = UnifiedStyleEngine() + + try: + count = json_strings_list.size() + for i in range(count): + json_str = json_strings_list.get(i) + if not json_str: continue + try: + data = json.loads(json_str) + engine.ingest_data(data) + except Exception as e: + print(f"Error parsing JSON chunk at index {i}: {e}") + continue + + except AttributeError: + for json_str in json_strings_list: + if not json_str: continue + try: + data = json.loads(json_str) + engine.ingest_data(data) + except: continue + + final_output = engine.compute_final_stylesheet() + return json.dumps(final_output) + + +# PROMPT GENERATION LOGIC + + +DEFAULT_OPTIONS = { + "width": 1024, + "height": 1024, + "steps": 30, + "guidance_scale": 7.5, + "sampler": "k_lms", + "seed": "random", +} + +def normalize_sheet(sheet: Dict[str, Any]) -> Dict[str, Dict[str, List[str]]]: + """ + Organizes the flat stylesheet tags into Primary and Secondary categories. + """ + results = sheet.get("results", sheet) + norm: Dict[str, Dict[str, List[str]]] = {} + + for cat, val in results.items(): + if cat in ("filename", "meta", "Color Palette"): + continue + + items = [] + if isinstance(val, list): + items = val + elif isinstance(val, dict) and "scores" in val: + items = val["scores"] + + try: + items_sorted = sorted( + items, key=lambda x: float(x.get("score", 0)), reverse=True + ) + except Exception: + items_sorted = items + + labels = [ + str(it.get("label")).strip() for it in items_sorted if it.get("label") + ] + if labels: + norm[cat] = {"Primary": labels[0], "Secondary": labels[1:]} + + return norm + +STYLE_WORDS = [ + "hand-drawn", "hand drawn", "sketch", "thin lines", "minimalistic", + "minimal", "flat", "cartoon", "comic", "pop-art", "surreal", "photorealistic", + "realistic", "painting", "oil painting", "pastel", "artistic", "illustration", + "line drawing", "line art", "3d effect", "texture", "pattern", "soft light", + "studio light", "dramatic", "cinematic", "moody", "vintage", "retro", "poster", + "render", "octane render", "unreal engine","digital illustration", +] + +style_words_regex = re.compile( + r"\b(" + "|".join(re.escape(w) for w in STYLE_WORDS) + r")\b", + flags=re.IGNORECASE, +) + +def remove_style_words(caption: str) -> str: + """ + Removes specific artistic style descriptions from the caption so they + don't conflict with the Stylesheet. + """ + if not caption: return caption + clean = style_words_regex.sub(" ", caption) + clean = re.sub(r"\s{2,}", " ", clean).strip() + return clean + +def extract_subject(caption_text: str) -> str: + """ + Heuristic to guess the main subject (first few words of the first sentence). + """ + txt = caption_text.strip() + if not txt: return "subject" + first_clause = re.split(r"[.?!]\s*", txt)[0] + return " ".join(first_clause.split()[:12]) + +def pick_style_phrases(norm: Dict[str, Dict[str, List[str]]], max_secondaries=3) -> List[str]: + """ + Extracts non-color features from the Stylesheet. + """ + cats = [ + "Style", "Background/Texture", "Lighting", "Composition", + "Era/Cultural Reference", "Material Look", "Typography", + ] + key_map = { + "Style": "art style", + "Background/Texture": "background texture", + "Lighting": "lighting", + "Composition": "composition", + "Era/Cultural Reference": "era", + "Material Look": "material", + "Typography": "typography", + } + + out: List[str] = [] + for cat in cats: + entry = norm.get(cat) + if not entry: continue + + prim = entry.get("Primary") + secs = entry.get("Secondary", []) + chosen: List[str] = [] + if prim: chosen.append(prim) + for s in secs: + if s and len(chosen) < (1 + max_secondaries): + chosen.append(s) + if chosen: + out.append(f"{key_map.get(cat, cat)}: {', '.join(chosen)}") + return out + +def style_to_tone(primary: str) -> str: + """ + Infers a general tone based on the primary style. + """ + s = (primary or "").lower() + if any(k in s for k in ["retro", "collage", "poster", "pop-art", "surreal"]): + return "illustrative, textured, poster-like" + if any(k in s for k in ["photoreal", "realistic", "film"]): + return "photorealistic, high-detail" + return "high detail, sharp focus" + +def generate_magic_prompt(stylesheet_json_str: str, caption: str, user_prompt: str) -> str: + try: + sheet = json.loads(stylesheet_json_str) + except Exception as e: + return f"{user_prompt}. The image features: {caption}" + + cleaned_caption = remove_style_words(caption) + + norm = normalize_sheet(sheet) + subject = extract_subject(cleaned_caption) + + style_phrases = pick_style_phrases(norm) + + style_primary = (norm.get("Style") or {}).get("Primary", "") + tone_hint = style_to_tone(style_primary) + + + style_instruction = "" + if style_phrases: + style_instruction = "Style elements: " + "; ".join(style_phrases) + + prompt_lines: List[str] = [ + f"Main subject: {subject}.", + f"Description: {cleaned_caption}.", + ] + + if user_prompt: + prompt_lines.append(f"User instruction: {user_prompt}.") + + if style_instruction: + prompt_lines.append(f"{style_instruction}.") + + prompt_lines += [ + f"Tone: {tone_hint}.", + "Note: Colors should be derived strictly from the Description.", + "Avoid: blurry, deformed, text, watermark, extra limbs, artifacts.", + ] + + prompt = "\n".join(prompt_lines) + + payload = { + "prompt": prompt, + "negative_prompt": "blurry, low resolution, deformed, text, watermark, extra limbs, artifacts", + "options": DEFAULT_OPTIONS, + "provenance": {"generated_from": "final_output"}, + "weights": {"caption": 0.7, "style_sheet": 0.2, "user": 0.1}, + } + + return prompt + " " + json.dumps(payload) +\ No newline at end of file diff --git a/models/subject.py b/models/subject.py @@ -0,0 +1,180 @@ +# -*- coding: utf-8 -*- +"""BiRefNet.ipynb + +Automatically generated by Colab. + +Original file is located at + https://colab.research.google.com/drive/1izs_-Hld8X4wtGylu8y-96PJb04Nc3kz +""" + +# Commented out IPython magic to ensure Python compatibility. +!pip install -r https://raw.githubusercontent.com/ZhengPeng7/BiRefNet/main/requirements.txt +!git clone https://github.com/ZhengPeng7/BiRefNet.git + +# %cd BiRefNet + +"""Run This To Create The .pt file""" + +# Commented out IPython magic to ensure Python compatibility. +# %cd BiRefNet + +import torch +from torchvision import transforms +from models.birefnet import BiRefNet + +device = "cuda" if torch.cuda.is_available() else "cpu" +print("Using device:", device) + +birefnet = BiRefNet.from_pretrained("ZhengPeng7/BiRefNet") +birefnet.to(device) +birefnet.eval() +birefnet.half() + + +model_fp16_cpu = birefnet.to("cpu") +torch.save(model_fp16_cpu.state_dict(), "birefnet_fp16.pt") + +print("🔥 Saved fp16 model to birefnet_fp16.pt") + +image_size = (1024, 1024) +transform_image = transforms.Compose([ + transforms.Resize(image_size), + transforms.ToTensor(), + transforms.Normalize( + mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225], + ), +]) + +"""Inference Code""" + +import os +import io +import time +import numpy as np +from PIL import Image +import cv2 +import matplotlib.pyplot as plt + +import torch +from torchvision import transforms +from google.colab import files + + +device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +print("Using device:", device) + + +OUTPUT_DIR = "/content/output/birefnet_cutouts" +os.makedirs(OUTPUT_DIR, exist_ok=True) +print("Cutouts will be saved in:", OUTPUT_DIR) + + +image_size = (1024, 1024) +transform_image = transforms.Compose([ + transforms.Resize(image_size), + transforms.ToTensor(), + transforms.Normalize( + mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225], + ), +]) + + +if 'birefnet' not in globals(): + raise RuntimeError( + "Model `birefnet` not found. Load your BiRefNet model and weights first, e.g.:\n" + " birefnet = BiRefNet(...)\n" + " birefnet.load_state_dict(torch.load('ckpt.pth', map_location='cpu'))\n" + " birefnet.to(device).eval()\n" + ) + +birefnet = birefnet.to(device).eval() + + +def run_birefnet_on_image(pil_img: Image.Image, thresh: float = 0.5): + """ + Runs BiRefNet on a single PIL image. + Returns: + mask_bin : (H, W) uint8 (0 or 1) -- final model output mask + prob_resized : (H, W) float32 [0,1] -- probability map (not saved by default) + infer_time : float seconds + """ + orig_w, orig_h = pil_img.size + + + inp = transform_image(pil_img).unsqueeze(0).to(device) + try: + inp = inp.half() + except Exception: + pass + + with torch.no_grad(): + t0 = time.perf_counter() + preds = birefnet(inp)[-1] + preds = preds.float().sigmoid().cpu() + t1 = time.perf_counter() + + prob = preds[0, 0].numpy().astype(np.float32) + prob_resized = cv2.resize(prob, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR) + mask_bin = (prob_resized >= thresh).astype(np.uint8) + + return mask_bin, prob_resized, (t1 - t0) + + +def make_binary_cutout(img_rgb: np.ndarray, mask_bin: np.ndarray): + """ + img_rgb: HxWx3 uint8 + mask_bin: HxW uint8 {0,1} + Returns RGBA HxWx4 uint8 with alpha=255 on foreground, 0 on background. + """ + H, W = img_rgb.shape[:2] + if mask_bin.shape != (H, W): + mask_bin = cv2.resize(mask_bin, (W, H), interpolation=cv2.INTER_NEAREST) + + alpha_u8 = (mask_bin * 255).astype(np.uint8) + rgba = np.dstack([img_rgb, alpha_u8]) + return rgba + +#upload +print("🔼 Upload one or more images to segment with BiRefNet...") +uploaded = files.upload() + +THRESH = 0.5 +KEEP_LARGEST = True + +for fname, file_data in uploaded.items(): + try: + print(f"\n=== Processing: {fname} ===") + pil_img = Image.open(io.BytesIO(file_data)).convert("RGB") + img_np = np.array(pil_img) + + mask_bin, prob_map, t_inf = run_birefnet_on_image(pil_img, thresh=THRESH) + print(f"BiRefNet inference time: {t_inf:.3f} s") + + if KEEP_LARGEST: + num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask_bin, connectivity=8) + if num_labels > 1: + areas = stats[1:, cv2.CC_STAT_AREA] + max_idx = 1 + np.argmax(areas) + mask_bin = (labels == max_idx).astype(np.uint8) + + rgba_cutout = make_binary_cutout(img_np, mask_bin) + + base = os.path.splitext(fname)[0] + out_path = os.path.join(OUTPUT_DIR, base + "_birefnet_cutout.png") + Image.fromarray(rgba_cutout).save(out_path, format="PNG") + print(" Saved final model cutout:", out_path) + + fig, axs = plt.subplots(1, 3, figsize=(12, 4)) + axs[0].imshow(img_np); axs[0].set_title("Original"); axs[0].axis("off") + axs[1].imshow(mask_bin, cmap="gray"); axs[1].set_title("Final model mask"); axs[1].axis("off") + axs[2].imshow(rgba_cutout); axs[2].set_title("Cutout (PNG, no background)"); axs[2].axis("off") + plt.tight_layout(); plt.show() + + files.download(out_path) + + except Exception as e: + print(f" Error with {fname}: {e}") + +print("\n Done! All cutouts saved in:", OUTPUT_DIR) +\ No newline at end of file diff --git a/models/testing_data_gen_template.py b/models/testing_data_gen_template.py @@ -0,0 +1,152 @@ +# -*- coding: utf-8 -*- +"""testing_data_gen_template.ipynb + +Automatically generated by Colab. + +Original file is located at + https://colab.research.google.com/drive/1iDu3v3m3iOKyRs6IggnOPbN6tKKa9SRV +""" + +import numpy as np +import pandas as pd + +!pip install requests tqdm + +pip install icrawler + +import os +import shutil +import urllib3 +import logging +from icrawler import ImageDownloader +from icrawler.builtin import BingImageCrawler + +# --- CONFIGURATION --- +OUTPUT_ROOT = "Feature_Dataset_HQ" +TARGET_COUNT = 100 + +urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) + +class No403Filter(logging.Filter): + def filter(self, record): + return "403" not in record.getMessage() + +logging.getLogger("downloader").addFilter(No403Filter()) + +class RobustDownloader(ImageDownloader): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.session.verify = False + self.session.headers.update({ + 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36', + 'Referer': 'https://www.bing.com/' + }) +#classes according to the feature +SEARCH_MAPPING = { + "soft-light": "soft light, gentle illumination, low contrast, wrap-around light", + "specular-highlights": "strong specular highlights, shiny skin, glossy reflections, bright hotspots", + "backlit": "backlit lighting from behind the subject, rim light, silhouette effect, sun flare", + "studio-lighting": "professional studio lighting, controlled illumination, three-point lighting setup", + "flat-lighting": "flat lighting, low contrast, even illumination, no shadows, cloudy day look", + "dramatic-contrast": "dramatic high-contrast lighting, chiaroscuro, hard shadows, noir style", + "diffused": "soft diffused lighting, cloudy sky light, light box look, scattered light, no hard edges" +} + +def count_files(folder): + if not os.path.exists(folder): return 0 + return len([f for f in os.listdir(folder) if f.endswith(('.jpg', '.png', '.jpeg'))]) + +def trim_excess(folder, limit): + files = sorted([f for f in os.listdir(folder) if f.endswith(('.jpg', '.png', '.jpeg'))]) + if len(files) > limit: + print(f" -> Overshot! Trimming {len(files) - limit} extra images...") + for f in files[limit:]: + try: os.remove(os.path.join(folder, f)) + except: pass + +if __name__ == "__main__": + if not os.path.exists(OUTPUT_ROOT): + os.makedirs(OUTPUT_ROOT) + + print(f"Starting Recursive Download (Target: {TARGET_COUNT})...") + + for category, query in SEARCH_MAPPING.items(): + print(f"\n--- Processing: {category} ---") + save_dir = os.path.join(OUTPUT_ROOT, category) + + # Ensure folder exists to prevent FileNotFoundError + if not os.path.exists(save_dir): + os.makedirs(save_dir) + + # Initial Search Depth + search_depth = 200 + + while True: + current_count = count_files(save_dir) + + # 1. CHECK CONDITION + if current_count >= TARGET_COUNT: + print(f" -> Success! Found {current_count} images.") + break + + # 2. INCREASE DEPTH + # If we already have some images, we need to search deeper to find NEW ones + print(f" -> Have {current_count}/{TARGET_COUNT}. Increasing search depth to {search_depth}...") + + crawler = BingImageCrawler( + downloader_cls=RobustDownloader, + feeder_threads=1, + parser_threads=2, + downloader_threads=4, + storage={'root_dir': save_dir}, + log_level='ERROR' + ) + + # 3. SAFETY OFFSET + # If folder is empty, offset=0. If files exist, offset='auto' to append. + if current_count == 0: + offset_param = 0 + else: + offset_param = 'auto' + + crawler.crawl( + keyword=query, + max_num=search_depth, + filters={'type': 'photo'}, + file_idx_offset=offset_param + ) + + # Increase the search buffer for the next loop + # If 200 wasn't enough, next time try 700, then 1200... + search_depth += 500 + + # Safety break + if search_depth > 5000: + print(" -> Hit safety limit (5000). Stopping.") + break + + # Final Trim + trim_excess(save_dir, TARGET_COUNT) + print(f" Final Count: {count_files(save_dir)}") + + print(f"\n Material Look Dataset (HQ) Download Complete!") + +import shutil +import os + +# --- CONFIGURATION --- +FOLDER_TO_ZIP = "MaterialLook_Dataset_HQ" +OUTPUT_FILENAME = "MaterialLook_Dataset" # .zip will be added automatically + +if __name__ == "__main__": + if os.path.exists(FOLDER_TO_ZIP): + print(f"Compressing '{FOLDER_TO_ZIP}'...") + + # This creates the zip file + shutil.make_archive(OUTPUT_FILENAME, 'zip', FOLDER_TO_ZIP) + + # Check size for confirmation + zip_size = os.path.getsize(f"{OUTPUT_FILENAME}.zip") / (1024 * 1024) + print(f" Success! Created {OUTPUT_FILENAME}.zip ({zip_size:.2f} MB)") + else: + print(f" Error: Folder '{FOLDER_TO_ZIP}' not found.") +\ No newline at end of file diff --git a/models/text&mat/dinov2_vits14_fp16.pt b/models/text&mat/dinov2_vits14_fp16.pt Binary files differ. diff --git a/models/text&mat/texture&mat.py b/models/text&mat/texture&mat.py @@ -0,0 +1,441 @@ +# -*- coding: utf-8 -*- +"""Untitled7.ipynb + +Automatically generated by Colab. + +Original file is located at + https://colab.research.google.com/drive/1u-5cHHckvZPMjIFJ5YJcCAGzDWVPRl91 +""" + +import numpy as np +import pandas as pd + +import torch +import torch.nn as nn +import numpy as np +import os + +CENTROIDS_PATH = 'texture_centroids_augmented.pth_path' +OUTPUT_PATH = 'texture_centroids_augmented.pt_path' + +class DataWrapper(nn.Module): + def __init__(self, data_dict): + super().__init__() + self.data = torch.jit.Attribute(data_dict, dict[str, torch.Tensor]) + + def forward(self): + return self.data + +if os.path.exists(CENTROIDS_PATH): + print("Loading .pth file...") + + raw_data = torch.load(CENTROIDS_PATH, map_location='cpu', weights_only=False) + + print("Processing data...") + clean_dict = {} + for k, v in raw_data.items(): + if isinstance(v, np.ndarray): + clean_dict[k] = torch.from_numpy(v).float() + else: + clean_dict[k] = v.float() + + wrapper = DataWrapper(clean_dict) + scripted_wrapper = torch.jit.script(wrapper) + + scripted_wrapper.save(OUTPUT_PATH) + print(f"Success! Saved as {OUTPUT_PATH}") +else: + print("File not found.") + +import torch +import torchvision.transforms as T +from PIL import Image +import numpy as np +import os +import glob +from typing import Dict, List + +DATASET_ROOT = 'Texture_Dataset_path' +OUTPUT_PATH = 'texture_centroids_s14.pt' +DEVICE = torch.device('cuda' if torch.cuda.is_available() else "cpu") + + +print(f"Loading DINOv2-S14 on {DEVICE}...") +model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14') +model.to(DEVICE) +model.eval() + +transform = T.Compose([ + T.Resize(256, interpolation=T.InterpolationMode.BICUBIC), + T.CenterCrop(224), + T.ToTensor(), + T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), +]) + + + +def get_folder_centroid(folder_path): + image_files = glob.glob(os.path.join(folder_path, "*")) + valid_exts = ('.jpg', '.jpeg', '.png', '.bmp', '.webp') + image_files = [f for f in image_files if f.lower().endswith(valid_exts)] + + if not image_files: return None + print(f"Processing '{os.path.basename(folder_path)}' ({len(image_files)} images)...") + + vectors = [] + for img_path in image_files: + try: + img = Image.open(img_path).convert('RGB') + img_t = transform(img).unsqueeze(0).to(DEVICE) + with torch.no_grad(): + output = model.forward_features(img_t) + vectors.append(output['x_norm_clstoken'].cpu()) + except Exception: continue + + if not vectors: return None + + all_vecs = torch.cat(vectors, dim=0) + centroid = torch.mean(all_vecs, dim=0) + centroid = torch.nn.functional.normalize(centroid, p=2, dim=0) + return centroid + + + +class CentroidsWrapper(torch.nn.Module): + def __init__(self, data_dict): + super().__init__() + self.keys: List[str] = list(data_dict.keys()) + self.tensors = torch.nn.ParameterList([ + torch.nn.Parameter(data_dict[k]) for k in self.keys + ]) + + def forward(self) -> Dict[str, torch.Tensor]: + result: Dict[str, torch.Tensor] = {} + for i, tensor in enumerate(self.tensors): + result[self.keys[i]] = tensor + return result + + +# EXECUTION + + +if not os.path.exists(DATASET_ROOT): + print(f" Error: Dataset path not found: {DATASET_ROOT}") +else: + centroids_dict = {} + folders = sorted([d for d in os.listdir(DATASET_ROOT) if os.path.isdir(os.path.join(DATASET_ROOT, d))]) + + print(f"Found {len(folders)} texture categories.") + + for category in folders: + full_path = os.path.join(DATASET_ROOT, category) + centroid = get_folder_centroid(full_path) + if centroid is not None: + centroids_dict[category] = centroid + + if centroids_dict: + print(f"\nSaving {len(centroids_dict)} centroids to {OUTPUT_PATH}...") + + wrapper = CentroidsWrapper(centroids_dict) + scripted_wrapper = torch.jit.script(wrapper) + scripted_wrapper.save(OUTPUT_PATH) + + print(" SUCCESS! File saved.") + print(f"Output: {OUTPUT_PATH}") + else: + print(" Failed to generate any centroids.") + +import torch +import torchvision.transforms as T +import torch.nn.functional as F +from PIL import Image +import numpy as np +import os + +DEVICE = torch.device('cuda' if torch.cuda.is_available() else "cpu") +PATCH_SIZE = 14 + +CENTROIDS_PATH = 'CENTROIDS_PATH' +TEST_IMAGE_PATH = 'TEST_IMAGE_PATH' + +class TextureScoreScanner: + def __init__(self, centroids_dict): + self.device = DEVICE + print(f"Loading DINOv2 (S14 - Small) on {self.device}...") + self.model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14') + self.model = self.model.half().to(self.device) + self.model.eval() + + self.texture_names = [] + matrix_list = [] + print("Processing centroids...") + for name, tensor in centroids_dict.items(): + if tensor.dim() > 1: tensor = tensor.squeeze() + tensor = tensor.half().to(self.device) + + self.texture_names.append(name) + matrix_list.append(tensor) + + if len(matrix_list) > 0: + matrix = torch.stack(matrix_list) + self.centroid_matrix = F.normalize(matrix, p=2, dim=1) + + if self.centroid_matrix.shape[1] != 384: + print(f" WARNING: Dimension Mismatch! S14 expects 384, got {self.centroid_matrix.shape[1]}") + else: + raise ValueError("No valid centroids found in the .pt file.") + + def preprocess(self, img_path): + if not os.path.exists(img_path): raise FileNotFoundError(f"Image not found: {img_path}") + img = Image.open(img_path).convert('RGB') + w, h = img.size + + new_w = (w // PATCH_SIZE) * PATCH_SIZE + new_h = (h // PATCH_SIZE) * PATCH_SIZE + + transform = T.Compose([ + T.Resize((new_h, new_w)), + T.ToTensor(), + T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), + ]) + return transform(img).unsqueeze(0).half().to(self.device) + + def get_scores(self, img_path): + input_tensor = self.preprocess(img_path) + + with torch.no_grad(): + features = self.model.forward_features(input_tensor) + + patches = features["x_norm_patchtokens"].squeeze(0) + patches_norm = F.normalize(patches, p=2, dim=1) + + similarity_matrix = torch.mm(patches_norm, self.centroid_matrix.T) + + class_scores = torch.mean(similarity_matrix, dim=0).cpu().numpy() + + results = {} + for i, score in enumerate(class_scores): + results[self.texture_names[i]] = float(score) + + return sorted(results.items(), key=lambda item: item[1], reverse=True) + + def get_filtered_predictions(self, sorted_scores, max_count=5, relative_thresh=0.85, absolute_floor=0.15): + if not sorted_scores: return [] + + final_list = [] + top_score = sorted_scores[0][1] + + for name, score in sorted_scores: + if len(final_list) >= max_count: break + + if score < absolute_floor: continue + if score < (top_score * relative_thresh): break + + final_list.append((name, score)) + + return final_list + +# --- MAIN EXECUTION --- + +if os.path.exists(CENTROIDS_PATH): + print(f"Loading dictionary from {CENTROIDS_PATH}...") + + loaded_container = torch.jit.load(CENTROIDS_PATH, map_location='cpu') + centroids_data = loaded_container() + + scanner = TextureScoreScanner(centroids_data) + + print(f"\nAnalyzing: {TEST_IMAGE_PATH}") + if os.path.exists(TEST_IMAGE_PATH): + raw_scores = scanner.get_scores(TEST_IMAGE_PATH) + + smart_results = scanner.get_filtered_predictions( + raw_scores, + max_count=3, # Show max 3 textures + relative_thresh=0.85, # Must be 85% as good as the top match + absolute_floor=0.15 # Minimum similarity score + ) + + print("\n" + "="*40) + print(f"DETECTED TEXTURES ({len(smart_results)})") + print("="*40) + + if len(smart_results) > 0: + top_score_val = raw_scores[0][1] + for name, score in smart_results: + percentage = (score / top_score_val) * 100 + print(f"• {name.upper().ljust(15)} : Score {score:.4f} (Conf: {percentage:.0f}%)") + else: + print("No significant texture match found.") + else: + print(f"Error: Test image not found at {TEST_IMAGE_PATH}") + +else: + print(f" Error: Centroids file '{CENTROIDS_PATH}' not found.") + print("Please run the training/generation script first.") + +import torch +import os + +MODEL_NAME = 'dinov2_vits14' +OUTPUT_FILENAME = "dinov2_vits14_fp16.pt" + +print(f" Downloading {MODEL_NAME} (FP32) from Torch Hub...") +model = torch.hub.load('facebookresearch/dinov2', MODEL_NAME) + +print(" Converting model to Half Precision (FP16)...") +model = model.half() + +print(f" Saving weights to {OUTPUT_FILENAME}...") +torch.save(model.state_dict(), OUTPUT_FILENAME) + +# Verify +file_size_mb = os.path.getsize(OUTPUT_FILENAME) / (1024 * 1024) +print("-" * 30) +print(f" Done! Saved to: {os.path.abspath(OUTPUT_FILENAME)}") +print(f" File Size: {file_size_mb:.2f} MB") +print("-" * 30) + +import torch +import torchvision.transforms as T +import torch.nn.functional as F +from PIL import Image +import numpy as np +import os + +DEVICE = torch.device('cuda' if torch.cuda.is_available() else "cpu") +PATCH_SIZE = 14 + +CENTROIDS_PATH = 'centroids_path' +TEST_IMAGE_PATH = 'test_image_path' + +MODEL_WEIGHTS_PATH = '/kaggle/working/dinov2_vits14_fp16.pt' + +class TextureScoreScanner: + def __init__(self, centroids_dict): + self.device = DEVICE + print(f"Loading DINOv2 (S14 - Small) on {self.device}...") + + self.model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14', pretrained=False) + + self.model = self.model.half().to(self.device) + + if os.path.exists(MODEL_WEIGHTS_PATH): + print(f"Loading local weights from: {MODEL_WEIGHTS_PATH}") + state_dict = torch.load(MODEL_WEIGHTS_PATH, map_location=self.device) + self.model.load_state_dict(state_dict) + else: + raise FileNotFoundError(f"Weights file not found at: {MODEL_WEIGHTS_PATH}") + + self.model.eval() + + self.texture_names = [] + matrix_list = [] + + print("Processing centroids...") + for name, tensor in centroids_dict.items(): + if tensor.dim() > 1: tensor = tensor.squeeze() + + tensor = tensor.half().to(self.device) + + self.texture_names.append(name) + matrix_list.append(tensor) + + if len(matrix_list) > 0: + matrix = torch.stack(matrix_list) + self.centroid_matrix = F.normalize(matrix, p=2, dim=1) + + if self.centroid_matrix.shape[1] != 384: + print(f"⚠️ WARNING: Dimension Mismatch! S14 expects 384, got {self.centroid_matrix.shape[1]}") + else: + raise ValueError("No valid centroids found in the .pt file.") + + def preprocess(self, img_path): + if not os.path.exists(img_path): raise FileNotFoundError(f"Image not found: {img_path}") + img = Image.open(img_path).convert('RGB') + w, h = img.size + + new_w = (w // PATCH_SIZE) * PATCH_SIZE + new_h = (h // PATCH_SIZE) * PATCH_SIZE + + transform = T.Compose([ + T.Resize((new_h, new_w)), + T.ToTensor(), + T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), + ]) + + return transform(img).unsqueeze(0).half().to(self.device) + + def get_scores(self, img_path): + input_tensor = self.preprocess(img_path) + + with torch.no_grad(): + features = self.model.forward_features(input_tensor) + + patches = features["x_norm_patchtokens"].squeeze(0) + patches_norm = F.normalize(patches, p=2, dim=1) + + similarity_matrix = torch.mm(patches_norm, self.centroid_matrix.T) + + class_scores = torch.mean(similarity_matrix, dim=0).float().cpu().numpy() + + results = {} + for i, score in enumerate(class_scores): + results[self.texture_names[i]] = float(score) + + return sorted(results.items(), key=lambda item: item[1], reverse=True) + + def get_filtered_predictions(self, sorted_scores, max_count=5, relative_thresh=0.85, absolute_floor=0.15): + if not sorted_scores: return [] + + final_list = [] + top_score = sorted_scores[0][1] + + for name, score in sorted_scores: + if len(final_list) >= max_count: break + + if score < absolute_floor: continue + if score < (top_score * relative_thresh): break + + final_list.append((name, score)) + + return final_list + +# --- MAIN EXECUTION --- + +if os.path.exists(CENTROIDS_PATH): + print(f"Loading dictionary from {CENTROIDS_PATH}...") + + loaded_container = torch.jit.load(CENTROIDS_PATH, map_location='cpu') + centroids_data = loaded_container() + + scanner = TextureScoreScanner(centroids_data) + + print(f"\nAnalyzing: {TEST_IMAGE_PATH}") + if os.path.exists(TEST_IMAGE_PATH): + raw_scores = scanner.get_scores(TEST_IMAGE_PATH) + + smart_results = scanner.get_filtered_predictions( + raw_scores, + max_count=3, + relative_thresh=0.85, + absolute_floor=0.05 # Kept at 0.05 as discussed + ) + + print("\n" + "="*40) + print(f"DETECTED TEXTURES ({len(smart_results)})") + print("="*40) + + if len(smart_results) > 0: + top_score_val = raw_scores[0][1] + for name, score in smart_results: + percentage = (score / top_score_val) * 100 + print(f"• {name.upper().ljust(15)} : Score {score:.4f} (Conf: {percentage:.0f}%)") + else: + print("No significant texture match found.") + else: + print(f"Error: Test image not found at {TEST_IMAGE_PATH}") + +else: + print(f" Error: Centroids file '{CENTROIDS_PATH}' not found.") + diff --git a/models/text&mat/texture_centroids_s14.pt b/models/text&mat/texture_centroids_s14.pt Binary files differ. diff --git a/models/typo.py b/models/typo.py @@ -0,0 +1,77 @@ +# -*- coding: utf-8 -*- +"""Untitled5.ipynb + +Automatically generated by Colab. + +Original file is located at + https://colab.research.google.com/drive/11sdpsNaYwetetOotIX3GKtxgaQOt7S0Y +""" + +!pip install easyocr -q + +import numpy as np +from PIL import Image +import easyocr +import pickle +import os +import json +from tensorflow.keras.models import model_from_json, Model +from tensorflow.keras.utils import get_custom_objects + +TEST_IMAGE_PATH = 'test image.png' +MODEL_JSON = 'fannet.json' +MODEL_WEIGHTS = 'fannet_weights.h5' +FONT_DB_PKL = 'font_database.pkl' + +if os.path.exists(FONT_DB_PKL) and os.path.exists(MODEL_JSON): + with open(FONT_DB_PKL, 'rb') as f: + font_db = pickle.load(f) + + get_custom_objects()['Model'] = Model + + with open(MODEL_JSON, 'r') as f: + model_json_content = f.read() + + fannet_model = model_from_json(model_json_content) + fannet_model.load_weights(MODEL_WEIGHTS) + + encoder = Model(inputs=fannet_model.input, outputs=fannet_model.get_layer('dense_1').output) + reader = easyocr.Reader(['en'], gpu=True) + + img_pil = Image.open(TEST_IMAGE_PATH) + detections = reader.readtext(TEST_IMAGE_PATH) + + print(f"{'TEXT':<20} | {'PREDICTED FONT':<25} | {'CONFIDENCE'}") + print("-" * 65) + + for bbox, text, conf in detections: + tl, tr, br, bl = bbox + x_min = max(0, int(min(tl[0], bl[0]))) + y_min = max(0, int(min(tl[1], tr[1]))) + x_max = min(img_pil.width, int(max(tr[0], br[0]))) + y_max = min(img_pil.height, int(max(bl[1], br[1]))) + + if x_max - x_min > 5 and y_max - y_min > 5: + crop = img_pil.crop((x_min, y_min, x_max, y_max)) + + img = crop.convert('L').resize((64, 64), Image.LANCZOS) + arr = np.array(img, dtype=np.float32) + arr = (arr < np.mean(arr)).astype(np.float32) + processed = np.expand_dims(arr, axis=(0, -1)) + + dummy_char = np.zeros((1, 26, 1)) + features = encoder.predict([processed, dummy_char], verbose=0)[0] + + best_font = "Unknown" + min_dist = float('inf') + + for name, db_feat in font_db.items(): + dist = np.linalg.norm(features - db_feat) + if dist < min_dist: + min_dist = dist + best_font = name + + font_conf = 1.0 / (1.0 + min_dist) + print(f"{text:<20} | {best_font:<25} | {font_conf:.4f}") +else: + print("Files not found. Please upload test image.png, fannet.json, fannet_weights.h5, and font_database.pkl") +\ No newline at end of file