creek

  1"""
  2Lighting
  3Original file: https://colab.research.google.com/drive/1edtbuTtWprRBw_lOvGJuEEDu5HGc-KA3
  4"""
  5
  6!pip install git+https://github.com/openai/CLIP.git --q
  7!pip install torch torchvision pillow --q
  8
  9## Build Text Bank
 10
 11import torch
 12import clip
 13from PIL import Image
 14import numpy as np
 15
 16LIGHTING_CLASSES = [
 17    "soft-light",
 18    "specular-highlights",
 19    "backlit",
 20    "studio-lighting",
 21    "flat-lighting",
 22    "dramatic-contrast",
 23    "diffused",
 24]
 25
 26LIGHTING_PHRASES = {
 27    "soft-light": "soft light",
 28    "specular-highlights": "strong specular highlights",
 29    "backlit": "backlit lighting from behind the subject",
 30    "studio-lighting": "professional studio lighting",
 31    "flat-lighting": "flat, low-contrast lighting",
 32    "dramatic-contrast": "dramatic high-contrast lighting",
 33    "diffused": "soft diffused lighting with low shadows",
 34}
 35
 36PROMPT_TEMPLATES = [
 37    "a photo with {}",
 38    "a portrait lit with {}",
 39    "a high quality photograph with {}",
 40    "a cinematic shot using {}",
 41    "an image captured in {}",
 42    "a studio photo with {}",
 43    "a product photo taken with {}",
 44    "a landscape scene under {}",
 45    "a close-up shot with {}",
 46    "a professional photography setup using {}",
 47]
 48
 49
 50device = "cuda" if torch.cuda.is_available() else "cpu"
 51model, preprocess = clip.load("ViT-B/32", device=device)
 52
 53
 54def build_text_bank(output_path="lighting_text_bank.npz"):
 55    all_prompt_texts = []
 56    all_prompt_labels = []
 57    all_prompt_features = []
 58
 59    class_centroids = []
 60
 61    with torch.no_grad():
 62        for lighting in LIGHTING_CLASSES:
 63            phrase = LIGHTING_PHRASES[lighting]
 64
 65            prompts = [tpl.format(phrase) for tpl in PROMPT_TEMPLATES]
 66
 67            tokens = clip.tokenize(prompts).to(device)
 68            text_features = model.encode_text(tokens)
 69            text_features = text_features / text_features.norm(dim=-1, keepdim=True)
 70
 71            all_prompt_texts.extend(prompts)
 72            all_prompt_labels.extend([lighting] * len(prompts))
 73            all_prompt_features.append(text_features.cpu().numpy())
 74
 75            centroid = text_features.mean(dim=0)
 76            centroid = centroid / centroid.norm()
 77            class_centroids.append(centroid.cpu().numpy())
 78
 79    all_prompt_features = np.concatenate(all_prompt_features, axis=0)
 80    class_centroids = np.stack(class_centroids, axis=0)
 81
 82    np.savez(
 83        output_path,
 84        prompt_texts=np.array(all_prompt_texts),
 85        prompt_labels=np.array(all_prompt_labels),
 86        prompt_features=all_prompt_features,
 87        class_centroids=class_centroids,
 88        class_names=np.array(LIGHTING_CLASSES),
 89    )
 90    print(f"Saved text bank to {output_path}")
 91
 92
 93if __name__ == "__main__":
 94    build_text_bank()
 95
 96
 97## Classify Lighting
 98
 99import numpy as np
100import torch
101import clip
102from PIL import Image
103
104
105def get_image_embedding(image_path: str) -> torch.Tensor:
106    image = Image.open(image_path).convert("RGB")
107    img_tensor = preprocess(image).unsqueeze(0).to(device)
108
109    with torch.no_grad():
110        img_feat = model.encode_image(img_tensor)
111        img_feat = img_feat / img_feat.norm(dim=-1, keepdim=True)
112
113    return img_feat.squeeze(0)
114
115
116def load_text_bank(path="lighting_text_bank.npz"):
117    data = np.load(path, allow_pickle=True)
118    centroids = data["class_centroids"]
119    class_names = list(data["class_names"])
120    centroids = centroids / np.linalg.norm(centroids, axis=1, keepdims=True)
121    return centroids, class_names
122
123
124def classify_lighting(
125    image_path: str, text_bank_path="lighting_text_bank.npz", threshold: float = 0.25
126):
127    centroids, class_names = load_text_bank(text_bank_path)
128
129    img_feat = get_image_embedding(image_path)
130    img_np = img_feat.cpu().numpy()[None, :]
131
132    sims = img_np @ centroids.T
133    sims = sims.squeeze(0)
134
135    sorted_indices = np.argsort(sims)[::-1]
136
137    top_n = 3
138    results = []
139    for i in range(min(top_n, len(class_names))):
140        idx = sorted_indices[i]
141        lighting_style = class_names[idx]
142        similarity = float(sims[idx])
143        results.append({"lighting_style": lighting_style, "similarity": similarity})
144
145    return results
146
147
148if __name__ == "__main__":
149    img_path = "/content/test.png"
150    result = classify_lighting(img_path)
151    print(result)
152
153## Demo Runner
154
155!pip install ftfy regex tqdm
156
157import torch
158import clip
159from PIL import Image
160import numpy as np
161
162device = "cuda" if torch.cuda.is_available() else "cpu"
163model, preprocess = clip.load("ViT-B/32", device=device)
164
165print("Device:", device)
166
167LIGHTING_CLASSES = [
168    "soft-light",
169    "specular-highlights",
170    "backlit",
171    "studio-lighting",
172    "flat-lighting",
173    "dramatic-contrast",
174    "diffused",
175]
176
177
178LIGHTING_PHRASES = {
179    "soft-light": "soft light",
180    "specular-highlights": "strong specular highlights",
181    "backlit": "backlit lighting from behind the subject",
182    "studio-lighting": "professional studio lighting",
183    "flat-lighting": "flat, low-contrast lighting",
184    "dramatic-contrast": "dramatic high-contrast lighting",
185    "diffused": "soft diffused lighting with low shadows",
186}
187
188
189PROMPT_TEMPLATES = [
190    "a photo with {}",
191    "a portrait lit with {}",
192    "a high quality photograph with {}",
193    "a cinematic shot using {}",
194    "an image captured in {}",
195    "a studio photo with {}",
196    "a product photo taken with {}",
197    "a landscape scene under {}",
198    "a close-up shot with {}",
199    "a professional photography setup using {}",
200]
201
202
203def build_text_bank(output_path="lighting_text_bank.npz"):
204    all_prompt_texts = []
205    all_prompt_labels = []
206    all_prompt_features = []
207
208    class_centroids = []
209
210    model.eval()
211    with torch.no_grad():
212        for lighting in LIGHTING_CLASSES:
213            phrase = LIGHTING_PHRASES[lighting]
214
215            prompts = [tpl.format(phrase) for tpl in PROMPT_TEMPLATES]
216            tokens = clip.tokenize(prompts).to(device)
217
218            text_features = model.encode_text(tokens)
219            text_features = text_features / text_features.norm(dim=-1, keepdim=True)
220
221            all_prompt_texts.extend(prompts)
222            all_prompt_labels.extend([lighting] * len(prompts))
223            all_prompt_features.append(text_features.cpu().numpy())
224
225            centroid = text_features.mean(dim=0)
226            centroid = centroid / centroid.norm()
227            class_centroids.append(centroid.cpu().numpy())
228
229    all_prompt_features = np.concatenate(all_prompt_features, axis=0)
230    class_centroids = np.stack(class_centroids, axis=0)
231
232    np.savez(
233        output_path,
234        prompt_texts=np.array(all_prompt_texts),
235        prompt_labels=np.array(all_prompt_labels),
236        prompt_features=all_prompt_features,
237        class_centroids=class_centroids,
238        class_names=np.array(LIGHTING_CLASSES),
239    )
240    print(f"Saved text bank → {output_path}")
241
242
243build_text_bank()
244
245
246def get_image_embedding(image_path: str) -> torch.Tensor:
247    image = Image.open(image_path).convert("RGB")
248    img_tensor = preprocess(image).unsqueeze(0).to(device)
249
250    model.eval()
251    with torch.no_grad():
252        img_feat = model.encode_image(img_tensor)
253        img_feat = img_feat / img_feat.norm(dim=-1, keepdim=True)
254
255    return img_feat.squeeze(0)
256
257
258DEFAULT_THRESHOLD = 0.28
259
260
261def load_text_bank(path="lighting_text_bank.npz"):
262    bank = np.load(path, allow_pickle=True)
263    centroids = torch.tensor(bank["class_centroids"]).to(device)
264    class_names = bank["class_names"]
265
266    centroids = centroids / centroids.norm(dim=-1, keepdim=True)
267    return centroids, class_names
268
269
270def classify_lighting(image_path, threshold=DEFAULT_THRESHOLD):
271    img_feat = get_image_embedding(image_path).to(device)
272
273    centroids, class_names = load_text_bank()
274
275    img_feat = img_feat.to(centroids.dtype)
276
277    img_feat = img_feat / img_feat.norm()
278
279    sims = centroids @ img_feat
280    best_score, best_idx = torch.max(sims, dim=0)
281
282    best_score = float(best_score.detach().cpu())
283    sims_list = sims.detach().cpu().tolist()
284
285    sims_dict = {str(cls): float(s) for cls, s in zip(class_names, sims_list)}
286
287    if best_score < threshold:
288        predicted = "no_specific_lighting"
289    else:
290        predicted = str(class_names[best_idx])
291
292    return predicted, best_score, sims_dict
293
294
295# Runner
296
297from google.colab import files
298
299uploaded = files.upload()
300
301image_path = list(uploaded.keys())[0]
302
303pred, score, sims = classify_lighting(image_path)
304
305print("Prediction:", pred)
306print("Best score:", score)
307print("\nAll similarities:")
308for k, v in sims.items():
309    print(f"{k:20s}: {v:.3f}")