import os import glob import random import shutil import torch from PIL import Image from tqdm import tqdm from transformers import Owlv2Processor, Owlv2ForObjectDetection # --------------------------------------------------------------------- # フェーズ2: トレンドのアバター服飾タグリスト（AI提案） # ※これらのキーワード（英語）をテキストプロンプトとしてゼロショットAIに投げて枠を探させます。 # --------------------------------------------------------------------- TEXT_PROMPTS = [ "gothic dress", "cyberpunk clothes", "techwear", "maid outfit", "swimsuit", "bikini", "school uniform", "casual wear", "frill skirt", "hoodie", "jacket", "shorts", "twintails hair", "ponytail hair", "side ponytail hair", "braids hair", "short hair", "bob hair", "long hair", "straight hair", "curly hair", "cat ears", # nekomimi "rabbit ears", # usagimimi "animal ears", # kemonomimi "sheep ears", "fox ears", "horns", "hair clip", "headset", "halo", "crown", "glasses", "goggles", "choker", "ribbon", "hat", ] # BOOTHのユーザーが使うであろうタグ名と、データセット内での一貫性を保つため、 # もし見つかった場合は以下のID（または新規ID）としてyoloデータセットにマッピングします。 # ※今回は簡単のため、見つかった単語のインデックスを独自の一時的なクラスIDとして付与し、 # あとで統合データセットを作るときに data.yaml を自動更新する仕組みにします。 def perform_zero_shot_annotation(num_samples=1000, confidence_threshold=0.1): backend_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) # Check device device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Using device: {device}") # Load Owlv2 model and processor (powerful open-vocabulary object detector) # We use owlv2-base-patch16-ensemble as it has good zero-shot accuracy model_id = "google/owlv2-base-patch16-ensemble" print(f"Loading {model_id} ...") try: processor = Owlv2Processor.from_pretrained(model_id) model = Owlv2ForObjectDetection.from_pretrained(model_id).to(device) except Exception as e: print(f"Error loading model: {e}") print("Please ensure you have internet access to download the weights on first run.") return raw_images_dir = os.path.join(backend_dir, "scraper", "data", "raw_images") output_dataset_dir = os.path.join(backend_dir, "yolo_dataset", "zero_shot_auto") images_out_dir = os.path.join(output_dataset_dir, "images") labels_out_dir = os.path.join(output_dataset_dir, "labels") os.makedirs(images_out_dir, exist_ok=True) os.makedirs(labels_out_dir, exist_ok=True) all_images = glob.glob(os.path.join(raw_images_dir, "*.jpg")) if not all_images: print("No raw images found.") return random.shuffle(all_images) samples = all_images[:num_samples] print(f"Starting zero-shot annotation for {len(samples)} images against {len(TEXT_PROMPTS)} tags...") successful_annotated = 0 for img_path in tqdm(samples): filename = os.path.basename(img_path) img_name, ext = os.path.splitext(filename) try: image = Image.open(img_path).convert("RGB") # owl vit expects prompts formatted slightly specifically # format: "a photo of an object" or just "object" -> "a photo of a [item]" works best texts = [[f"a photo of a {prompt}" for prompt in TEXT_PROMPTS]] inputs = processor(text=texts, images=image, return_tensors="pt").to(device) with torch.no_grad(): outputs = model(**inputs) # Target image sizes (height, width) to rescale bounding boxes back to original size target_sizes = torch.tensor([image.size[::-1]]).to(device) results = processor.image_processor.post_process_object_detection(outputs=outputs, target_sizes=target_sizes, threshold=confidence_threshold) # Retrieve predictions for the first image i = 0 text = texts[i] # No threshold filter here to see best raw scores results_raw = processor.image_processor.post_process_object_detection(outputs=outputs, target_sizes=target_sizes, threshold=0.0) boxes, scores, labels = results_raw[i]["boxes"], results_raw[i]["scores"], results_raw[i]["labels"] if len(scores) > 0: print(f"[{filename}] Max score: {scores.max().item():.3f} for label '{TEXT_PROMPTS[labels[scores.argmax()].item()]}'") else: print(f"[{filename}] No boxes generated at all.") has_detections = False label_lines = [] # Apply our actual custom threshold for box, score, label in zip(boxes, scores, labels): if score.item() < confidence_threshold: continue box = [round(j, 2) for j in box.tolist()] # box is [xmin, ymin, xmax, ymax] x_min, y_min, x_max, y_max = box img_w, img_h = image.size # YOLO format: cx, cy, w, h normalized cx = ((x_min + x_max) / 2) / img_w cy = ((y_min + y_max) / 2) / img_h w = (x_max - x_min) / img_w h = (y_max - y_min) / img_h # clamp to 0-1 cx = max(0, min(cx, 1)) cy = max(0, min(cy, 1)) w = max(0, min(w, 1)) h = max(0, min(h, 1)) cls_id = label.item() # This corresponds to the index in TEXT_PROMPTS label_lines.append(f"{cls_id} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}") has_detections = True if has_detections: dest_img = os.path.join(images_out_dir, filename) shutil.copy2(img_path, dest_img) label_path = os.path.join(labels_out_dir, f"{img_name}.txt") with open(label_path, "w", encoding="utf-8") as f: f.write("\n".join(label_lines) + "\n") successful_annotated += 1 except Exception as e: print(f"Error processing {filename}: {e}") print(f"Zero-shot auto-annotation complete! {successful_annotated} images successfully annotated.") # Save the custom classes.txt so we know what they map to with open(os.path.join(output_dataset_dir, "classes.txt"), "w", encoding="utf-8") as f: f.write("\n".join(TEXT_PROMPTS) + "\n") if __name__ == "__main__": # Run on 1000 images with expanded tags for quality over quantity in this iteration perform_zero_shot_annotation(num_samples=1000, confidence_threshold=0.1)