Update Batch_Inference.py

Batch_Inference.py

@@ -10,13 +10,13 @@ from tqdm import tqdm
 
 
 HF_USERNAME = "Subh75"
-HF_ORGNAME="LeafNet75"
+HF_ORGNAME = "LeafNet75"
 MODEL_NAME = "Leaf-Annotate-v2"
 HF_MODEL_REPO_ID = f"{HF_ORGNAME}/{MODEL_NAME}"
 
-# Set to your original image  and output folder respectively
-INPUT_IMAGE_DIR = "toast" 
-OUTPUT_MASK_DIR = "masks"
+# Set to your original image and output folder respectively
+INPUT_IMAGE_DIR = "newimgs/images"
+OUTPUT_MASK_DIR = "newimgs/masks"
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 IMG_SIZE = 256
@@ -26,11 +26,11 @@ CONFIDENCE_THRESHOLD = 0.5
 def load_model_from_hub(repo_id: str):
     """Loads the interactive segmentation model from the Hub."""
     print(f"Loading model '{repo_id}' from Hugging Face Hub...")
-    
+
     model = smp.Unet(
         encoder_name="mobilenet_v2",
         encoder_weights=None,
-        in_channels=4, # RGB + Scribble
+        in_channels=4,  # RGB + Scribble
         classes=1,
     )
 
@@ -41,71 +41,85 @@ def load_model_from_hub(repo_id: str):
     print("Model loaded successfully.")
     return model
 
+
 def predict_scribble(model, pil_image, scribble_mask):
     """Runs inference using a scribble and returns a binary mask."""
-    # Preprocess image and scribble
-    img_resized = np.array(pil_image.resize((IMG_SIZE, IMG_SIZE), Image.Resampling.BILINEAR))
-    scribble_resized = cv2.resize(scribble_mask, (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_NEAREST)
+    img_resized = np.array(
+        pil_image.resize((IMG_SIZE, IMG_SIZE), Image.Resampling.BILINEAR)
+    )
+    scribble_resized = cv2.resize(
+        scribble_mask, (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_NEAREST
+    )
 
-    img_tensor = torch.from_numpy(img_resized.astype(np.float32)).permute(2, 0, 1) / 255.0
-    scribble_tensor = torch.from_numpy(scribble_resized.astype(np.float32)).unsqueeze(0) / 255.0
+    img_tensor = (
+        torch.from_numpy(img_resized.astype(np.float32)).permute(2, 0, 1) / 255.0
+    )
+    scribble_tensor = (
+        torch.from_numpy(scribble_resized.astype(np.float32)).unsqueeze(0) / 255.0
+    )
 
     input_tensor = torch.cat([img_tensor, scribble_tensor], dim=0).unsqueeze(0).to(DEVICE)
-    
+
     with torch.no_grad():
         output = model(input_tensor)
-    
+
     probs = torch.sigmoid(output)
     binary_mask_resized = (probs > CONFIDENCE_THRESHOLD).float().squeeze().cpu().numpy()
-    
-    final_mask = cv2.resize(binary_mask_resized, (pil_image.width, pil_image.height), interpolation=cv2.INTER_NEAREST)
+
+    final_mask = cv2.resize(
+        binary_mask_resized, (pil_image.width, pil_image.height), interpolation=cv2.INTER_NEAREST
+    )
     return (final_mask * 255).astype(np.uint8)
 
+
 def main():
-    """
-    Main function to run batch inference on a folder of images.
-    """
+    """Main function to run batch inference on a folder of images."""
     if not os.path.isdir(INPUT_IMAGE_DIR):
         print(f"Error: Input directory not found at '{INPUT_IMAGE_DIR}'")
-        print("Please update the 'INPUT_IMAGE_DIR' variable in the script.")
         return
-        
+
     os.makedirs(OUTPUT_MASK_DIR, exist_ok=True)
-    
+
     model = load_model_from_hub(HF_MODEL_REPO_ID)
-    
-    image_files = [f for f in os.listdir(INPUT_IMAGE_DIR) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
+
+    image_files = [
+        f for f in os.listdir(INPUT_IMAGE_DIR) if f.lower().endswith((".png", ".jpg", ".jpeg"))
+    ]
 
     print(f"\nFound {len(image_files)} images to process.")
-    
+
     for filename in tqdm(image_files, desc="Generating Masks"):
         image_path = os.path.join(INPUT_IMAGE_DIR, filename)
-        
+
         try:
             original_image = Image.open(image_path).convert("RGB")
             h, w = original_image.height, original_image.width
 
+            # Create a dummy scribble (center line)
             scribble = np.zeros((h, w), dtype=np.uint8)
             center_x, center_y = w // 2, h // 2
-            length = int(min(w, h) * 0.2) 
-            
+            length = int(min(w, h) * 0.2)
+
             start_point = (center_x - length // 2, center_y)
             end_point = (center_x + length // 2, center_y)
             cv2.line(scribble, start_point, end_point, 255, thickness=25)
-            
+
+            # Predict mask
             predicted_mask = predict_scribble(model, original_image, scribble)
-            
+
             mask_image = Image.fromarray(predicted_mask)
-            
+
+            # Keep same base name, save as .png in OUTPUT_MASK_DIR
             base_name = Path(filename).stem
-            output_path = os.path.join(OUTPUT_MASK_DIR, f"{base_name}_mask.png")
-            
+            output_path = os.path.join(OUTPUT_MASK_DIR, f"{base_name}.png")
+
             mask_image.save(output_path)
 
         except Exception as e:
-            print(f"\nCould not process {filename}. Error: {e}")
+            print(f"\n❌ Could not process {filename}. Error: {e}")
+
+    print(f"\n Done! Masks saved in '{OUTPUT_MASK_DIR}' with same names as input images.")
 
-    print(f"\n Done! All generated masks have been saved to the '{OUTPUT_MASK_DIR}' folder.")
 
 if __name__ == "__main__":
-    main()
+    main()