Update app.py

app.py

@@ -2391,84 +2391,156 @@ def extract_images_from_pdf(pdf_stream: io.BytesIO):
         raise RuntimeError(f"❌ Error in extract_images_from_pdf: {str(e)}")
 
 # def similarity_matching(input_json_path: str, project_folder: str) -> str:
-def similarity_matching(sprites_data: dict, project_folder: str) -> str: 
-    """
-    Finds the best matching assets for input sprites using a FAISS index and builds a Scratch project.
-
-    Args:
-        sprites_data (str): JSON string containing sprite metadata and base64 image data.
-        project_folder (str): The directory where the final project.json and assets will be saved.
-
-    Returns:
-        str: The path to the generated project.json file.
-    """
-    logger.info("🔍 Running similarity matching with FAISS...")
+def similarity_matching(sprites_data: str, project_folder: str) -> str:
+    logger.info("🔍 Running similarity matching…")
     os.makedirs(project_folder, exist_ok=True)
-    
-    project_json_path = os.path.join(project_folder, 'project.json')
-    # Lazily load the model and FAISS index on the first call
-    load_model_and_index()
-
-    # ==================================================== #
-    #   1. PREPARE INPUT SPRITE IMAGES FROM BASE64 DATA    #
-    # ==================================================== #
-    # The input is a JSON string, so we first need to parse it.
-    #sprites_data_dict = json.loads(sprites_data)
-    
-    # ...
-    # The 'sprites_data' variable is already a dictionary.
-    
-    pil_images = []
-    # Use sprites_data directly
-    for sid, sprite in sprites_data.items(): 
-    #for sid, sprite in sprites_data_dict.items():
-        base64_string = sprite["base64"]
-        # Decode the base64 string to bytes, then open as a PIL Image
-        image_data = base64.b64decode(base64_string.split(',')[-1])
-        image = Image.open(BytesIO(image_data)).convert("RGB")
-        pil_images.append(image)
-        
-    if not pil_images:
-        logger.warning("No images found in the input sprites_data. Aborting.")
-        # Handle case with no images if necessary
-        return ""
 
-    # =========================================== #
-    #   2. ENCODE INPUT IMAGES & SEARCH IN FAISS  #
-    # =========================================== #
-    logger.info(f"Encoding {len(pil_images)} input sprite(s)...")
-    
-    # The model can encode a list of PIL images directly in a batch
-    query_embeddings = MODEL.encode(pil_images, convert_to_tensor=True, show_progress_bar=False)
-    
-    # Ensure the embeddings are float32 numpy arrays for FAISS
-    query_embeddings_np = query_embeddings.cpu().numpy().astype(np.float32)
-
-    # We only need the single best match for each sprite, so k=1
-    k = 1
-    logger.info(f"Searching for the top {k} match for each sprite in FAISS index...")
-    distances, indices = FAISS_INDEX.search(query_embeddings_np, k)
-
-    # `indices` is a 2D array, e.g., [[23], [105], [42]]. We want a flat list of these indices.
-    most_similar_indices = indices.flatten()
-
-    # ========================================================= #
-    #   3. COPY MATCHED ASSETS & BUILD THE FINAL project.json   #
-    #   (This part remains the same as your original logic)     #
-    # ========================================================= #
-    logger.info("Processing matched assets and building project.json...")
-    
+    # ----------------------------------------
+    # CHANGED: define normalized base-paths so startswith() checks work
     backdrop_base_path = os.path.normpath(str(BACKDROP_DIR))
-    sprite_base_path = os.path.normpath(str(SPRITE_DIR))
-    
-    project_data = []
+    sprite_base_path   = os.path.normpath(str(SPRITE_DIR))
+    code_blocks_path = os.path.normpath(str(CODE_BLOCKS_DIR))
+    # ----------------------------------------
+
+    project_json_path = os.path.join(project_folder, "project.json")
+
+    # ==============================
+    #    READ SPRITE METADATA
+    # ==============================
+    sprite_ids, sprite_base64 = [], []
+    for sid, sprite in sprites_data.items():
+        sprite_ids.append(sid)
+        sprite_base64.append(sprite["base64"])
+
+    sprite_images_bytes = []
+    for b64 in sprite_base64:
+        img = Image.open(BytesIO(base64.b64decode(b64.split(",")[-1]))).convert("RGB")
+        buffer = BytesIO()
+        img.save(buffer, format="PNG")
+        buffer.seek(0)
+        sprite_images_bytes.append(buffer)
+
+    # =========================================
+    #  Build the list of all candidate images
+    #  (REPLACED: now loaded from FAISS paths JSON)
+    # =========================================
+    # NOTE: previously you had a hardcoded `folder_image_paths` sample.
+    # We'll load image paths from BLOCKS_DIR/image_paths.json (faiss index companion).
+    try:
+        index_path = os.path.join(str(BLOCKS_DIR), "faiss_index 2.bin")
+        paths_json_path = os.path.join(str(BLOCKS_DIR), "image_paths 2.json")
+        # load image paths (these should be full paths or relative paths you used when building the index)
+        with open(paths_json_path, "r", encoding="utf-8") as f:
+            image_paths = json.load(f)
+        folder_image_paths = [os.path.normpath(str(p)) for p in image_paths]
+    except Exception as e:
+        # Fallback: if image_paths.json not available, fall back to any existing embed.json usage
+        logger.error(f"Failed to load FAISS paths JSON from {paths_json_path}: {e}")
+        # try to load previous embedding_json and recover its paths (if present)
+        try:
+            with open(f"{BLOCKS_DIR}/embed.json", "r", encoding="utf-8") as f:
+                embedding_json = json.load(f)
+            # embedding_json entries expected to have a 'path' or similar field; try common keys
+            folder_image_paths = []
+            for idx_entry in embedding_json:
+                if "path" in idx_entry:
+                    folder_image_paths.append(os.path.normpath(str(idx_entry["path"])))
+                elif "filename" in idx_entry:
+                    folder_image_paths.append(os.path.normpath(str(idx_entry["filename"])))
+                else:
+                    # If no path available, synthesize a placeholder to avoid crashing later
+                    folder_image_paths.append(os.path.normpath(str(BLOCKS_DIR)))
+            logger.warning("Using embed.json fallback to populate folder_image_paths.")
+        except Exception as e2:
+            logger.error(f"Failed to load embed.json fallback: {e2}")
+            folder_image_paths = []
+
+    # -----------------------------------------
+    #  Load FAISS index and SentenceTransformer model
+    #  and perform search for each sprite image
+    # -----------------------------------------
+    most_similar_indices = None
+    try:
+        import faiss  # local import in case not available globally
+        from sentence_transformers import SentenceTransformer
+        # load model (this may take time but ensures correct embedder)
+        logger.info("Loading SentenceTransformer 'clip-ViT-L-14' model for embeddings...")
+        st_model = SentenceTransformer('clip-ViT-L-14')
+
+        logger.info(f"Loading FAISS index from {index_path} ...")
+        index = faiss.read_index(index_path)
+
+        # Convert BytesIO images to PIL images for batch encoding
+        pil_images = []
+        for buf in sprite_images_bytes:
+            buf.seek(0)
+            pil = Image.open(buf).convert("RGB")
+            pil_images.append(pil)
+
+        if len(pil_images) == 0:
+            raise ValueError("No sprite images available to encode.")
+
+        logger.info(f"Encoding {len(pil_images)} sprite images with SentenceTransformer...")
+        # encode returns numpy array (n_samples, dim)
+        sprite_embeddings = st_model.encode(pil_images, batch_size=8, convert_to_numpy=True)
+        # ensure float32 and contiguous
+        sprite_embeddings = np.ascontiguousarray(sprite_embeddings.astype(np.float32))
+
+        # number of neighbors to request (we only need the top-1 to emulate previous behavior)
+        k = 1
+        logger.info(f"Performing FAISS search (k={k}) for all sprites...")
+        distances, indices = index.search(sprite_embeddings, k)  # distances.shape = (N, k); indices.shape = (N, k)
+        # take top-1 index for each sprite
+        if indices is not None and indices.shape[0] > 0:
+            most_similar_indices = indices[:, 0].tolist()
+        else:
+            most_similar_indices = []
+    except Exception as e:
+        logger.error(f"FAISS / SentenceTransformer matching failed: {e}. Falling back to simple in-memory similarity if possible.")
+        # fallback: if we previously loaded embedding_json, use that in-memory approach (cosine)
+        try:
+            if 'embedding_json' not in locals():
+                with open(f"{BLOCKS_DIR}/embed.json", "r", encoding="utf-8") as f:
+                    embedding_json = json.load(f)
+            # build img_matrix from embedding_json entries (assumes "embeddings" key)
+            img_matrix = np.array([img["embeddings"] for img in embedding_json], dtype=np.float32)
+            # compute embeddings for sprites using any available embed function (init_dinov2/embed_bytesio_list previously used)
+            init_dinov2()
+            sprite_matrix = embed_bytesio_list(sprite_images_bytes, batch_size=8)
+            sprite_matrix = l2_normalize_rows(sprite_matrix)
+            img_matrix = l2_normalize_rows(img_matrix)
+            similarity = np.matmul(sprite_matrix, img_matrix.T)
+            most_similar_indices = np.argmax(similarity, axis=1).tolist()
+            logger.info("Fallback in-memory similarity search succeeded.")
+        except Exception as e2:
+            logger.error(f"Fallback similarity also failed: {e2}")
+            most_similar_indices = []
+
+    # Ensure we have a list with same length as number of sprites (or empty)
+    if most_similar_indices is None:
+        most_similar_indices = []
+
+    # =========================================
+    #  Copy matched sprite assets + collect data
+    # =========================================
+    project_data   = []
     copied_folders = set()
-    
+
     for sprite_idx, matched_idx in enumerate(most_similar_indices):
-        # IMAGE_PATHS is the list loaded from image_paths.json
-        matched_image_path = IMAGE_PATHS[matched_idx]
-        matched_folder = os.path.dirname(matched_image_path)
+        if matched_idx is None:
+            continue
 
+        # protect against out-of-range indices
+        try:
+            matched_image_path = folder_image_paths[int(matched_idx)]
+        except Exception as e:
+            logger.warning(f"Matched index {matched_idx} invalid: {e}")
+            continue
+
+        matched_image_path = os.path.normpath(str(matched_image_path))
+        matched_folder     = os.path.dirname(matched_image_path)
+
+        # CHANGED: use our new normalized sprite_base_path
         if not matched_folder.startswith(sprite_base_path):
             continue
 
@@ -2482,33 +2554,39 @@ def similarity_matching(sprites_data: dict, project_folder: str) -> str:
             logger.warning(f"No sprite.json in {matched_folder}")
             continue
 
-        with open(sprite_json_path, 'r') as f:
+        with open(sprite_json_path, 'r', encoding='utf-8') as f:
             sprite_info = json.load(f)
-        
-        # Copy all non-matched files
+        # copy all non‐matched files
         for fname in os.listdir(matched_folder):
             if fname in (os.path.basename(matched_image_path), 'sprite.json'):
                 continue
-            shutil.copy2(os.path.join(matched_folder, fname), os.path.join(project_folder, fname))
+            try:
+                shutil.copy2(os.path.join(matched_folder, fname),
+                             os.path.join(project_folder, fname))
+            except Exception as e:
+                logger.error(f"Failed to copy sprite asset {fname} from {matched_folder}: {e}")
         project_data.append(sprite_info)
 
-    # (The rest of your backdrop handling and project.json finalization logic would go here,
-    # unchanged from your original script. I am omitting it for brevity but you should include it.)
-    # =========================================
-    #  Copy matched backdrop assets + collect
-    # =========================================
     # =========================================
     #  Copy matched backdrop assets + collect
     # =========================================
     backdrop_data = []
     copied_backdrop_folders = set()
     for backdrop_idx, matched_idx in enumerate(most_similar_indices):
-        # FIXED: Use the correct variable name
-        matched_image_path = IMAGE_PATHS[matched_idx]
+        if matched_idx is None:
+            continue
+
+        # protect against out-of-range indices
+        try:
+            matched_image_path = folder_image_paths[int(matched_idx)]
+        except Exception as e:
+            logger.warning(f"Matched index {matched_idx} invalid for backdrop loop: {e}")
+            continue
+
+        matched_image_path = os.path.normpath(str(matched_image_path))
         matched_folder     = os.path.dirname(matched_image_path)
-        # ... rest of the loop
         matched_filename = os.path.basename(matched_image_path)
-        
+
         # CHANGED: use our new normalized backdrop_base_path
         if not matched_folder.startswith(backdrop_base_path):
             continue
@@ -2529,14 +2607,11 @@ def similarity_matching(sprites_data: dict, project_folder: str) -> str:
             logger.info(f"✅ Copied matched backdrop image {matched_filename} to {project_folder}")
         except Exception as e:
             logger.error(f"❌ Failed to copy matched backdrop {matched_image_path}: {e}")
-        
+
         # copy non‐matched files
         for fname in os.listdir(matched_folder):
-            # if fname in (os.path.basename(matched_image_path), 'project.json'):
             if fname in {matched_filename, 'project.json'}:
                 continue
-            # shutil.copy2(os.path.join(matched_folder, fname),
-            #              os.path.join(project_folder, fname))
             src = os.path.join(matched_folder, fname)
             dst = os.path.join(project_folder, fname)
             if os.path.isfile(src):
@@ -2549,7 +2624,7 @@ def similarity_matching(sprites_data: dict, project_folder: str) -> str:
         # append the stage‐target from its project.json
         pj = os.path.join(matched_folder, 'project.json')
         if os.path.exists(pj):
-            with open(pj, 'r') as f:
+            with open(pj, 'r', encoding='utf-8') as f:
                 bd_json = json.load(f)
             for tgt in bd_json.get("targets", []):
                 if tgt.get("isStage"):
@@ -2557,7 +2632,6 @@ def similarity_matching(sprites_data: dict, project_folder: str) -> str:
         else:
             logger.warning(f"No project.json in {matched_folder}")
 
-
     # =========================================
     #  Merge into final Scratch project.json
     # =========================================
@@ -2585,31 +2659,31 @@ def similarity_matching(sprites_data: dict, project_folder: str) -> str:
                 if key not in seen_costumes:
                     seen_costumes.add(key)
                     all_costumes.append(costume)
-    
+
             if i == 0:
                 sounds = bd.get("sounds", [])
         stage_obj={
             "isStage": True,
             "name": "Stage",
             "objName": "Stage",
-            "variables": {}, 
-            "lists": {}, 
+            "variables": {},
+            "lists": {},
             "broadcasts": {},
-            "blocks": {}, 
+            "blocks": {},
             "comments": {},
             "currentCostume": 1 if len(all_costumes) > 1 else 0,
             "costumes": all_costumes,
             "sounds": sounds,
-            "volume": 100, 
+            "volume": 100,
             "layerOrder": 0,
-            "tempo": 60, 
+            "tempo": 60,
             "videoTransparency": 50,
             "videoState": "on",
             "textToSpeechLanguage": None
         }
         final_project["targets"].insert(0, stage_obj)
     else:
-        logger.warning("⚠️ No backdrop matched. Using default static backdrop.")
+        logger.warning(⚠️ No backdrop matched. Using default static backdrop.")
         default_backdrop_path = BACKDROP_DIR / "cd21514d0531fdffb22204e0ec5ed84a.svg"
         default_backdrop_name = "cd21514d0531fdffb22204e0ec5ed84a.svg"
 
@@ -2623,7 +2697,7 @@ def similarity_matching(sprites_data: dict, project_folder: str) -> str:
             logger.info(f"✅ Default backdrop sound copied to project: {default_backdrop_sound_name}")
         except Exception as e:
             logger.error(f"❌ Failed to copy default backdrop: {e}")
-            
+
         stage_obj={
             "isStage": True,
             "name": "Stage",
@@ -2664,10 +2738,11 @@ def similarity_matching(sprites_data: dict, project_folder: str) -> str:
         }
         final_project["targets"].insert(0, stage_obj)
 
-    with open(project_json_path, 'w') as f:
+    with open(project_json_path, 'w', encoding='utf-8') as f:
         json.dump(final_project, f, indent=2)
 
     return project_json_path
+    
 # def similarity_matching(sprites_data: str, project_folder: str) -> str:
 #     logger.info("🔍 Running similarity matching…")
 #     os.makedirs(project_folder, exist_ok=True)