LeafNet75/Leaf-Annotate-v2

1import os2from pathlib import Path3import torch4from PIL import Image5import numpy as np6import cv27import segmentation_models_pytorch as smp8from huggingface_hub import hf_hub_download9from tqdm import tqdm10 11 12HF_USERNAME = "Subh75"13HF_ORGNAME = "LeafNet75"14MODEL_NAME = "Leaf-Annotate-v2"15HF_MODEL_REPO_ID = f"{HF_ORGNAME}/{MODEL_NAME}"16 17# Set to your original image and output folder respectively18INPUT_IMAGE_DIR = "newimgs/images"19OUTPUT_MASK_DIR = "newimgs/masks"20 21DEVICE = "cuda" if torch.cuda.is_available() else "cpu"22IMG_SIZE = 25623CONFIDENCE_THRESHOLD = 0.524 25 26def load_model_from_hub(repo_id: str):27    """Loads the interactive segmentation model from the Hub."""28    print(f"Loading model '{repo_id}' from Hugging Face Hub...")29 30    model = smp.Unet(31        encoder_name="mobilenet_v2",32        encoder_weights=None,33        in_channels=4,  # RGB + Scribble34        classes=1,35    )36 37    model_weights_path = hf_hub_download(repo_id=repo_id, filename="best_model.pth")38    model.load_state_dict(torch.load(model_weights_path, map_location=DEVICE))39    model.to(DEVICE)40    model.eval()41    print("Model loaded successfully.")42    return model43 44 45def predict_scribble(model, pil_image, scribble_mask):46    """Runs inference using a scribble and returns a binary mask."""47    img_resized = np.array(48        pil_image.resize((IMG_SIZE, IMG_SIZE), Image.Resampling.BILINEAR)49    )50    scribble_resized = cv2.resize(51        scribble_mask, (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_NEAREST52    )53 54    img_tensor = (55        torch.from_numpy(img_resized.astype(np.float32)).permute(2, 0, 1) / 255.056    )57    scribble_tensor = (58        torch.from_numpy(scribble_resized.astype(np.float32)).unsqueeze(0) / 255.059    )60 61    input_tensor = torch.cat([img_tensor, scribble_tensor], dim=0).unsqueeze(0).to(DEVICE)62 63    with torch.no_grad():64        output = model(input_tensor)65 66    probs = torch.sigmoid(output)67    binary_mask_resized = (probs > CONFIDENCE_THRESHOLD).float().squeeze().cpu().numpy()68 69    final_mask = cv2.resize(70        binary_mask_resized, (pil_image.width, pil_image.height), interpolation=cv2.INTER_NEAREST71    )72    return (final_mask * 255).astype(np.uint8)73 74 75def main():76    """Main function to run batch inference on a folder of images."""77    if not os.path.isdir(INPUT_IMAGE_DIR):78        print(f"Error: Input directory not found at '{INPUT_IMAGE_DIR}'")79        return80 81    os.makedirs(OUTPUT_MASK_DIR, exist_ok=True)82 83    model = load_model_from_hub(HF_MODEL_REPO_ID)84 85    image_files = [86        f for f in os.listdir(INPUT_IMAGE_DIR) if f.lower().endswith((".png", ".jpg", ".jpeg"))87    ]88 89    print(f"\nFound {len(image_files)} images to process.")90 91    for filename in tqdm(image_files, desc="Generating Masks"):92        image_path = os.path.join(INPUT_IMAGE_DIR, filename)93 94        try:95            original_image = Image.open(image_path).convert("RGB")96            h, w = original_image.height, original_image.width97 98            # Create a dummy scribble (center line)99            scribble = np.zeros((h, w), dtype=np.uint8)100            center_x, center_y = w // 2, h // 2101            length = int(min(w, h) * 0.2)102 103            start_point = (center_x - length // 2, center_y)104            end_point = (center_x + length // 2, center_y)105            cv2.line(scribble, start_point, end_point, 255, thickness=25)106 107            # Predict mask108            predicted_mask = predict_scribble(model, original_image, scribble)109 110            mask_image = Image.fromarray(predicted_mask)111 112            # Keep same base name, save as .png in OUTPUT_MASK_DIR113            base_name = Path(filename).stem114            output_path = os.path.join(OUTPUT_MASK_DIR, f"{base_name}.png")115 116            mask_image.save(output_path)117 118        except Exception as e:119            print(f"\n Could not process {filename}. Error: {e}")120 121    print(f"\n Done! Masks saved in '{OUTPUT_MASK_DIR}' with same names as input images.")122 123 124if __name__ == "__main__":125    main()126